Merge pull request #650 from guaacoelho/cache_enhancement

Removes the creation of the solver from within the shots iteration loop.

Author: mrava87

pylops/waveeqprocessing/twoway.py

@@ -229,11 +229,13 @@ def updatesrc(self, wav):
             time_range=self.geometry.time_axis,
         )
 
-    def _srcillumination_oneshot(self, isrc: int) -> Tuple[NDArray, NDArray]:
+    def _srcillumination_oneshot(self, solver: AcousticWaveSolverType, isrc: int) -> Tuple[NDArray, NDArray]:
         """Source wavefield and illumination for one shot
 
         Parameters
         ----------
+        solver : :obj:`AcousticWaveSolver`
+            Devito's solver object.
         isrc : :obj:`int`
             Index of source to model
 
@@ -245,21 +247,10 @@ def _srcillumination_oneshot(self, isrc: int) -> Tuple[NDArray, NDArray]:
             Source illumination
 
         """
-        # create geometry for single source
-        geometry = AcquisitionGeometry(
-            self.model,
-            self.geometry.rec_positions,
-            self.geometry.src_positions[isrc, :],
-            self.geometry.t0,
-            self.geometry.tn,
-            f0=self.geometry.f0,
-            src_type=self.geometry.src_type,
-        )
-        solver = AcousticWaveSolver(self.model, geometry, space_order=self.space_order)
 
         # assign source location to source object with custom wavelet
         if hasattr(self, "wav"):
-            self.wav.coordinates.data[0, :] = self.geometry.src_positions[isrc, :]
+            self.wav.coordinates.data[0, :] = solver.geometry.src_positions[:]
 
         # source wavefield
         u0 = solver.forward(
@@ -279,13 +270,27 @@ def srcillumination_allshots(self, savewav: bool = False) -> None:
             Save source wavefield (``True``) or not (``False``)
 
         """
+        # create geometry for single source
+        geometry = AcquisitionGeometry(
+            self.model,
+            self.geometry.rec_positions,
+            self.geometry.src_positions[0, :],
+            self.geometry.t0,
+            self.geometry.tn,
+            f0=self.geometry.f0,
+            src_type=self.geometry.src_type,
+        )
+
+        solver = AcousticWaveSolver(self.model, geometry, space_order=self.space_order)
+
         nsrc = self.geometry.src_positions.shape[0]
         if savewav:
             self.src_wavefield = []
         self.src_illumination = np.zeros(self.model.shape)
 
         for isrc in range(nsrc):
-            src_wav, src_ill = self._srcillumination_oneshot(isrc)
+            solver.geometry.src_positions = self.geometry.src_positions[isrc, :]
+            src_wav, src_ill = self._srcillumination_oneshot(solver, isrc)
             if savewav:
                 self.src_wavefield.append(src_wav)
             self.src_illumination += src_ill
@@ -359,11 +364,13 @@ def _born_allshots(self, dm: NDArray) -> NDArray:
         dtot = np.array(dtot).reshape(nsrc, d.shape[0], d.shape[1])
         return dtot
 
-    def _bornadj_oneshot(self, isrc, dobs):
+    def _bornadj_oneshot(self, solver: AcousticWaveSolverType, isrc, dobs):
         """Adjoint born modelling for one shot
 
         Parameters
         ----------
+        solver : :obj:`AcousticWaveSolver`
+            Devito's solver object.
         isrc : :obj:`float`
             Index of source to model
         dobs : :obj:`np.ndarray`
@@ -375,25 +382,13 @@ def _bornadj_oneshot(self, isrc, dobs):
             Model
 
         """
-        # create geometry for single source
-        geometry = AcquisitionGeometry(
-            self.model,
-            self.geometry.rec_positions,
-            self.geometry.src_positions[isrc, :],
-            self.geometry.t0,
-            self.geometry.tn,
-            f0=self.geometry.f0,
-            src_type=self.geometry.src_type,
-        )
         # create boundary data
         recs = self.geometry.rec.copy()
         recs.data[:] = dobs.T[:]
 
-        solver = AcousticWaveSolver(self.model, geometry, space_order=self.space_order)
-
         # assign source location to source object with custom wavelet
         if hasattr(self, "wav"):
-            self.wav.coordinates.data[0, :] = self.geometry.src_positions[isrc, :]
+            self.wav.coordinates.data[0, :] = solver.geometry.src_positions[:]
 
         # source wavefield
         if hasattr(self, "src_wavefield"):
@@ -423,11 +418,25 @@ def _bornadj_allshots(self, dobs: NDArray) -> NDArray:
             Model
 
         """
+        # create geometry for single source
+        geometry = AcquisitionGeometry(
+            self.model,
+            self.geometry.rec_positions,
+            self.geometry.src_positions[0, :],
+            self.geometry.t0,
+            self.geometry.tn,
+            f0=self.geometry.f0,
+            src_type=self.geometry.src_type,
+        )
+
         nsrc = self.geometry.src_positions.shape[0]
         mtot = np.zeros(self.model.shape, dtype=np.float32)
 
+        solver = AcousticWaveSolver(self.model, geometry, space_order=self.space_order)
+
         for isrc in range(nsrc):
-            m = self._bornadj_oneshot(isrc, dobs[isrc])
+            solver.geometry.src_positions = self.geometry.src_positions[isrc, :]
+            m = self._bornadj_oneshot(solver, isrc, dobs[isrc])
             mtot += self._crop_model(m.data, self.model.nbl)
         return mtot