ci: modified test suite to run on cupy

Makefile

@@ -1,7 +1,7 @@
 PIP := $(shell command -v pip3 2> /dev/null || command which pip 2> /dev/null)
 PYTHON := $(shell command -v python3 2> /dev/null || command which python 2> /dev/null)
 
-.PHONY: install dev-install install_conda dev-install_conda tests doc docupdate servedoc lint typeannot coverage
+.PHONY: install dev-install dev-install_gpu install_conda dev-install_conda dev-install_conda_arm tests doc docupdate servedoc lint typeannot coverage
 
 pipcheck:
 ifndef PIP
@@ -24,6 +24,11 @@ dev-install:
 	$(PIP) install -r requirements-dev.txt &&\
 	$(PIP) install -r requirements-torch.txt && $(PIP) install -e .
 
+dev-install_gpu:
+	make pipcheck
+	$(PIP) install -r requirements-dev-gpu.txt &&\
+	$(PIP) install -e .
+
 install_conda:
 	conda env create -f environment.yml && conda activate pylops && pip install .
 
@@ -33,6 +38,9 @@ dev-install_conda:
 dev-install_conda_arm:
 	conda env create -f environment-dev-arm.yml && conda activate pylops && pip install -e .
 
+dev-install_conda_gpu:
+	conda env create -f environment-dev-gpu.yml && conda activate pylops_gpu && pip install -e .
+
 tests:
 	make pythoncheck
 	pytest

environment-dev-arm.yml

@@ -5,7 +5,7 @@ channels:
   - numba
   - pytorch
 dependencies:
-  - python>=3.6.4
+  - python>=3.9.0
   - pip
   - numpy>=1.21.0
   - scipy>=1.11.0

environment-dev-gpu.yml

@@ -0,0 +1,34 @@
+name: pylops_gpu
+channels:
+  - conda-forge
+  - defaults
+  - numba
+dependencies:
+  - python>=3.11.0
+  - pip
+  - numpy>=1.21.0
+  - scipy>=1.11.0
+  - cupy
+  - sympy
+  - matplotlib
+  - ipython
+  - pytest
+  - Sphinx
+  - numpydoc
+  - numba
+  - icc_rt
+  - pre-commit
+  - autopep8
+  - isort
+  - black
+  - pip:
+    - torch
+    - pytest-runner
+    - setuptools_scm
+    - pydata-sphinx-theme
+    - sphinx-gallery
+    - nbsphinx
+    - sphinxemoji
+    - image
+    - flake8
+    - mypy

environment-dev.yml

@@ -5,7 +5,7 @@ channels:
   - numba
   - pytorch
 dependencies:
-  - python>=3.6.4
+  - python>=3.9.0
   - pip
   - numpy>=1.21.0
   - scipy>=1.11.0

environment.yml

@@ -2,6 +2,6 @@ name: pylops
 channels:
   - defaults
 dependencies:
-  - python>=3.6.4
+  - python>=3.9.0
   - numpy>=1.21.0
   - scipy>=1.14.0

pytests/test_avo.py

@@ -1,8 +1,18 @@
-import numpy as np
+import os
+
+if int(os.environ.get("TEST_CUPY_PYLOPS", 0)):
+    import cupy as np
+    from cupy.testing import assert_array_almost_equal
+
+    backend = "cupy"
+else:
+    import numpy as np
+    from numpy.testing import assert_array_almost_equal
+
+    backend = "numpy"
+import numpy as npp
 import pytest
-from numpy.testing import assert_array_almost_equal
 from scipy.signal import filtfilt
-from scipy.sparse.linalg import lsqr
 
 from pylops.avo.avo import (
     akirichards,
@@ -13,7 +23,9 @@
     zoeppritz_scattering,
 )
 from pylops.avo.prestack import AVOLinearModelling
+from pylops.optimization.basic import lsqr
 from pylops.utils import dottest
+from pylops.utils.backend import to_numpy
 
 np.random.seed(0)
 
@@ -24,16 +36,20 @@
 # Create medium parameters for multiple contrasts
 nt0 = 201
 dt0 = 0.004
-t0 = np.arange(nt0) * dt0
-vp = 1200 + np.arange(nt0) + filtfilt(np.ones(5) / 5.0, 1, np.random.normal(0, 80, nt0))
-vs = 600 + vp / 2 + filtfilt(np.ones(5) / 5.0, 1, np.random.normal(0, 20, nt0))
-rho = 1000 + vp + filtfilt(np.ones(5) / 5.0, 1, np.random.normal(0, 30, nt0))
-m = (np.stack((np.log(vp), np.log(vs), np.log(rho)), axis=1)).ravel()
+t0 = npp.arange(nt0) * dt0
+vp = (
+    1200
+    + npp.arange(nt0)
+    + filtfilt(npp.ones(5) / 5.0, 1, npp.random.normal(0, 80, nt0))
+)
+vs = 600 + vp / 2 + filtfilt(npp.ones(5) / 5.0, 1, npp.random.normal(0, 20, nt0))
+rho = 1000 + vp + filtfilt(npp.ones(5) / 5.0, 1, npp.random.normal(0, 30, nt0))
+m = npp.stack((npp.log(vp), npp.log(vs), npp.log(rho)), axis=1).ravel()
 
 # Angles
 ntheta = 21
 thetamin, thetamax = 0, 40
-theta = np.linspace(thetamin, thetamax, ntheta)
+theta = npp.linspace(thetamin, thetamax, ntheta)
 
 # Parameters
 par1 = {"vsvp": 0.5, "linearization": "akirich"}  # constant vsvp
@@ -47,16 +63,16 @@ def test_zoeppritz():
     `<https://www.crewes.org/ResearchLinks/ExplorerPrograms/ZE/index.html>`_
     as benchmark
     """
-    r_zoep = zoeppritz_scattering(vp1, vs1, rho1, vp0, vs0, rho0, theta[0])
+    r_zoep = zoeppritz_scattering(vp1, vs1, rho1, vp0, vs0, rho0, np.asarray(theta)[:1])
     rpp_zoep = zoeppritz_element(
-        vp1, vs1, rho1, vp0, vs0, rho0, theta[0], element="PdPu"
+        vp1, vs1, rho1, vp0, vs0, rho0, np.asarray(theta)[:1], element="PdPu"
     )
-    rpp_zoep1 = zoeppritz_pp(vp1, vs1, rho1, vp0, vs0, rho0, theta[0])
+    rpp_zoep1 = zoeppritz_pp(vp1, vs1, rho1, vp0, vs0, rho0, np.asarray(theta)[:1])
 
     assert r_zoep.shape == (4, 4, 1)
-    assert r_zoep[0, 0] == pytest.approx(0.04658, rel=1e-3)
-    assert rpp_zoep == pytest.approx(0.04658, rel=1e-3)
-    assert rpp_zoep1 == pytest.approx(0.04658, rel=1e-3)
+    assert to_numpy(r_zoep[0, 0, 0]) == pytest.approx(0.04658, rel=1e-3)
+    assert to_numpy(rpp_zoep) == pytest.approx(0.04658, rel=1e-3)
+    assert to_numpy(rpp_zoep1) == pytest.approx(0.04658, rel=1e-3)
 
 
 def test_zoeppritz_and_approx_zeroangle():
@@ -66,22 +82,24 @@ def test_zoeppritz_and_approx_zeroangle():
     ai1, si1, _ = vp1 * rho1, vs1 * rho1, vp1 / vs1
 
     # Zoeppritz
-    rpp_zoep = zoeppritz_pp(vp1, vs1, rho1, vp0, vs0, rho0, theta[0])
-    rpp_zoep_approx = approx_zoeppritz_pp(vp1, vs1, rho1, vp0, vs0, rho0, theta[0])
+    rpp_zoep = zoeppritz_pp(vp1, vs1, rho1, vp0, vs0, rho0, np.asarray(theta)[:1])
+    rpp_zoep_approx = approx_zoeppritz_pp(
+        vp1, vs1, rho1, vp0, vs0, rho0, np.asarray(theta)[:1]
+    )
 
     # Aki Richards
-    rvp = np.log(vp0) - np.log(vp1)
-    rvs = np.log(vs0) - np.log(vs1)
-    rrho = np.log(rho0) - np.log(rho1)
+    rvp = np.asarray(np.log(vp0) - np.log(vp1))
+    rvs = np.asarray(np.log(vs0) - np.log(vs1))
+    rrho = np.asarray(np.log(rho0) - np.log(rho1))
 
-    G1, G2, G3 = akirichards(theta[0], vs1 / vp1)
+    G1, G2, G3 = akirichards(np.asarray(theta)[:1], vs1 / vp1)
     rpp_aki = G1 * rvp + G2 * rvs + G3 * rrho
 
     # Fatti
-    rai = np.log(ai0) - np.log(ai1)
-    rsi = np.log(si0) - np.log(si1)
+    rai = np.asarray(np.log(ai0) - np.log(ai1))
+    rsi = np.asarray(np.log(si0) - np.log(si1))
 
-    G1, G2, G3 = fatti(theta[0], vs1 / vp1)
+    G1, G2, G3 = fatti(np.asarray(theta)[:1], vs1 / vp1)
     rpp_fatti = G1 * rai + G2 * rsi + G3 * rrho
 
     assert_array_almost_equal(rpp_zoep, rpp_zoep_approx, decimal=3)
@@ -97,22 +115,24 @@ def test_zoeppritz_and_approx_multipleangles():
     ai1, si1 = vp1 * rho1, vs1 * rho1
 
     # Zoeppritz
-    rpp_zoep = zoeppritz_pp(vp1, vs1, rho1, vp0, vs0, rho0, theta)
-    rpp_zoep_approx = approx_zoeppritz_pp(vp1, vs1, rho1, vp0, vs0, rho0, theta)
+    rpp_zoep = zoeppritz_pp(vp1, vs1, rho1, vp0, vs0, rho0, np.asarray(theta))
+    rpp_zoep_approx = approx_zoeppritz_pp(
+        vp1, vs1, rho1, vp0, vs0, rho0, np.asarray(theta)
+    )
 
     # Aki Richards
-    rvp = np.log(vp0) - np.log(vp1)
-    rvs = np.log(vs0) - np.log(vs1)
-    rrho = np.log(rho0) - np.log(rho1)
+    rvp = np.asarray(np.log(vp0) - np.log(vp1))
+    rvs = np.asarray(np.log(vs0) - np.log(vs1))
+    rrho = np.asarray(np.log(rho0) - np.log(rho1))
 
-    G1, G2, G3 = akirichards(theta, vs1 / vp1)
+    G1, G2, G3 = akirichards(np.asarray(theta), vs1 / vp1)
     rpp_aki = G1 * rvp + G2 * rvs + G3 * rrho
 
     # Fatti
-    rai = np.log(ai0) - np.log(ai1)
-    rsi = np.log(si0) - np.log(si1)
+    rai = np.asarray(np.log(ai0) - np.log(ai1))
+    rsi = np.asarray(np.log(si0) - np.log(si1))
 
-    G1, G2, G3 = fatti(theta, vs1 / vp1)
+    G1, G2, G3 = fatti(np.asarray(theta), vs1 / vp1)
     rpp_fatti = G1 * rai + G2 * rsi + G3 * rrho
 
     assert_array_almost_equal(rpp_zoep, rpp_zoep_approx, decimal=3)
@@ -124,11 +144,21 @@ def test_zoeppritz_and_approx_multipleangles():
 def test_AVOLinearModelling(par):
     """Dot-test and inversion for AVOLinearModelling"""
     AVOop = AVOLinearModelling(
-        theta, vsvp=par["vsvp"], nt0=nt0, linearization=par["linearization"]
+        np.asarray(theta),
+        vsvp=par["vsvp"] if isinstance(par["vsvp"], float) else np.asarray(par["vsvp"]),
+        nt0=nt0,
+        linearization=par["linearization"],
     )
-    assert dottest(AVOop, ntheta * nt0, 3 * nt0)
+    assert dottest(AVOop, ntheta * nt0, 3 * nt0, backend=backend)
 
     minv = lsqr(
-        AVOop, AVOop * m, damp=1e-20, iter_lim=1000, atol=1e-8, btol=1e-8, show=0
+        AVOop,
+        AVOop * np.asarray(m),
+        x0=np.zeros_like(m),
+        damp=1e-20,
+        niter=1000,
+        atol=1e-8,
+        btol=1e-8,
+        show=0,
     )[0]
     assert_array_almost_equal(m, minv, decimal=3)