Simplify output to (n_sensor, Nt) everywhere, add reshape_to_grid helper

Keep the simplest possible output contract: all time-series are
(n_sensor, Nt) with sensor points in C-flattened order. Aggregates
are (n_sensor,). No automatic reshaping.

Users who want spatial structure call reshape_to_grid(data, grid_shape)
which converts (n_sensor, Nt) → (*grid_shape, Nt).

Removes _reshape_sensor_to_grid and all the time-first reshaping logic.
Integration tests use _c_to_f_reorder to compare C-flat Python output
against F-flat MATLAB references.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: waltsims

kwave/kspaceFirstOrder.py

@@ -157,10 +157,9 @@ def kspaceFirstOrder(
         dict: Recorded sensor data keyed by field name (e.g.
         ``"p"``, ``"p_final"``, ``"ux"``, ``"uy"``).
 
-        Sensor time-series are C-ordered.  When the sensor mask covers the
-        entire grid, time-series fields are returned as
-        ``(Nt, *grid_shape)`` (time-first).  For partial masks the shape is
-        ``(n_sensor, Nt)`` with sensor points in C-flattened order.
+        All time-series are ``(n_sensor, Nt)`` with sensor points in
+        C-flattened order.  Use :func:`reshape_to_grid` to recover spatial
+        structure for full-grid masks.
     """
     if device not in ("cpu", "gpu"):
         raise ValueError(f"device must be 'cpu' or 'gpu', got {device!r}")
@@ -178,8 +177,6 @@ def kspaceFirstOrder(
 
     # --- Shared pre-processing (both backends) ---
 
-    user_grid_shape = tuple(int(n) for n in kgrid.N)
-
     if not pml_inside:
         kgrid, medium, source, sensor = _expand_for_pml_outside(kgrid, medium, source, sensor, pml_size)
 
@@ -250,35 +247,28 @@ def kspaceFirstOrder(
     if not pml_inside:
         result = _strip_pml(result, pml_size, kgrid.dim)
 
-    result = _reshape_sensor_to_grid(result, sensor, user_grid_shape)
-
     return result
 
 
-def _reshape_sensor_to_grid(result, sensor, user_grid_shape):
-    """Reshape sensor time-series to (Nt, *grid_shape) when the sensor covers the full user grid.
+def reshape_to_grid(data, grid_shape):
+    """Reshape flat sensor data to grid shape.
 
-    After PML expansion the sensor mask is padded with zeros, so we check
-    against the user's original grid shape (before expansion).
+    Convenience helper for full-grid sensor masks where ``n_sensor``
+    equals the total number of grid points.
+
+    Args:
+        data: sensor array — ``(n_sensor, Nt)`` time-series or
+            ``(n_sensor,)`` aggregate.
+        grid_shape: tuple of grid dimensions, e.g. ``(Nx, Ny)``.
+
+    Returns:
+        For time-series: ``(*grid_shape, Nt)``
+        For aggregates:  ``(*grid_shape)``
     """
-    user_numel = int(np.prod(user_grid_shape))
-
-    mask = getattr(sensor, "mask", None) if sensor is not None else None
-    if mask is None:
-        n_sensor = user_numel
-    elif _is_cartesian_mask(mask, len(user_grid_shape)):
-        return result
-    else:
-        n_sensor = int(np.asarray(mask, dtype=bool).sum())
-
-    if n_sensor != user_numel:
-        return result
-
-    for key, val in result.items():
-        if not isinstance(val, np.ndarray):
-            continue
-        if val.ndim == 2 and val.shape[0] == n_sensor:
-            result[key] = val.T.reshape(-1, *user_grid_shape)
-        elif val.ndim == 1 and val.shape[0] == n_sensor:
-            result[key] = val.reshape(user_grid_shape)
-    return result
+    data = np.asarray(data)
+    if data.ndim == 2:
+        n_sensor, Nt = data.shape
+        return data.reshape(*grid_shape, Nt)
+    elif data.ndim == 1:
+        return data.reshape(grid_shape)
+    return data

tests/integration/conftest.py

@@ -33,48 +33,44 @@ def _load(name):
     return _load
 
 
-def _to_matlab_shape(py_val, mat_val):
-    """Reshape C-order Python output to match MATLAB F-order reference shape.
+def _c_to_f_reorder(py_val, grid_shape):
+    """Reorder sensor rows from C-flat to F-flat ordering for MATLAB comparison.
 
-    The new API returns full-grid time-series as (Nt, *grid_shape) in C-order.
-    MATLAB references store them as (n_sensor, Nt) in F-flat order.
-    Similarly, aggregates are (*grid_shape) vs MATLAB (n_sensor,).
+    Python (C-order) and MATLAB (F-order) flatten grid points in different
+    orders.  Both produce (n_sensor, Nt) but the row ordering differs.
+    This builds a permutation to reorder Python rows to match MATLAB.
     """
-    if py_val.shape == mat_val.shape:
+    if grid_shape is None or len(grid_shape) < 2:
         return py_val
-
-    # Time-series ≥3D: (Nt, *grid_shape) → (n_sensor, Nt) with F-order flatten
-    if py_val.ndim >= 3 and mat_val.ndim == 2:
-        Nt = py_val.shape[0]
-        # Move time axis last, then F-order flatten the grid dims
-        return np.moveaxis(py_val, 0, -1).reshape(-1, Nt, order="F")
-
-    # Time-series 1D: (Nt, N) → (N, Nt) — both 2D but transposed
-    if py_val.ndim == 2 and mat_val.ndim == 2 and py_val.shape == mat_val.shape[::-1]:
-        return py_val.T
-
-    # Aggregates: (*grid_shape) → (n_sensor,) with F-order flatten
-    if py_val.ndim >= 2 and mat_val.ndim == 1:
-        return py_val.ravel(order="F")
-
+    c_indices = np.arange(int(np.prod(grid_shape)))
+    f_indices = np.ravel_multi_index(np.unravel_index(c_indices, grid_shape), grid_shape, order="F")
+    # f_indices[i] = where C-flat point i lands in F-flat order
+    # We need the inverse: for each F-flat position, which C-flat row?
+    inv = np.argsort(f_indices)
+    if py_val.ndim == 2 and py_val.shape[0] == len(c_indices):
+        return py_val[inv]
+    if py_val.ndim == 1 and py_val.shape[0] == len(c_indices):
+        return py_val[inv]
     return py_val
 
 
-def assert_fields_close(result, ref, fields, *, rtol=1e-10, atol=1e-12):
+def assert_fields_close(result, ref, fields, *, rtol=1e-10, atol=1e-12, grid_shape=None):
     """Compare Python result dict against MATLAB reference arrays.
 
     Args:
         result: dict from kspaceFirstOrder()
         ref: dict from scipy.io.loadmat()
         fields: list of (python_key, matlab_key) tuples
+        grid_shape: tuple of grid dims for C→F row reordering (full-grid masks only)
         rtol, atol: tolerances passed to np.testing.assert_allclose
     """
     for py_key, mat_key in fields:
         assert py_key in result, f"Python result missing key '{py_key}'"
         assert mat_key in ref, f"MATLAB reference missing key '{mat_key}'"
         py_val = np.atleast_1d(np.squeeze(np.asarray(result[py_key])))
         mat_val = np.atleast_1d(np.squeeze(np.asarray(ref[mat_key])))
-        py_val = _to_matlab_shape(py_val, mat_val)
+        if grid_shape is not None:
+            py_val = _c_to_f_reorder(py_val, grid_shape)
         assert py_val.shape == mat_val.shape, f"Shape mismatch for {py_key}: Python {py_val.shape} vs MATLAB {mat_val.shape}"
         np.testing.assert_allclose(
             py_val, mat_val, rtol=rtol, atol=atol, err_msg=f"Field '{py_key}' differs from MATLAB reference '{mat_key}'"

tests/integration/test_ivp_2D.py

@@ -42,4 +42,5 @@ def test_ivp_2D_vs_matlab(load_matlab_ref):
         result,
         ref,
         [("p", "sensor_data_p")],
+        grid_shape=(128, 128),
     )

tests/test_c_order.py

@@ -9,7 +9,7 @@
 from kwave.kmedium import kWaveMedium
 from kwave.ksensor import kSensor
 from kwave.ksource import kSource
-from kwave.kspaceFirstOrder import _is_cartesian_mask, _reshape_sensor_to_grid
+from kwave.kspaceFirstOrder import reshape_to_grid
 
 # ---------------------------------------------------------------------------
 # _fix_output_order (CppSimulation)
@@ -111,55 +111,34 @@ def test_1d_aggregate_reorder(self):
 
 
 # ---------------------------------------------------------------------------
-# _reshape_sensor_to_grid
+# reshape_to_grid helper
 # ---------------------------------------------------------------------------
 
 
-class TestReshapeSensorToGrid:
-    def test_cartesian_mask_unchanged(self):
-        """Cartesian sensor masks should not be reshaped."""
-        sensor = SimpleNamespace(mask=np.array([[0.0, 1e-3], [0.0, 0.0]]))  # (2, 2) Cartesian
-        result = {"p": np.arange(20).reshape(2, 10)}
-        out = _reshape_sensor_to_grid(result, sensor, (64, 64))
-        assert out["p"].shape == (2, 10)  # unchanged
-
-    def test_partial_mask_unchanged(self):
-        """Partial binary masks should not be reshaped."""
-        mask = np.zeros((8, 8), dtype=bool)
-        mask[0, 0] = True
-        mask[4, 4] = True
-        sensor = SimpleNamespace(mask=mask)
-        result = {"p": np.arange(20).reshape(2, 10)}
-        out = _reshape_sensor_to_grid(result, sensor, (8, 8))
-        assert out["p"].shape == (2, 10)  # unchanged
-
-    def test_full_grid_reshaped(self):
-        """Full-grid binary mask is reshaped to (Nt, *grid_shape)."""
-        sensor = SimpleNamespace(mask=np.ones((4, 6), dtype=bool))
-        Nt = 5
-        result = {"p": np.arange(120).reshape(24, Nt)}
-        out = _reshape_sensor_to_grid(result, sensor, (4, 6))
-        assert out["p"].shape == (Nt, 4, 6)
-
-    def test_aggregate_reshaped(self):
-        """1D aggregates reshaped to grid_shape for full-grid masks."""
-        sensor = SimpleNamespace(mask=np.ones((4, 6), dtype=bool))
-        result = {"p_max": np.arange(24)}
-        out = _reshape_sensor_to_grid(result, sensor, (4, 6))
-        assert out["p_max"].shape == (4, 6)
-
-    def test_sensor_none(self):
-        """sensor=None means full grid."""
-        result = {"p": np.arange(80).reshape(16, 5)}
-        out = _reshape_sensor_to_grid(result, None, (4, 4))
-        assert out["p"].shape == (5, 4, 4)
-
-    def test_non_array_values_unchanged(self):
-        """Non-ndarray values in result are passed through."""
-        sensor = SimpleNamespace(mask=np.ones((4, 4), dtype=bool))
-        result = {"p": np.arange(80).reshape(16, 5), "metadata": "hello"}
-        out = _reshape_sensor_to_grid(result, sensor, (4, 4))
-        assert out["metadata"] == "hello"
+class TestReshapeToGrid:
+    def test_time_series_2d(self):
+        """(n_sensor, Nt) → (*grid_shape, Nt)."""
+        data = np.arange(120).reshape(24, 5)
+        out = reshape_to_grid(data, (4, 6))
+        assert out.shape == (4, 6, 5)
+
+    def test_aggregate_1d(self):
+        """(n_sensor,) → (*grid_shape)."""
+        data = np.arange(24)
+        out = reshape_to_grid(data, (4, 6))
+        assert out.shape == (4, 6)
+
+    def test_3d_grid(self):
+        """Works with 3D grids."""
+        data = np.arange(60).reshape(60, 1)
+        out = reshape_to_grid(data, (3, 4, 5))
+        assert out.shape == (3, 4, 5, 1)
+
+    def test_passthrough_higher_dim(self):
+        """Higher-dim arrays pass through unchanged."""
+        data = np.arange(120).reshape(2, 3, 4, 5)
+        out = reshape_to_grid(data, (4, 6))
+        assert out.shape == (2, 3, 4, 5)
 
 
 # ---------------------------------------------------------------------------

tests/test_kspaceFirstOrder.py

@@ -61,8 +61,7 @@ def test_python_backend_runs(self, sim_2d):
         kgrid, medium, source, sensor = sim_2d
         result = kspaceFirstOrder(kgrid, medium, source, sensor, backend="python")
         assert "p" in result
-        # Full-grid sensor → (Nt, *grid_shape) in C-order
-        assert result["p"].shape == (int(kgrid.Nt), 64, 64)
+        assert result["p"].shape == (int(sensor.mask.sum()), int(kgrid.Nt))
 
     def test_cpp_save_only(self, sim_2d):
         kgrid, medium, source, sensor = sim_2d

tests/test_native_solver.py

@@ -43,7 +43,7 @@ def test_p0_source(self, grid_2d):
         result = kspaceFirstOrder(
             grid_2d, kWaveMedium(sound_speed=1500), _p0_source((64, 64)), kSensor(mask=np.ones((64, 64), dtype=bool)), backend="python"
         )
-        assert result["p"].shape == (10, 64, 64)
+        assert result["p"].shape == (64 * 64, 10)
         assert np.max(np.abs(result["p"])) > 0
 
     def test_heterogeneous_medium(self, grid_2d):
@@ -56,7 +56,7 @@ def test_heterogeneous_medium(self, grid_2d):
             kSensor(mask=np.ones((64, 64), dtype=bool)),
             backend="python",
         )
-        assert result["p"].shape == (10, 64, 64)
+        assert result["p"].shape == (64 * 64, 10)
 
     def test_absorption(self, grid_2d):
         result = kspaceFirstOrder(
@@ -66,7 +66,7 @@ def test_absorption(self, grid_2d):
             kSensor(mask=np.ones((64, 64), dtype=bool)),
             backend="python",
         )
-        assert result["p"].shape == (10, 64, 64)
+        assert result["p"].shape == (64 * 64, 10)
 
     def test_pml_auto(self):
         kgrid = kWaveGrid(Vector([128, 128]), Vector([0.1e-3, 0.1e-3]))
@@ -84,7 +84,7 @@ def test_record_aggregates(self, grid_2d):
         sensor = kSensor(mask=np.ones((64, 64), dtype=bool))
         sensor.record = ["p", "p_max", "p_rms"]
         result = kspaceFirstOrder(grid_2d, kWaveMedium(sound_speed=1500), _p0_source((64, 64)), sensor, backend="python")
-        assert result["p_max"].shape == (64, 64)
+        assert result["p_max"].shape == (64 * 64,)
         assert "p_rms" in result
 
 
@@ -119,7 +119,7 @@ def test_nonlinearity_bona(self, grid_2d):
             kSensor(mask=np.ones((64, 64), dtype=bool)),
             backend="python",
         )
-        assert result["p"].shape == (10, 64, 64)
+        assert result["p"].shape == (64 * 64, 10)
 
     def test_stokes_absorption(self, grid_2d):
         result = kspaceFirstOrder(
@@ -129,7 +129,7 @@ def test_stokes_absorption(self, grid_2d):
             kSensor(mask=np.ones((64, 64), dtype=bool)),
             backend="python",
         )
-        assert result["p"].shape == (10, 64, 64)
+        assert result["p"].shape == (64 * 64, 10)
 
     def test_dirichlet_pressure_source(self, grid_1d):
         source = kSource()
@@ -146,16 +146,16 @@ def test_velocity_recording(self, grid_2d):
         sensor = kSensor(mask=np.ones((64, 64), dtype=bool))
         sensor.record = ["p", "ux", "uy", "ux_max", "uy_rms", "ux_final", "p_final"]
         result = kspaceFirstOrder(grid_2d, kWaveMedium(sound_speed=1500), _p0_source((64, 64)), sensor, backend="python")
-        assert result["ux"].shape == (10, 64, 64)
-        assert result["ux_max"].shape == (64, 64)
+        assert result["ux"].shape == (64 * 64, 10)
+        assert result["ux_max"].shape == (64 * 64,)
         assert "ux_final" in result and "p_final" in result
 
     def test_intensity_recording(self, grid_2d):
         sensor = kSensor(mask=np.ones((64, 64), dtype=bool))
         sensor.record = ["p", "ux", "uy", "Ix", "Iy", "Ix_avg", "Iy_avg"]
         result = kspaceFirstOrder(grid_2d, kWaveMedium(sound_speed=1500), _p0_source((64, 64)), sensor, backend="python")
-        assert result["Ix"].shape == (10, 64, 64)
-        assert result["Ix_avg"].shape == (64, 64)
+        assert result["Ix"].shape == (64 * 64, 10)
+        assert result["Ix_avg"].shape == (64 * 64,)
 
     def test_record_start_index(self, grid_1d):
         source = kSource()
@@ -172,7 +172,7 @@ def test_sensor_none_records_everywhere(self, grid_1d):
         source.p0 = np.zeros(64)
         source.p0[32] = 1.0
         result = kspaceFirstOrder(grid_1d, kWaveMedium(sound_speed=1500), source, None, backend="python", pml_inside=True)
-        assert result["p"].shape == (20, 64)
+        assert result["p"].shape == (64, 20)
 
 
 class TestCppSaveOnly: