GEOPY-2910: Reduce chunking of sensitivities for TEM inversions

[domfournier]

GEOPY-2910 - Reduce chunking of sensitivities for TEM inversions

follow up on fromer #150 and #152 that got reverted

[Copilot]

Pull request overview

This PR targets GEOPY-2910 by adjusting how sensitivity (Jacobian) blocks are chunked/computed in Dask-based simulations, aiming to reduce chunking overhead for TEM inversions.

Changes:

• Refactors potential-fields sensitivity block computation and disk persistence to use Dask array/Zarr workflows.
• Updates TEM time-domain sensitivity assembly to reduce per-chunk overhead by stacking/solving larger derivative blocks.
• Adds an optimization flag to parallel block partitioning for TEM sensitivity computation.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 3 comments.

File	Description
simpeg/dask/potential_fields/base.py	Reworks sensitivity block evaluation, distributed-client handling, and disk (zarr) persistence logic.
simpeg/dask/electromagnetics/time_domain/simulation.py	Refactors derivative block handling and row assembly to reduce chunking overhead in TEM Jacobian computation.

Suppressed comments (3)

simpeg/dask/potential_fields/base.py:83

• This client.submit(...) path should not be used when store_sensitivities == "disk", because the disk-writing code expects dask arrays (from_delayed) rather than distributed Futures. Gate this branch so disk mode always uses the delayed/from_delayed path.

    for count, block in enumerate(block_split):
        if client and worker:
            row = client.submit(

simpeg/dask/potential_fields/base.py:110

• client.gather(rows) assumes rows is a list of Futures, but in disk mode the loop should be building dask arrays for to_zarr. If this branch runs in disk mode it will error or do the wrong thing; gate it the same way as the submit path.

    if client and worker:
        kernel = client.gather(rows)

simpeg/dask/electromagnetics/time_domain/simulation.py:481

• When local_ind is empty, the function continues without advancing colm_count. That misaligns the column window (colm_count : colm_count + n_rec) for all subsequent receivers in the block, producing incorrect sensitivities. colm_count should advance by n_rec regardless of whether any data are kept for this receiver/time mask.

        if len(local_ind) < 1:
            row_block = np.zeros(
                (len(ind_array[1]), simulation.model.size), dtype=np.float32
            )
            rows.append(row_block)

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Copilot]

Pull request overview

Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.

Suppressed comments (2)

simpeg/dask/electromagnetics/time_domain/simulation.py:482

• When this chunk has no active time columns, continue skips the new column-offset update. The next chunk then slices field_derivs from the skipped chunk's columns, producing incorrect TEM sensitivities whenever an inactive chunk precedes an active one. Advance colm_count before continuing.

        if len(local_ind) < 1:
            row_block = np.zeros(
                (len(ind_array[1]), simulation.model.size), dtype=np.float32
            )
            rows.append(row_block)
            continue

simpeg/dask/potential_fields/base.py:118

• With an assigned worker, rows contains distributed Future objects even after they have been gathered into kernel. Passing those futures to dask.array.concatenate fails because they do not expose array dimensions/chunks, so disk-backed sensitivities break in the worker-pinned distributed path. Concatenate the gathered/lazy kernel selected above instead.

        j_matrix = array.concatenate(rows, axis=0)