GEOPY-2910: b

simpeg/dask/electromagnetics/time_domain/simulation.py

@@ -88,17 +88,16 @@ def compute_J(self, m, f=None):
         f, Ainv = self.fields(m=m, return_Ainv=True)
 
     client, worker = self._get_client_worker()
-
     ftype = self._fieldType + "Solution"
     n_cells = m.size
-
     simulation_times = np.r_[0, np.cumsum(self.time_steps)] + self.t0
     data_times = self.survey.source_list[0].receiver_list[0].times
     compute_row_size = np.ceil(self.max_chunk_size / (m.shape[0] * 8.0 * 1e-6))
     blocks = get_parallel_blocks(
         self.survey.source_list,
         compute_row_size,
         thread_count=self.n_threads(client=client, worker=worker),
+        optimize=True,
     )
     fields_array = f[:, ftype, :]
 
@@ -156,7 +155,6 @@ def compute_J(self, m, f=None):
                 AdiagTinv,
                 ATinv_df_duT_v[ind],
                 time_mask,
-                client,
             )
 
         if client:
@@ -166,7 +164,7 @@ def compute_J(self, m, f=None):
 
         for block_ind in range(len(blocks)):
 
-            if len(block) == 0:
+            if len(blocks[block_ind]) == 0:
                 continue
 
             if client:
@@ -337,48 +335,53 @@ def get_field_deriv_block(
     AdiagTinv,
     ATinv_df_duT_v,
     time_mask,
-    client,
 ):
     """
     Stack the blocks of field derivatives for a given timestep and call the direct solver.
     """
-    if len(ATinv_df_duT_v) == 0:
-        ATinv_df_duT_v = [[] for _ in block]
+    if len(block) == 0:
+        return None
 
     Asubdiag = None
     if tInd < self.nT - 1:
         Asubdiag = self.getAsubdiag(tInd + 1)
 
-    updated_ATinv_df_duT_v = []
-
-    for (_, (rx_ind, _, shape)), field_deriv, ATinv_chunk in zip(
-        block, field_derivs, ATinv_df_duT_v
-    ):
+    time_blocks = []
+    colm_indices = []
+    colm_count = 0
+    for (_, (rx_ind, _, shape)), field_deriv in zip(block, field_derivs):
 
         # Cut out early data
         time_check = np.kron(time_mask, np.ones(shape, dtype=bool))[rx_ind]
         local_ind = np.arange(rx_ind.shape[0])[time_check]
 
-        if len(ATinv_chunk) == 0:
+        if len(ATinv_df_duT_v) == 0:
             # last timestep (first to be solved)
             time_block = field_deriv.toarray()[:, local_ind]
-            shape = (
-                field_deriv.shape[0],
-                len(rx_ind),
-            )
-            ATinv_chunk = np.zeros(shape, dtype=np.float32)
         else:
             time_block = np.asarray(
-                field_deriv[:, local_ind] - Asubdiag.T * ATinv_chunk[:, local_ind]
+                field_deriv[:, local_ind]
+                - Asubdiag.T
+                * ATinv_df_duT_v[:, colm_count : colm_count + rx_ind.shape[0]][
+                    :, local_ind
+                ]
             )
 
-        if time_block.ndim == 2 and time_block.shape[1] > 0:
-            solve = (AdiagTinv * time_block).reshape(time_block.shape)
-            ATinv_chunk[:, local_ind] = solve
+        time_blocks.append(time_block)
+        colm_indices.append(local_ind + colm_count)
+        colm_count += rx_ind.shape[0]
+
+    if len(ATinv_df_duT_v) == 0:
+        ATinv_df_duT_v = np.zeros((field_deriv.shape[0], colm_count), dtype=np.float32)
 
-        updated_ATinv_df_duT_v.append(ATinv_chunk)
+    block_array = np.hstack(time_blocks).reshape((ATinv_df_duT_v.shape[0], -1))
+    if block_array.shape[1] > 0:
+        solve = AdiagTinv * block_array
+        ATinv_df_duT_v[:, np.hstack(colm_indices).flatten()] = solve.reshape(
+            (-1, block_array.shape[1])
+        )
 
-    return updated_ATinv_df_duT_v
+    return ATinv_df_duT_v
 
 
 def block_deriv(
@@ -462,11 +465,14 @@ def compute_rows(
     Compute the rows of the sensitivity matrix for a given source and receiver.
     """
     rows = []
-    for ind, (address, ind_array) in enumerate(blocks[block_ind]):
+    colm_count = 0
+    for address, ind_array in blocks[block_ind]:
         # for (address, ind_array), field_derivs in zip(chunks, ATinv_df_duT_v):
         src = simulation.survey.source_list[address[0]]
         time_check = np.kron(time_mask, np.ones(ind_array[2], dtype=bool))[ind_array[0]]
-        local_ind = np.arange(len(ind_array[0]))[time_check]
+
+        n_rec = len(ind_array[0])
+        local_ind = np.arange(n_rec)[time_check]
 
         if len(local_ind) < 1:
             row_block = np.zeros(
@@ -478,18 +484,24 @@ def compute_rows(
         dAsubdiagT_dm_v = simulation.getAsubdiagDeriv(
             tInd,
             fields[:, address[0], tInd],
-            field_derivs[block_ind][ind][:, local_ind],
+            field_derivs[block_ind][:, colm_count : colm_count + n_rec][:, local_ind],
             adjoint=True,
         )
 
         dRHST_dm_v = simulation.getRHSDeriv(
-            tInd + 1, src, field_derivs[block_ind][ind][:, local_ind], adjoint=True
+            tInd + 1,
+            src,
+            field_derivs[block_ind][:, colm_count : colm_count + n_rec][:, local_ind],
+            adjoint=True,
         )  # on nodes of time mesh
 
         un_src = fields[:, address[0], tInd + 1]
         # cell centered on time mesh
         dAT_dm_v = simulation.getAdiagDeriv(
-            tInd, un_src, field_derivs[block_ind][ind][:, local_ind], adjoint=True
+            tInd,
+            un_src,
+            field_derivs[block_ind][:, colm_count : colm_count + n_rec][:, local_ind],
+            adjoint=True,
         )
         row_block = np.zeros(
             (len(ind_array[1]), simulation.model.size), dtype=np.float32
@@ -506,6 +518,8 @@ def compute_rows(
         else:
             Jmatrix[ind_array[1], :] += row_block
 
+        colm_count += n_rec
+
 
 def evaluate_dpred_block(indices, sources, mesh, time_mesh, fields):
     """

simpeg/dask/potential_fields/base.py

@@ -4,10 +4,8 @@
 
 import os
 from dask import delayed, array, compute
-
 from dask.diagnostics import ProgressBar
-
-import zarr
+from dask.distributed import Client
 
 
 _chunk_format = "row"
@@ -41,42 +39,15 @@ def residual(self, m, dobs, f=None):
     return self.dpred(m, f=f) - dobs
 
 
-def block_compute(sim, rows, components, j_matrix, count):
+def block_compute(sim, rows, components):
     block = []
     for row in rows:
         block.append(sim.evaluate_integral(row, components))
 
     if sim.store_sensitivities == "forward_only":
         return np.hstack(block)
 
-    values = np.vstack(block)
-    return storage_formatter(values, count, j_matrix)
-
-
-def storage_formatter(
-    rows: np.ndarray,
-    count: int,
-    j_matrix: zarr.Array | None = None,
-):
-    """
-    Format the storage of the sensitivity matrix.
-
-    :param rows: List of dask arrays representing blocks of the sensitivity matrix.
-    :param count: Current row count offset.
-    :param j_matrix: Zarr array to store the sensitivity matrix on disk, if applicable
-
-    :return: If j_matrix is provided, returns None after storing the rows; otherwise,
-    returns the stacked rows as a NumPy array.
-    """
-
-    if isinstance(j_matrix, zarr.Array):
-        j_matrix.set_orthogonal_selection(
-            (np.arange(count, count + rows.shape[0]), slice(None)),
-            rows.astype(np.float32),
-        )
-        return None
-
-    return rows
+    return np.vstack(block)
 
 
 def linear_operator(self):
@@ -86,44 +57,34 @@ def linear_operator(self):
         n_cells *= 3
 
     if self.store_sensitivities == "disk":
-
         if os.path.exists(self.sensitivity_path):
             return array.from_zarr(self.sensitivity_path)
 
-        Jmatrix = zarr.open(
-            self.sensitivity_path,
-            mode="w",
-            shape=(self.survey.nD, n_cells),
-            chunks=(self.max_chunk_size, n_cells),
-        )
-    else:
-        Jmatrix = None
-
     n_components = len(self.survey.source_list[0].receiver_list[0].components)
     n_blocks = np.ceil(
         (n_cells * n_components * self.survey.receiver_locations.shape[0] * 8.0 * 1e-6)
         / self.max_chunk_size
     )
     block_split = np.array_split(self.survey.receiver_locations, n_blocks)
-
     client, worker = self._get_client_worker()
 
-    if client:
+    if client is None:
+        client = Client()
+
+    if client and worker:
         sim = client.scatter(self, workers=worker)
     else:
         delayed_compute = delayed(block_compute)
 
     rows = []
     count = 0
-    for block in block_split:
-        if client:
+    for count, block in enumerate(block_split):
+        if client and worker:
             row = client.submit(
                 block_compute,
                 sim,
                 block,
                 self.survey.source_list[0].receiver_list[0].components,
-                Jmatrix,
-                count,
                 workers=worker,
             )
 
@@ -132,8 +93,6 @@ def linear_operator(self):
                 self,
                 block,
                 self.survey.source_list[0].receiver_list[0].components,
-                Jmatrix,
-                count,
             )
             row = array.from_delayed(
                 chunk,
@@ -147,14 +106,22 @@ def linear_operator(self):
         count += block.shape[0]
         rows.append(row)
 
-    if client:
+    if client and worker:
         kernel = client.gather(rows)
-    else:
+    elif forward_only:
         with ProgressBar():
             kernel = compute(rows)[0]
+    else:
+        kernel = rows
 
-    if self.store_sensitivities == "disk" and os.path.exists(self.sensitivity_path):
-        return array.from_zarr(self.sensitivity_path)
+    if self.store_sensitivities == "disk":
+        j_matrix = array.concatenate(rows, axis=0)
+
+        with ProgressBar():
+            j_matrix = j_matrix.to_zarr(
+                self.sensitivity_path, return_stored=True, compute=True
+            )
+        return j_matrix
 
     if forward_only:
         return np.hstack(kernel)