fix: join per-parameter adjoint gradients in _adjoint_array

CHANGELOG.md

@@ -13,6 +13,7 @@ and versions follow [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 - **`analysis = "subspace2"`**, the ensemble-transform IES of Raanes, Stordal & Evensen (2019), on ES-MDA, LM-EnRML and GN-EnRML. It solves for the `ne × ne` transform `W` directly, starting from `W = I`, and uses the analytic data covariance through `scale_data` rather than the ensemble representation `E Eᵀ` that `subspace` uses — so it takes no SVD and reads neither `energy` nor `iteration.energy`. It is exactly `margis` with the marginalised error scale `Ratio` fixed at 1, i.e. with the data uncertainty taken as known; `tests/assimilation/test_subspace2.py` pins that identity bit-for-bit. Ported from upstream 6f313d7, with two departures: the transform is initialised at iteration 0 rather than 1 (schemes here count from 0, so the reference version never initialises and dies on the first call), and the whitened observations are recomputed every call instead of cached on the first. ES-MDA redraws its observations and their scale at every assimilation step, so the cache would drive later steps with the first step's observations whitened by the first step's factor — the run still completes and the misfit still falls, which is why it would have gone unnoticed. On the characterisation case it is the difference between a posterior misfit of 175.7 and 201.1; the cached, lower number is the wrong one. The 39 existing golden arrays are unchanged.
 
 ### Fixed
+- **Adjoints with one gradient per parameter reach the analysis again.** A simulator may return each member's adjoint frame with `(datatype, parameter)` columns, one gradient per parameter in each cell. The frame path joined those gradients through `PETDataFrame._to_singlelevel_columns`; the `DataLayout`-based `_adjoint_array` that replaced it in #166 handed numpy the per-parameter Series as it was, so `forecast()` stopped with `ValueError: setting an array element with a sequence.` The gradients are now joined in column order, as before. Frames with one column per data type are unaffected.
 - **`autoadaloc` with `projection = "ensemble"` runs, and computes the update it describes.** The `approx` analysis built the projected anomalies as `X2 @ D_anom`, an `(nr, ne) @ (nd, ne)` product that raises `ValueError` unless `nd == ne`, and that would count the innovations twice even when it ran. They are now `V_r Σ_r X2`, the transform the non-localized branch already uses, so an all-ones taper gives the non-localized step exactly (#170).
 - **A `LOCALIZATION` block that names no mode runs again.** Localization was selected by which keyword appeared in the block -- `autoadaloc`, `localanalysis`, `dist_loc`, a pickled mask file, or none of them for the parallel update -- and the rewrite replaced that with a required `name` key. Every config written before the rewrite therefore stopped at startup with `Localization config has no 'name'`, naming three strings its author had never seen, and `pet migrate` did not cover the block. The mode is now inferred from the keyword that used to select it, and an explicit `name` still wins.
 - **A crashed simulation costs the optimizer its trial point, not the whole run.** `EnsembleOptimizationBase.function` raised `RuntimeError` whenever `calc_prediction` reported failure, so one control vector the simulator could not run ended the optimization and discarded every iteration before it. It now reports `inf`, which no backtracking step can improve on, so the point is rejected and the run continues from the last good iterate — the behaviour `765d637` removed. A crashed *single-point* evaluation additionally leaves `stateF` at its last good value: the gradient is `enF - repeat(stateF, nr)`, so writing `inf` there would have poisoned every later gradient with `inf`/`NaN` instead of rejecting one point. The separate abort when *every* member of a forecast fails is unchanged — there is nothing left to compute a gradient from.

src/pipt/ensembles/forecast.py

@@ -138,19 +138,33 @@ def _adjoint_array(self):
         members = self.member_adjoints
         if not members:
             return None
-        cells = {row: [np.asarray(member.loc[row.label, row.datatype], dtype=float).reshape(row.size, -1)
-                       for member in members] for row in self.data_layout.rows}
-        nx = next(iter(cells.values()))[0].shape[1]
-        out = np.empty((self.data_layout.nd, nx, len(members)))
-        for row, blocks in cells.items():
-            for j, block in enumerate(blocks):
-                out[row.rows, :, j] = block
+
+        rows = self.data_layout.rows
+        # The layout's rows are contiguous and in order, so stacking each member's
+        # blocks row by row gives its (nd, nx) Jacobian; members go last.
+        jacobians = [np.vstack([self._adjoint_block(member, row) for row in rows]) for member in members]
+        out = np.stack(jacobians, axis=-1)
+
         if self.data_df.is_scaled:
+            # Scaled data is (d - min) / (max - min), so its sensitivity is d's divided by (max - min).
             span = self.data_df.scale_max - self.data_df.scale_min
-            for row in self.data_layout.rows:
-                out[row.rows] = (out[row.rows] - 0) / span[row.datatype]
+            for row in rows:
+                out[row.rows] /= span[row.datatype]
         return out
 
+    @staticmethod
+    def _adjoint_block(member, row):
+        """The sensitivities of one observed cell in one member's adjoint frame, as ``(row.size, nx)``.
+
+        With ``(datatype, parameter)`` columns the cell holds one gradient per
+        parameter; they are joined in column order into one gradient over the
+        whole state.
+        """
+        cell = member.loc[row.label, row.datatype]
+        if isinstance(cell, pd.Series):
+            cell = np.concatenate([np.ravel(gradient) for gradient in cell.values])
+        return np.asarray(cell, dtype=float).reshape(row.size, -1)
+
     def _record_positions(self):
         """Where each observed label sits in a member's records: the simulator's ``true_order``, else the label itself."""
         order = getattr(self.sim, "true_order", None)

tests/assimilation/test_adjoints.py

@@ -1,6 +1,7 @@
 """Adjoints reach the analysis as an ``(nd, nx, ne)`` array aligned with the prediction rows."""
 
 import numpy as np
+import pandas as pd
 
 from input_output import read_config
 from misc.structures import PETDataFrame
@@ -26,6 +27,34 @@ def test_the_adjoint_array_is_what_the_frame_path_stacked(tmp_path, monkeypatch)
     assert ensemble.adjoints.shape == (ensemble.pred_data.nd, ensemble.enX.shape[0], NE)
 
 
+def _per_parameter(frame, parameters):
+    """``frame`` with each cell's gradient split over ``(datatype, parameter)`` columns, as ``parameters`` slices it."""
+    columns = pd.MultiIndex.from_tuples([(datatype, name) for datatype in frame.columns for name in parameters])
+    out = pd.DataFrame(index=frame.index, columns=columns, dtype=object)
+    for label in frame.index:
+        for datatype in frame.columns:
+            for name, part in parameters.items():
+                out.at[label, (datatype, name)] = np.asarray(frame.at[label, datatype])[part]
+    return out
+
+
+def test_per_parameter_adjoint_columns_are_joined_into_one_gradient(tmp_path, monkeypatch):
+    # A simulator may give one gradient per parameter, under (datatype, parameter)
+    # columns. The array must be as if each cell held those gradients joined in
+    # column order; the lookup used to hand numpy a Series of arrays and fail.
+    monkeypatch.chdir(tmp_path)
+    report_points = _write_synthetic_case(ne=NE)
+    cfg_da, cfg_sim, cfg_ens = read_config.read(_write_config("adjpar", "esmda", "approx", report_points, ne=NE))
+    cfg_sim["compute_adjoints"] = True
+    ensemble = AssimilationEnsemble(cfg_da, cfg_ens, VanDerPolOscillator(cfg_sim))
+    ensemble.forecast(ensemble.enX)
+    joined = ensemble.adjoints
+
+    parameters = {"x": slice(0, 2), "mu": slice(2, None)}
+    ensemble.member_adjoints = [_per_parameter(frame, parameters) for frame in ensemble.member_adjoints]
+    np.testing.assert_array_equal(ensemble._adjoint_array(), joined)
+
+
 def test_without_adjoints_the_ensemble_carries_none(tmp_path, monkeypatch):
     monkeypatch.chdir(tmp_path)
     report_points = _write_synthetic_case(ne=NE)