Merge pull request #681 from GeophyAI/fix_jax

Support JAX backend in ChirpRadon2D/3D by removing in-place assignment.

Author: mrava87

README.md

@@ -152,3 +152,4 @@ A list of video tutorials to learn more about PyLops:
 * Alex Rakowski, alex-rakowski
 * David Sollberger, solldavid
 * Gustavo Coelho, guaacoelho
+* Shaowen Wang, GeophyAI

docs/source/credits.rst

@@ -23,4 +23,5 @@ Contributors
 *  `Fedor Goncharov <https://github.com/fedor-goncharov>`_, fedor-goncharov
 *  `Alex Rakowski <https://github.com/alex-rakowski>`_, alex-rakowski
 *  `David Sollberger <https://github.com/solldavid>`_, solldavid
-*  `Gustavo Coelho <https://github.com/guaacoelho>`_, guaacoelho
+*  `Gustavo Coelho <https://github.com/guaacoelho>`_, guaacoelho
+*  `Shaowen Wang <https://github.com/GeophyAI>`_, GeophyAI

docs/source/gpu.rst

@@ -473,11 +473,11 @@ Signal processing:
    * - :class:`pylops.signalprocessing.ChirpRadon2D`
      - |:white_check_mark:|
      - |:white_check_mark:|
-     - |:red_circle:|
+     - |:white_check_mark:|
    * - :class:`pylops.signalprocessing.ChirpRadon3D`
      - |:white_check_mark:|
      - |:white_check_mark:|
-     - |:red_circle:|
+     - |:white_check_mark:|
    * - :class:`pylops.signalprocessing.Sliding1D`
      - |:white_check_mark:|
      - |:white_check_mark:|

pylops/signalprocessing/_chirpradon2d.py

@@ -1,4 +1,4 @@
-from pylops.utils.backend import get_array_module
+from pylops.utils.backend import get_array_module, inplace_set
 from pylops.utils.typing import NDArray
 
 
@@ -62,7 +62,7 @@ def _chirp_radon_2d(
 
     # perform transform
     h = ncp.zeros((2 * nx, nt)).astype(cdtype)
-    h[0:nx, :] = ncp.fft.fftn(data, axes=(1,)) * K
+    h = inplace_set(ncp.fft.fftn(data, axes=(1,)) * K, h, (slice(0, nx), slice(None)))
     g = ncp.fft.ifftn(
         ncp.fft.fftn(h, axes=(0,)) * ncp.fft.fftn(K0, axes=(0,)), axes=(0,)
     )

pylops/signalprocessing/_chirpradon3d.py

@@ -1,6 +1,6 @@
 import numpy as np
 
-from pylops.utils.backend import get_array_module
+from pylops.utils.backend import get_array_module, inplace_set
 from pylops.utils.typing import NDArray
 
 try:
@@ -88,8 +88,11 @@ def _chirp_radon_3d(
 
     # perform transform
     h = ncp.zeros((2 * ny, 2 * nx, nt)).astype(cdtype)
-    h[0:ny, 0:nx, :] = ncp.fft.fftn(data, axes=(2,)) * K1 * K2
-
+    h = inplace_set(
+        ncp.fft.fftn(data, axes=(2,)) * K1 * K2,
+        h,
+        (slice(0, ny), slice(0, nx), slice(None)),
+    )
     g = ncp.fft.ifftn(
         ncp.fft.fftn(h, axes=(1,)) * ncp.fft.fftn(K01, axes=(1,)), axes=(1,)
     )