Gpu support using cupy

.github/workflows/lint-and-test.yml

@@ -48,7 +48,7 @@ jobs:
       run: |
         python -m pip install flake8 pytest
         if [ -f requirements.txt ]; then pip install -r requirements.txt; fi
-        python -m pip install -e .[jax]
+        python -m pip install -e .
     - name: Lint with flake8
       run: |
         # stop the build if there are Python syntax errors or undefined names

docs/api/solver.rst

@@ -26,6 +26,9 @@ in the dictionary to ensure current conservation.
 .. autoclass:: tdgl.SolverOptions
     :members:
 
+.. autoclass:: tdgl.TDGLSolver
+    :members:
+
 .. autoenum:: tdgl.solver.options.SparseSolver
 
 .. autoclass:: tdgl.Parameter

docs/background.rst

@@ -465,7 +465,7 @@ The strategy for updating the induced vector potential to converge to a self-con
 
 
 The integer index :math:`s` counts the number of iterations performed in the self-consistent calculation.
-The parameters :math:`\alpha\in(0,\infty)` and :math:`\beta\in(0,1)` in :eq:`polyak` can be set by the user,
+The parameters :math:`\alpha\in(0,\infty)` and :math:`\beta\in(0,1]` in :eq:`polyak` can be set by the user,
 and the initial conditions for :eq:`polyak` are :math:`\mathbf{A}^{n,0}_{\mathrm{induced},ij} = \mathbf{A}^{n-1}_{\mathrm{induced},ij}`
 and :math:`\mathbf{v}^{n,0}_{ij} = \mathbf{0}`. The iterative application of :eq:`polyak` terminates when the relative change in the
 induced vector potential between iterations falls below a user-defined tolerance.

docs/installation.rst

@@ -53,16 +53,17 @@ To install an editable version of ``pyTDGL`` for development, run:
 
   :ref:`Contributing to pyTDGL <about/contributing:Contributing>`
 
-Optional dependencies
----------------------
+
+Alternative sparse solvers
+--------------------------
 
 ``tdgl`` supports multiple solvers for sparse systems of linear equations: `SuperLU <https://portal.nersc.gov/project/sparse/superlu/>`_ (the default),
 `UMFPACK <https://people.engr.tamu.edu/davis/suitesparse.html>`_, and `MKL PARDISO <https://www.intel.com/content/www/us/en/docs/onemkl/developer-reference-c/2023-0/onemkl-pardiso-parallel-direct-sparse-solver-iface.html>`_.
 
 ``SuperLU``, the default solver, is included in ``scipy`` and therefore requires no additional installation.
 
 The other solvers, ``UMFPACK`` and ``PARDISO``, may outperform ``SuperLU`` on certain problems and certain CPU hardware.
-In particular, ``PARDISO`` is optimized for Intel CPUs and, unlike ``SuperLU`` and ``UMFPACK``, is multithreaded.
+In particular, ``PARDISO`` is optimized for Intel CPUs and, unlike ``SuperLU`` and ``UMFPACK``, ``PARDISO`` is multithreaded.
 This means that ``PARDISO`` may perform best when solving models with very large meshes on Intel CPUs.
 
 Your mileage may vary, so we encourage you to try the different solvers if you are looking to optimize the run time of your ``tdgl`` simulations.
@@ -99,10 +100,32 @@ Installing PARDISO
   # or conda install -c conda-forge pypardiso
   # or pip install tdgl[pardiso]
 
-Verify the installation
------------------------
+GPU acceleration
+----------------
+
+For users with an NVIDIA or AMD GPU, ``tdgl`` can be accelerated using the `CuPy <https://cupy.dev/>`_ library.
+First install the appropriate version of ``cupy`` for your GPU hardware and driver version
+(see installation instructions `here <https://docs.cupy.dev/en/stable/install.html>`_).
+Then set the ``gpu`` attribute of :class:`tdgl.SolverOptions` to ``True``. Setting ``tdgl.SolverOptions.gpu = True``
+means that essentially all portions of the simulation *except* the sparse linear solve used to compute the scalar electric potential
+:math:`\mu(\mathbf{r}, t)` will be performed on the GPU. The sparse linear solve will be performed on the CPU using the solver specified by
+``tdgl.SolverOptions.sparse_solver`` (by default, ``SuperLU``). One can also perform the sparse linear solve on the GPU using
+``cupy`` by setting ``tdgl.SolverOptions.sparse_solver = "cupy"``, however emperically it seems that this is slower than
+performing the sparse linear solve on the CPU.
+
+Due to overheads related to transferring data between the CPU and GPU, it is expected that ``cupy`` will provide
+a significant speedup only for models with relatively large meshes and/or models that include `screening <notebooks/screening.ipynb>`_.
+Please open a `GitHub issue <https://github.com/loganbvh/py-tdgl/issues>`_ if you have any problems using ``tdgl`` with ``cupy``.
+
+.. note::
+
+  Note that ``cupy`` support for AMD GPUs is `currently experimental <https://docs.cupy.dev/en/stable/install.html#using-cupy-on-amd-gpu-experimental>`_.
+
+
+Verifying the installation
+--------------------------
 
-To verify your installation by running the ``tdgl`` test suite,
+If you would like to verify your installation by running the ``tdgl`` test suite,
 execute the following command in a terminal:
 
 .. code-block:: bash

docs/notebooks/logo.ipynb

@@ -233,7 +233,7 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "Constructing Voronoi polygons: 100%|█████| 4347/4347 [00:00<00:00, 22715.27it/s]\n"
+      "Constructing Voronoi polygons: 100%|██████████████████████████████████████████████████████████████████████████| 4347/4347 [00:00<00:00, 22681.73it/s]\n"
      ]
     },
     {
@@ -276,7 +276,7 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "Simulating: 100%|██████████████████████████▉| 800/800 [02:00<00:00,  6.63tau/s ]\n"
+      "Simulating: 100%|███████████████████████████████████████████████████████████████████████████████████████████████▉| 800/800 [01:55<00:00,  6.90tau/s ]\n"
      ]
     }
    ],
@@ -9471,7 +9471,7 @@
     {
      "data": {
       "text/html": [
-       "<table><tr><th>Software</th><th>Version</th></tr><tr><td>tdgl</td><td>0.5.1; git revision 03859ea [2023-09-21]</td></tr><tr><td>Numpy</td><td>1.24.3</td></tr><tr><td>SciPy</td><td>1.10.1</td></tr><tr><td>matplotlib</td><td>3.7.1</td></tr><tr><td>jax</td><td>None</td></tr><tr><td>numba</td><td>0.57.1</td></tr><tr><td>IPython</td><td>8.14.0</td></tr><tr><td>Python</td><td>3.10.11 | packaged by conda-forge | (main, May 10 2023, 19:01:19) [Clang 14.0.6 ]</td></tr><tr><td>OS</td><td>posix [darwin]</td></tr><tr><td>Number of CPUs</td><td>Physical: 10, Logical: 10</td></tr><tr><td>BLAS Info</td><td>OPENBLAS</td></tr><tr><td colspan='2'>Thu Sep 21 15:35:01 2023 PDT</td></tr></table>"
+       "<table><tr><th>Software</th><th>Version</th></tr><tr><td>tdgl</td><td>0.6.0; git revision e6d3811 [2023-09-27]</td></tr><tr><td>Numpy</td><td>1.24.3</td></tr><tr><td>SciPy</td><td>1.10.1</td></tr><tr><td>matplotlib</td><td>3.7.1</td></tr><tr><td>cupy</td><td>None</td></tr><tr><td>numba</td><td>0.57.1</td></tr><tr><td>IPython</td><td>8.14.0</td></tr><tr><td>Python</td><td>3.10.11 | packaged by conda-forge | (main, May 10 2023, 19:01:19) [Clang 14.0.6 ]</td></tr><tr><td>OS</td><td>posix [darwin]</td></tr><tr><td>Number of CPUs</td><td>Physical: 10, Logical: 10</td></tr><tr><td>BLAS Info</td><td>OPENBLAS</td></tr><tr><td colspan='2'>Wed Sep 27 14:58:29 2023 PDT</td></tr></table>"
       ],
       "text/plain": [
        "<IPython.core.display.HTML object>"

docs/notebooks/mesh.ipynb

@@ -29,7 +29,7 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "Constructing Voronoi polygons: 100%|████████| 525/525 [00:00<00:00, 6318.46it/s]\n"
+      "Constructing Voronoi polygons: 100%|████████████████████████████████████████████████████████████████████████████| 525/525 [00:00<00:00, 21170.12it/s]\n"
      ]
     }
    ],

docs/notebooks/polygons.ipynb

@@ -469,9 +469,9 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "Constructing Voronoi polygons: 100%|███████| 459/459 [00:00<00:00, 16340.20it/s]\n",
-      "Constructing Voronoi polygons: 100%|███████| 459/459 [00:00<00:00, 17271.97it/s]\n",
-      "Constructing Voronoi polygons: 100%|███████| 459/459 [00:00<00:00, 17162.65it/s]\n"
+      "Constructing Voronoi polygons: 100%|████████████████████████████████████████████████████████████████████████████| 459/459 [00:00<00:00, 16417.53it/s]\n",
+      "Constructing Voronoi polygons: 100%|████████████████████████████████████████████████████████████████████████████| 459/459 [00:00<00:00, 17100.75it/s]\n",
+      "Constructing Voronoi polygons: 100%|████████████████████████████████████████████████████████████████████████████| 459/459 [00:00<00:00, 17710.51it/s]\n"
      ]
     },
     {
@@ -520,7 +520,7 @@
     {
      "data": {
       "text/html": [
-       "<table><tr><th>Software</th><th>Version</th></tr><tr><td>tdgl</td><td>0.5.1; git revision 03859ea [2023-09-21]</td></tr><tr><td>Numpy</td><td>1.24.3</td></tr><tr><td>SciPy</td><td>1.10.1</td></tr><tr><td>matplotlib</td><td>3.7.1</td></tr><tr><td>jax</td><td>None</td></tr><tr><td>numba</td><td>0.57.1</td></tr><tr><td>IPython</td><td>8.14.0</td></tr><tr><td>Python</td><td>3.10.11 | packaged by conda-forge | (main, May 10 2023, 19:01:19) [Clang 14.0.6 ]</td></tr><tr><td>OS</td><td>posix [darwin]</td></tr><tr><td>Number of CPUs</td><td>Physical: 10, Logical: 10</td></tr><tr><td>BLAS Info</td><td>OPENBLAS</td></tr><tr><td colspan='2'>Thu Sep 21 15:32:52 2023 PDT</td></tr></table>"
+       "<table><tr><th>Software</th><th>Version</th></tr><tr><td>tdgl</td><td>0.6.0; git revision e6d3811 [2023-09-27]</td></tr><tr><td>Numpy</td><td>1.24.3</td></tr><tr><td>SciPy</td><td>1.10.1</td></tr><tr><td>matplotlib</td><td>3.7.1</td></tr><tr><td>cupy</td><td>None</td></tr><tr><td>numba</td><td>0.57.1</td></tr><tr><td>IPython</td><td>8.14.0</td></tr><tr><td>Python</td><td>3.10.11 | packaged by conda-forge | (main, May 10 2023, 19:01:19) [Clang 14.0.6 ]</td></tr><tr><td>OS</td><td>posix [darwin]</td></tr><tr><td>Number of CPUs</td><td>Physical: 10, Logical: 10</td></tr><tr><td>BLAS Info</td><td>OPENBLAS</td></tr><tr><td colspan='2'>Wed Sep 27 14:54:08 2023 PDT</td></tr></table>"
       ],
       "text/plain": [
        "<IPython.core.display.HTML object>"

docs/notebooks/py-mesh.ipynb

@@ -29,7 +29,7 @@
     "import tdgl\n",
     "from tdgl.geometry import box, ensure_unique, close_curve\n",
     "\n",
-    "SAVE_FIGURES = True\n",
+    "SAVE_FIGURES = False\n",
     "image_path = os.path.join(os.pardir, \"images\")"
    ]
   },
@@ -184,7 +184,7 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "Constructing Voronoi polygons: 100%|█████| 4190/4190 [00:00<00:00, 21720.62it/s]\n"
+      "Constructing Voronoi polygons: 100%|██████████████████████████████████████████████████████████████████████████| 4190/4190 [00:00<00:00, 22042.31it/s]\n"
      ]
     },
     {
@@ -361,7 +361,7 @@
     {
      "data": {
       "text/html": [
-       "<table><tr><th>Software</th><th>Version</th></tr><tr><td>tdgl</td><td>0.5.1; git revision 03859ea [2023-09-21]</td></tr><tr><td>Numpy</td><td>1.24.3</td></tr><tr><td>SciPy</td><td>1.10.1</td></tr><tr><td>matplotlib</td><td>3.7.1</td></tr><tr><td>jax</td><td>None</td></tr><tr><td>numba</td><td>0.57.1</td></tr><tr><td>IPython</td><td>8.14.0</td></tr><tr><td>Python</td><td>3.10.11 | packaged by conda-forge | (main, May 10 2023, 19:01:19) [Clang 14.0.6 ]</td></tr><tr><td>OS</td><td>posix [darwin]</td></tr><tr><td>Number of CPUs</td><td>Physical: 10, Logical: 10</td></tr><tr><td>BLAS Info</td><td>OPENBLAS</td></tr><tr><td colspan='2'>Thu Sep 21 15:32:40 2023 PDT</td></tr></table>"
+       "<table><tr><th>Software</th><th>Version</th></tr><tr><td>tdgl</td><td>0.6.0; git revision e6d3811 [2023-09-27]</td></tr><tr><td>Numpy</td><td>1.24.3</td></tr><tr><td>SciPy</td><td>1.10.1</td></tr><tr><td>matplotlib</td><td>3.7.1</td></tr><tr><td>cupy</td><td>None</td></tr><tr><td>numba</td><td>0.57.1</td></tr><tr><td>IPython</td><td>8.14.0</td></tr><tr><td>Python</td><td>3.10.11 | packaged by conda-forge | (main, May 10 2023, 19:01:19) [Clang 14.0.6 ]</td></tr><tr><td>OS</td><td>posix [darwin]</td></tr><tr><td>Number of CPUs</td><td>Physical: 10, Logical: 10</td></tr><tr><td>BLAS Info</td><td>OPENBLAS</td></tr><tr><td colspan='2'>Wed Sep 27 14:53:58 2023 PDT</td></tr></table>"
       ],
       "text/plain": [
        "<IPython.core.display.HTML object>"

docs/notebooks/quickstart.ipynb

@@ -277,7 +277,7 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "Constructing Voronoi polygons: 100%|█████| 4633/4633 [00:00<00:00, 18399.75it/s]\n"
+      "Constructing Voronoi polygons: 100%|██████████████████████████████████████████████████████████████████████████| 4633/4633 [00:00<00:00, 21408.45it/s]\n"
      ]
     }
    ],
@@ -368,8 +368,8 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "Thermalizing: 100%|████████████████████████▉| 100/100 [00:11<00:00,  9.01tau/s ]\n",
-      "Simulating: 100%|██████████████████████████▉| 150/150 [00:26<00:00,  5.66tau/s ]\n"
+      "Thermalizing: 100%|█████████████████████████████████████████████████████████████████████████████████████████████▉| 100/100 [00:09<00:00, 10.37tau/s ]\n",
+      "Simulating: 100%|███████████████████████████████████████████████████████████████████████████████████████████████▉| 150/150 [00:26<00:00,  5.70tau/s ]\n"
      ]
     }
    ],
@@ -7478,7 +7478,7 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "Simulating: 100%|██████████████████████████▉| 200/200 [00:51<00:00,  3.91tau/s ]\n"
+      "Simulating: 100%|███████████████████████████████████████████████████████████████████████████████████████████████▉| 200/200 [00:50<00:00,  3.94tau/s ]\n"
      ]
     }
    ],
@@ -12256,7 +12256,7 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "Simulating: 100%|██████████████████████████▉| 200/200 [00:40<00:00,  4.96tau/s ]\n"
+      "Simulating: 100%|███████████████████████████████████████████████████████████████████████████████████████████████▉| 200/200 [00:39<00:00,  5.01tau/s ]\n"
      ]
     }
    ],
@@ -25570,7 +25570,7 @@
     {
      "data": {
       "text/html": [
-       "<table><tr><th>Software</th><th>Version</th></tr><tr><td>tdgl</td><td>0.5.1; git revision 03859ea [2023-09-21]</td></tr><tr><td>Numpy</td><td>1.24.3</td></tr><tr><td>SciPy</td><td>1.10.1</td></tr><tr><td>matplotlib</td><td>3.7.1</td></tr><tr><td>jax</td><td>None</td></tr><tr><td>numba</td><td>0.57.1</td></tr><tr><td>IPython</td><td>8.14.0</td></tr><tr><td>Python</td><td>3.10.11 | packaged by conda-forge | (main, May 10 2023, 19:01:19) [Clang 14.0.6 ]</td></tr><tr><td>OS</td><td>posix [darwin]</td></tr><tr><td>Number of CPUs</td><td>Physical: 10, Logical: 10</td></tr><tr><td>BLAS Info</td><td>OPENBLAS</td></tr><tr><td colspan='2'>Thu Sep 21 15:41:40 2023 PDT</td></tr></table>"
+       "<table><tr><th>Software</th><th>Version</th></tr><tr><td>tdgl</td><td>0.6.0; git revision e6d3811 [2023-09-27]</td></tr><tr><td>Numpy</td><td>1.24.3</td></tr><tr><td>SciPy</td><td>1.10.1</td></tr><tr><td>matplotlib</td><td>3.7.1</td></tr><tr><td>cupy</td><td>None</td></tr><tr><td>numba</td><td>0.57.1</td></tr><tr><td>IPython</td><td>8.14.0</td></tr><tr><td>Python</td><td>3.10.11 | packaged by conda-forge | (main, May 10 2023, 19:01:19) [Clang 14.0.6 ]</td></tr><tr><td>OS</td><td>posix [darwin]</td></tr><tr><td>Number of CPUs</td><td>Physical: 10, Logical: 10</td></tr><tr><td>BLAS Info</td><td>OPENBLAS</td></tr><tr><td colspan='2'>Wed Sep 27 15:04:27 2023 PDT</td></tr></table>"
       ],
       "text/plain": [
        "<IPython.core.display.HTML object>"