From 6216128d25b4d0281ce69948606b9580433e9dcc Mon Sep 17 00:00:00 2001 From: mrava87 Date: Wed, 3 Jul 2024 22:02:26 +0100 Subject: [PATCH] doc: added jax integration to doc --- docs/source/conf.py | 3 + docs/source/gpu.rst | 416 ++++++++++++++++++++++++++++++++++++---- pylops/torchoperator.py | 2 +- tutorials/ilsm.py | 2 +- tutorials/torchop.py | 4 +- 5 files changed, 391 insertions(+), 36 deletions(-) diff --git a/docs/source/conf.py b/docs/source/conf.py index caf745e5..c5e6536d 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -21,6 +21,7 @@ "numpydoc", "nbsphinx", "sphinx_gallery.gen_gallery", + "sphinxemoji.sphinxemoji", # 'sphinx.ext.napoleon', ] @@ -29,6 +30,8 @@ "python": ("https://docs.python.org/3/", None), "numpy": ("https://docs.scipy.org/doc/numpy/", None), "scipy": ("https://docs.scipy.org/doc/scipy/reference", None), + "cupy": ("https://docs.cupy.dev/en/stable/", None), + "jax": ("https://jax.readthedocs.io/en/latest", None), "sklearn": ("http://scikit-learn.org/stable/", None), "pandas": ("http://pandas.pydata.org/pandas-docs/stable/", None), "matplotlib": ("https://matplotlib.org/", None), diff --git a/docs/source/gpu.rst b/docs/source/gpu.rst index e861c567..7890785a 100755 --- a/docs/source/gpu.rst +++ b/docs/source/gpu.rst @@ -1,55 +1,389 @@ .. _gpu: -GPU Support -=========== +GPU / TPU Support +================= Overview -------- -PyLops supports computations on GPUs powered by `CuPy `_ (``cupy-cudaXX>=v13.0.0``). +From ``v1.12.0``, PyLops supports computations on GPUs powered by +`CuPy `_ (``cupy-cudaXX>=8.1.0``). This library must be installed *before* PyLops is installed. -.. note:: +From ``v2.3.0``, PyLops supports also computations on GPUs/TPUs powered by +`JAX `_. +This library must be installed *before* PyLops is installed. - Set environment variable ``CUPY_PYLOPS=0`` to force PyLops to ignore the ``cupy`` backend. - This can be also used if a previous (or faulty) version of ``cupy`` is installed in your system, - otherwise you will get an error when importing PyLops. +PyLops supports computations on GPUs powered by `CuPy `_ (``cupy-cudaXX>=v13.0.0``). +This library must be installed *before* PyLops is installed. +.. note:: + Set environment variables ``CUPY_PYLOPS=0`` and/or ``JAX_PYLOPS=0`` to force PyLops to ignore + ``cupy`` and ``jax`` backends. This can be also used if a previous version of ``cupy`` + or ``jax`` is installed in your system, otherwise you will get an error when importing PyLops. Apart from a few exceptions, all operators and solvers in PyLops can -seamlessly work with ``numpy`` arrays on CPU as well as with ``cupy`` arrays -on GPU. Users do simply need to consistently create operators and +seamlessly work with ``numpy`` arrays on CPU as well as with ``cupy/jax`` arrays +on GPU. For CuPy, users simply need to consistently create operators and provide data vectors to the solvers, e.g., when using :class:`pylops.MatrixMult` the input matrix must be a ``cupy`` array if the data provided to a solver is also ``cupy`` array. +For JAX, apart from following the procedure described for CuPy, a PyLops operator must also +be wrapped into a ``JaxOperator``. -.. warning:: - Some :class:`pylops.LinearOperator` methods are currently on GPU: +In the following, we provide a list of methods in :class:`pylops.LinearOperator` with their current status (available on CPU, +GPU with CuPy, and GPU with JAX. - - :meth:`pylops.LinearOperator.eigs` - - :meth:`pylops.LinearOperator.cond` - - :meth:`pylops.LinearOperator.tosparse` - - :meth:`pylops.LinearOperator.estimate_spectral_norm` +.. list-table:: + :widths: 50 25 25 25 + :header-rows: 1 -.. warning:: + * - Operator/method + - CPU + - GPU with CuPy + - GPU/TPU with JAX + * - :meth:`pylops.LinearOperator.eigs` + - |:white_check_mark:| + - |:red_circle:| + - |:red_circle:| + * - :meth:`pylops.LinearOperator.cond` + - |:white_check_mark:| + - |:red_circle:| + - |:red_circle:| + * - :meth:`pylops.LinearOperator.cond` + - |:white_check_mark:| + - |:red_circle:| + - |:red_circle:| + +Basic operators: + +.. list-table:: + :widths: 50 25 25 25 + :header-rows: 1 + + * - Operator/method + - CPU + - GPU with CuPy + - GPU/TPU with JAX + * - :class:`pylops.basicoperators.MatrixMult` + - |:white_check_mark:| + - |:white_check_mark:| + - |:white_check_mark:| + * - :class:`pylops.basicoperators.Identity` + - |:white_check_mark:| + - |:white_check_mark:| + - |:white_check_mark:| + * - :class:`pylops.basicoperators.Zero` + - |:white_check_mark:| + - |:white_check_mark:| + - |:white_check_mark:| + * - :class:`pylops.basicoperators.Diagonal` + - |:white_check_mark:| + - |:white_check_mark:| + - |:white_check_mark:| + * - :meth:`pylops.basicoperators.Transpose` + - |:white_check_mark:| + - |:white_check_mark:| + - |:white_check_mark:| + * - :class:`pylops.basicoperators.Flip` + - |:white_check_mark:| + - |:white_check_mark:| + - |:white_check_mark:| + * - :class:`pylops.basicoperators.Roll` + - |:white_check_mark:| + - |:white_check_mark:| + - |:white_check_mark:| + * - :class:`pylops.basicoperators.Pad` + - |:white_check_mark:| + - |:white_check_mark:| + - |:white_check_mark:| + * - :class:`pylops.basicoperators.Sum` + - |:white_check_mark:| + - |:white_check_mark:| + - |:white_check_mark:| + * - :class:`pylops.basicoperators.Symmetrize` + - |:white_check_mark:| + - |:white_check_mark:| + - |:white_check_mark:| + * - :class:`pylops.basicoperators.Restriction` + - |:white_check_mark:| + - |:white_check_mark:| + - |:white_check_mark:| + * - :class:`pylops.basicoperators.Regression` + - |:white_check_mark:| + - |:white_check_mark:| + - |:white_check_mark:| + * - :class:`pylops.basicoperators.LinearRegression` + - |:white_check_mark:| + - |:white_check_mark:| + - |:white_check_mark:| + * - :class:`pylops.basicoperators.CausalIntegration` + - |:white_check_mark:| + - |:white_check_mark:| + - |:white_check_mark:| + * - :class:`pylops.basicoperators.Spread` + - |:white_check_mark:| + - |:red_circle:| + - |:red_circle:| + * - :class:`pylops.basicoperators.VStack` + - |:white_check_mark:| + - |:white_check_mark:| + - |:white_check_mark:| + * - :class:`pylops.basicoperators.HStack` + - |:white_check_mark:| + - |:white_check_mark:| + - |:white_check_mark:| + * - :class:`pylops.basicoperators.Block` + - |:white_check_mark:| + - |:white_check_mark:| + - |:white_check_mark:| + * - :class:`pylops.basicoperators.BlockDiag` + - |:white_check_mark:| + - |:white_check_mark:| + - |:white_check_mark:| + + +Smoothing and derivatives: + +.. list-table:: + :widths: 50 25 25 25 + :header-rows: 1 + + * - Operator/method + - CPU + - GPU with CuPy + - GPU/TPU with JAX + * - :class:`pylops.basicoperators.FirstDerivative` + - |:white_check_mark:| + - |:white_check_mark:| + - |:white_check_mark:| + * - :class:`pylops.basicoperators.SecondDerivative` + - |:white_check_mark:| + - |:white_check_mark:| + - |:white_check_mark:| + * - :class:`pylops.basicoperators.Laplacian` + - |:white_check_mark:| + - |:white_check_mark:| + - |:white_check_mark:| + * - :class:`pylops.basicoperators.Gradient` + - |:white_check_mark:| + - |:white_check_mark:| + - |:white_check_mark:| + * - :class:`pylops.basicoperators.FirstDirectionalDerivative` + - |:white_check_mark:| + - |:white_check_mark:| + - |:white_check_mark:| + * - :class:`pylops.basicoperators.SecondDirectionalDerivative` + - |:white_check_mark:| + - |:white_check_mark:| + - |:white_check_mark:| + +Signal processing: + +.. list-table:: + :widths: 50 25 25 25 + :header-rows: 1 + + * - Operator/method + - CPU + - GPU with CuPy + - GPU/TPU with JAX + * - :class:`pylops.signalprocessing.Convolve1D` + - |:white_check_mark:| + - |:white_check_mark:| + - |:warning:| + * - :class:`pylops.signalprocessing.Convolve2D` + - |:white_check_mark:| + - |:white_check_mark:| + - |:white_check_mark:| + * - :class:`pylops.signalprocessing.ConvolveND` + - |:white_check_mark:| + - |:white_check_mark:| + - |:white_check_mark:| + * - :class:`pylops.signalprocessing.NonStationaryConvolve1D` + - |:white_check_mark:| + - |:white_check_mark:| + - |:white_check_mark:| + * - :class:`pylops.signalprocessing.NonStationaryFilters1D` + - |:white_check_mark:| + - |:white_check_mark:| + - |:white_check_mark:| + * - :class:`pylops.signalprocessing.NonStationaryConvolve2D` + - |:white_check_mark:| + - |:white_check_mark:| + - |:red_circle:| + * - :class:`pylops.signalprocessing.NonStationaryFilters2D` + - |:white_check_mark:| + - |:white_check_mark:| + - |:red_circle:| + * - :class:`pylops.signalprocessing.Interp` + - |:white_check_mark:| + - |:white_check_mark:| + - |:white_check_mark:| + * - :class:`pylops.signalprocessing.Bilinear` + - |:white_check_mark:| + - |:white_check_mark:| + - |:red_circle:| + * - :class:`pylops.signalprocessing.FFT` + - |:white_check_mark:| + - |:white_check_mark:| + - |:white_check_mark:| + * - :class:`pylops.signalprocessing.FFT2D` + - |:white_check_mark:| + - |:white_check_mark:| + - |:white_check_mark:| + * - :class:`pylops.signalprocessing.FFTND` + - |:white_check_mark:| + - |:white_check_mark:| + - |:white_check_mark:| + * - :class:`pylops.signalprocessing.Shift` + - |:white_check_mark:| + - |:white_check_mark:| + - |:white_check_mark:| + * - :class:`pylops.signalprocessing.DWT` + - |:white_check_mark:| + - |:red_circle:| + - |:red_circle:| + * - :class:`pylops.signalprocessing.DWT2D` + - |:white_check_mark:| + - |:red_circle:| + - |:red_circle:| + * - :class:`pylops.signalprocessing.DCT` + - |:white_check_mark:| + - |:red_circle:| + - |:red_circle:| + * - :class:`pylops.signalprocessing.Seislet` + - |:white_check_mark:| + - |:red_circle:| + - |:red_circle:| + * - :class:`pylops.signalprocessing.Radon2D` + - |:white_check_mark:| + - |:red_circle:| + - |:red_circle:| + * - :class:`pylops.signalprocessing.Radon3D` + - |:white_check_mark:| + - |:red_circle:| + - |:red_circle:| + * - :class:`pylops.signalprocessing.ChirpRadon2D` + - |:white_check_mark:| + - |:white_check_mark:| + - |:red_circle:| + * - :class:`pylops.signalprocessing.ChirpRadon3D` + - |:white_check_mark:| + - |:white_check_mark:| + - |:red_circle:| + * - :class:`pylops.signalprocessing.Sliding1D` + - |:white_check_mark:| + - |:white_check_mark:| + - |:white_check_mark:| + * - :class:`pylops.signalprocessing.Sliding2D` + - |:white_check_mark:| + - |:white_check_mark:| + - |:white_check_mark:| + * - :class:`pylops.signalprocessing.Sliding3D` + - |:white_check_mark:| + - |:white_check_mark:| + - |:white_check_mark:| + * - :class:`pylops.signalprocessing.Patch2D` + - |:white_check_mark:| + - |:white_check_mark:| + - |:white_check_mark:| + * - :class:`pylops.signalprocessing.Patch3D` + - |:white_check_mark:| + - |:white_check_mark:| + - |:white_check_mark:| + * - :class:`pylops.signalprocessing.Fredholm1` + - |:white_check_mark:| + - |:white_check_mark:| + - |:white_check_mark:| - Some operators are currently not available on GPU: +Wave-Equation processing - - :class:`pylops.Spread` - - :class:`pylops.signalprocessing.Radon2D` - - :class:`pylops.signalprocessing.Radon3D` - - :class:`pylops.signalprocessing.DWT` - - :class:`pylops.signalprocessing.DWT2D` - - :class:`pylops.signalprocessing.Seislet` - - :class:`pylops.waveeqprocessing.Demigration` - - :class:`pylops.waveeqprocessing.LSM` +.. list-table:: + :widths: 50 25 25 25 + :header-rows: 1 + + * - Operator/method + - CPU + - GPU with CuPy + - GPU/TPU with JAX + * - :class:`pylops.avo.avo.PressureToVelocity` + - |:white_check_mark:| + - |:white_check_mark:| + - |:white_check_mark:| + * - :class:`pylops.avo.avo.UpDownComposition2D` + - |:white_check_mark:| + - |:white_check_mark:| + - |:white_check_mark:| + * - :class:`pylops.avo.avo.UpDownComposition3D` + - |:white_check_mark:| + - |:white_check_mark:| + - |:white_check_mark:| + * - :class:`pylops.avo.avo.BlendingContinuous` + - |:white_check_mark:| + - |:white_check_mark:| + - |:white_check_mark:| + * - :class:`pylops.avo.avo.BlendingGroup` + - |:white_check_mark:| + - |:white_check_mark:| + - |:white_check_mark:| + * - :class:`pylops.avo.avo.BlendingHalf` + - |:white_check_mark:| + - |:white_check_mark:| + - |:white_check_mark:| + * - :class:`pylops.avo.avo.MDC` + - |:white_check_mark:| + - |:white_check_mark:| + - |:white_check_mark:| + * - :class:`pylops.avo.avo.Kirchhoff` + - |:white_check_mark:| + - |:red_circle:| + - |:red_circle:| + * - :class:`pylops.avo.avo.AcousticWave2D` + - |:white_check_mark:| + - |:red_circle:| + - |:red_circle:| + +Geophysical subsurface characterization: + +.. list-table:: + :widths: 50 25 25 25 + :header-rows: 1 + + * - Operator/method + - CPU + - GPU with CuPy + - GPU/TPU with JAX + * - :class:`pylops.avo.avo.AVOLinearModelling` + - |:white_check_mark:| + - |:white_check_mark:| + - |:white_check_mark:| + * - :class:`pylops.avo.poststack.PoststackLinearModelling` + - |:white_check_mark:| + - |:white_check_mark:| + - |:white_check_mark:| + * - :class:`pylops.avo.prestack.PrestackLinearModelling` + - |:white_check_mark:| + - |:white_check_mark:| + - |:warning:| + * - :class:`pylops.avo.prestack.PrestackWaveletModelling` + - |:white_check_mark:| + - |:white_check_mark:| + - |:warning:| .. warning:: - Some solvers are currently not available on GPU: - - :class:`pylops.optimization.sparsity.SPGL1` + 1. The JAX backend of the :class:`pylops.basicoperators.Convolve1D` operator + currently works only with 1d-arrays due to a different behaviour of + :meth:`scipy.signal.convolve` and :meth:`jax.scipy.signal.convolve` with + nd-arrays. + + 2. The JAX backend of the :class:`pylops.avo.prestack.PrestackLinearModelling` + operator currently works only with ``explicit=True`` due to the same issue as + in point 1 for the :class:`pylops.basicoperators.Convolve1D` operator employed + when ``explicit=False``. Example @@ -68,8 +402,7 @@ Finally, let's briefly look at an example. First we write a code snippet using y = Gop * x xest = Gop / y - -Now we write a code snippet using ``cupy`` arrays which PyLops will run on +Now we write a code snippet using ``cupy`` arrays which PyLops will run on your GPU: .. code-block:: python @@ -83,9 +416,28 @@ your GPU: xest = Gop / y The code is almost unchanged apart from the fact that we now use ``cupy`` arrays, -PyLops will figure this out! +PyLops will figure this out. + +Similarly, we write a code snippet using ``jax`` arrays which PyLops will run on +your GPU/TPU: + +.. code-block:: python + + ny, nx = 400, 400 + G = jnp.array(np.random.normal(0, 1, (ny, nx)).astype(np.float32)) + x = jnp.ones(nx, dtype=np.float32) + + Gop = JaxOperator(MatrixMult(G, dtype='float32')) + y = Gop * x + xest = Gop / y + + # Adjoint via AD + xadj = Gop.rmatvecad(x, y) + + +Again, the code is almost unchanged apart from the fact that we now use ``jax`` arrays, .. note:: - The CuPy backend is in active development, with many examples not yet in the docs. - You can find many `other examples `_ from the `PyLops Notebooks repository `_. + More examples for the CuPy and JAX backends be found `here `_ + and `here `_. \ No newline at end of file diff --git a/pylops/torchoperator.py b/pylops/torchoperator.py index 5e41f67f..1c4dc2da 100644 --- a/pylops/torchoperator.py +++ b/pylops/torchoperator.py @@ -14,7 +14,7 @@ else: torch_message = ( "Torch package not installed. In order to be able to use" - 'the twoway module run "pip install torch" or' + 'the torchoperator module run "pip install torch" or' '"conda install -c pytorch torch".' ) from pylops.utils.typing import TensorTypeLike diff --git a/tutorials/ilsm.py b/tutorials/ilsm.py index b4b016f3..1f394bfc 100755 --- a/tutorials/ilsm.py +++ b/tutorials/ilsm.py @@ -1,5 +1,5 @@ r""" -20. Image Domain Least-squares migration +19. Image Domain Least-squares migration ======================================== Seismic migration is the process by which seismic data are manipulated to create an image of the subsurface reflectivity. diff --git a/tutorials/torchop.py b/tutorials/torchop.py index c555573d..9e73d7b3 100755 --- a/tutorials/torchop.py +++ b/tutorials/torchop.py @@ -1,6 +1,6 @@ r""" -19. Automatic Differentiation -============================= +20. Torch Operator +================== This tutorial focuses on the use of :class:`pylops.TorchOperator` to allow performing Automatic Differentiation (AD) on chains of operators which can be: