From 77fda8742878cc53b1ec26b132509fe6259c19b7 Mon Sep 17 00:00:00 2001
From: domfournier <dominiquef@mirageoscience.com>
Date: Mon, 17 Aug 2026 13:39:54 -0700
Subject: [PATCH 1/2] Change behaviour to store blocks directly.

---
 simpeg/dask/potential_fields/base.py | 75 ++++++++--------------------
 1 file changed, 21 insertions(+), 54 deletions(-)

diff --git a/simpeg/dask/potential_fields/base.py b/simpeg/dask/potential_fields/base.py
index 20a3ee8dc6..d1589abb1d 100644
--- a/simpeg/dask/potential_fields/base.py
+++ b/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,7 +39,7 @@ 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))
@@ -49,34 +47,7 @@ def block_compute(sim, rows, components, j_matrix, count):
     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)

From 320e79fab3a88a9d8e93991bebdb7962c3312eb9 Mon Sep 17 00:00:00 2001
From: domfournier <dominiquef@mirageoscience.com>
Date: Mon, 24 Aug 2026 08:18:28 -0700
Subject: [PATCH 2/2] Apply suggestions from code review

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
---
 simpeg/dask/potential_fields/base.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/simpeg/dask/potential_fields/base.py b/simpeg/dask/potential_fields/base.py
index d1589abb1d..05f4231141 100644
--- a/simpeg/dask/potential_fields/base.py
+++ b/simpeg/dask/potential_fields/base.py
@@ -71,14 +71,16 @@ def linear_operator(self):
     if client is None:
         client = Client()
 
-    if client and worker:
+    if client and worker and self.store_sensitivities != "disk":
         sim = client.scatter(self, workers=worker)
     else:
         delayed_compute = delayed(block_compute)
 
     rows = []
     count = 0
-    for count, block in enumerate(block_split):
+    for block in block_split:
+        if len(block) == 0:
+            continue
         if client and worker:
             row = client.submit(
                 block_compute,