Merge pull request #600 from mrava87/fix-devitoimport

bug: move _CustomSource to private file

Author: mrava87

pylops/waveeqprocessing/_twoway.py

@@ -0,0 +1,43 @@
+from examples.seismic.utils import PointSource
+
+
+class _CustomSource(PointSource):
+    """Custom source
+
+    This class creates a Devito symbolic object that encapsulates a set of
+    sources with a user defined source signal wavelet ``wav``
+
+    Parameters
+    ----------
+    name : :obj:`str`
+        Name for the resulting symbol.
+    grid : :obj:`devito.types.grid.Grid`
+        The computational domain.
+    time_range : :obj:`examples.seismic.source.TimeAxis`
+        TimeAxis(start, step, num) object.
+    wav : :obj:`numpy.ndarray`
+        Wavelet of size
+
+    """
+
+    __rkwargs__ = PointSource.__rkwargs__ + ["wav"]
+
+    @classmethod
+    def __args_setup__(cls, *args, **kwargs):
+        kwargs.setdefault("npoint", 1)
+
+        return super().__args_setup__(*args, **kwargs)
+
+    def __init_finalize__(self, *args, **kwargs):
+        super().__init_finalize__(*args, **kwargs)
+
+        self.wav = kwargs.get("wav")
+
+        if not self.alias:
+            for p in range(kwargs["npoint"]):
+                self.data[:, p] = self.wavelet
+
+    @property
+    def wavelet(self):
+        """Return user-provided wavelet"""
+        return self.wav

pylops/waveeqprocessing/twoway.py

@@ -1,6 +1,6 @@
 __all__ = ["AcousticWave2D"]
 
-from typing import Tuple
+from typing import Any, NewType, Tuple
 
 import numpy as np
 
@@ -14,49 +14,12 @@
 if devito_message is None:
     from examples.seismic import AcquisitionGeometry, Model
     from examples.seismic.acoustic import AcousticWaveSolver
-    from examples.seismic.utils import PointSource
 
+    from ._twoway import _CustomSource
+else:
+    AcousticWaveSolver = Any
 
-class _CustomSource(PointSource):
-    """Custom source
-
-    This class creates a Devito symbolic object that encapsulates a set of
-    sources with a user defined source signal wavelet ``wav``
-
-    Parameters
-    ----------
-    name : :obj:`str`
-        Name for the resulting symbol.
-    grid : :obj:`devito.types.grid.Grid`
-        The computational domain.
-    time_range : :obj:`examples.seismic.source.TimeAxis`
-        TimeAxis(start, step, num) object.
-    wav : :obj:`numpy.ndarray`
-        Wavelet of size
-
-    """
-
-    __rkwargs__ = PointSource.__rkwargs__ + ["wav"]
-
-    @classmethod
-    def __args_setup__(cls, *args, **kwargs):
-        kwargs.setdefault("npoint", 1)
-
-        return super().__args_setup__(*args, **kwargs)
-
-    def __init_finalize__(self, *args, **kwargs):
-        super().__init_finalize__(*args, **kwargs)
-
-        self.wav = kwargs.get("wav")
-
-        if not self.alias:
-            for p in range(kwargs["npoint"]):
-                self.data[:, p] = self.wavelet
-
-    @property
-    def wavelet(self):
-        """Return user-provided wavelet"""
-        return self.wav
+AcousticWaveSolverType = NewType("AcousticWaveSolver", AcousticWaveSolver)
 
 
 class AcousticWave2D(LinearOperator):
@@ -327,7 +290,7 @@ def srcillumination_allshots(self, savewav: bool = False) -> None:
                 self.src_wavefield.append(src_wav)
             self.src_illumination += src_ill
 
-    def _born_oneshot(self, solver: AcousticWaveSolver, dm: NDArray) -> NDArray:
+    def _born_oneshot(self, solver: AcousticWaveSolverType, dm: NDArray) -> NDArray:
         """Born modelling for one shot
 
         Parameters