Epic #324: first-class MLX and reference-faithful fitting

.context/decisions/0001-torch-backend-natural-gradient-em.md

@@ -72,3 +72,19 @@ is now the sole PyTorch backend, and the public `AMICA` interface wraps it direc
 (the `backend=` selector and the basic-backend-only `debug`/`output_dir` fit args
 were dropped). The Context/Alternatives sections above describe the pre-removal state
 and are retained as the historical record.
+
+## Addendum (2026-09-22, issue #333)
+
+The #24 fix left `A` stored with each model's block equal to the TRANSPOSE of the reference's per-model mixing matrix
+(`get_mixing_matrix(h)` returns `A[:, comp_list[:, h]].T`), so a component is a ROW of its block, not a stored column.
+[ADR 0006](0006-component-orientation.md) states that convention precisely,
+records the `doscaling` defect it caused (stored columns were normalized; fixed in epic #324 Phase 7)
+and the planned move to a component-row layout for `share_comps` (Phase 8, issue #334).
+
+## Addendum (2026-09-23, issue #334)
+
+[ADR 0007](0007-component-row-layout.md) replaces the storage this ADR introduced:
+`A` is now `(n_comps, n)` with one component per row and `comp_list` indexing rows,
+so each model's block, and every result without a `share_comps` merge, is unchanged,
+while the share metric and merge act on components as the reference's do.
+Saved models convert without loss unless a merge had fired.

.context/decisions/0003-best-iterate-safeguard.md

@@ -76,10 +76,59 @@ latter joined in Phase 3, issue #289) -- see `pamica/mlx_impl/core.py`'s
 `_snapshot_params`/`_restore_params` and
 `pamica/tests/mlx_tests/test_mlx_keepbest.py`.
 
+## Re-measured after epic #324 (2026-09-23, issue #351)
+
+Epic #324 changed the default trajectory of every backend (component-row
+`doscaling` #333, 1-based schedule gates #335, the reference's iteration order
+and unconditional `A`-freeze #339/#345, the normalized initial `A` #341,
+single-precision constants #344), so the figures above describe code that no
+longer ships. The ensemble was re-run with the #51 protocol
+(`.context/issue-351/keep_best_ensemble.py`: the sample EEG, `n_models=2`, the
+#51 keyword set, N = 20, pamica seeds 0-19), with two changes:
+the reference is the pinned v0.3.3 native binary, seeded 0-19 and
+single-threaded (the #51 runs used the clock-seeded `amica15mac`), and each
+fit runs 300 iterations once, with the 100- and 200-iteration values read from
+its trajectory (checked equal to separate 100-iteration fits, with `keep_best`
+on and off, on two seeds on both sides).
+
+| budget | Fortran mean (sd) | pamica return-last mean (sd) | pamica `keep_best` mean (sd) | sd ratio | pamica minus Fortran | Kolmogorov-Smirnov (KS) p | restores |
+|---:|---:|---:|---:|---:|---:|---:|---:|
+| 100 | -3.3550 (0.0029) | -3.3542 (0.0029) | -3.3542 (0.0029) | 1.0x | +0.0008 | 0.83 | 0 of 20 |
+| 200 | -3.3418 (0.0030) | -3.3416 (0.0022) | -3.3416 (0.0022) | 0.75x | +0.0002 | 0.57 | 0 of 20 |
+| 300 | -3.3392 (0.0027) | -3.3391 (0.0020) | -3.3391 (0.0020) | 0.75x | +0.0001 | 0.98 | 1 of 20 |
+
+- The late overshoots that motivated this decision do not occur in these
+  fits: the largest likelihood decrease in any of the 20 pamica trajectories
+  is 1.2e-5, and none of them used a natural-gradient fallback. At 100
+  iterations the return-last and `keep_best` distributions are identical, and
+  their spread equals the reference's (sd ratio 1.0x, where #51 measured 12.7x
+  and 2.0x). Phase 7 of the epic (#340) traced most of the old overshoots to
+  the column-rule `doscaling` (#333).
+- A restore fired in one of the 20 fits, and only at the 300-iteration budget: seed 3,
+  whose best iterate (iteration 299) beat its last by 4.2e-6.
+- The mean gap is gone at every budget: pamica's mean lies 8.1e-4 above the
+  reference's at 100 iterations, 2.1e-4 at 200 and 1.0e-4 at 300, where #51
+  measured -0.009 at 100 and attributed it to convergence speed. The
+  bundled tier's clock-seeded ensemble (`benchmarks/reproduce_table1.py`,
+  `AMICA` seeds 1-20) agrees: -3.3541 against -3.3543, KS p 0.83. Refitting
+  that ensemble's pamica half with the code before the epic's changes to the fit (e38aa11) against
+  the same 20 reference fits gives -3.3627 (sd 0.006, KS p 1e-5), close to
+  the #27/#51 figures. Seven of those 20 fits stopped early on `min_dll`,
+  whose check counted likelihood dips as small gains until #339 (their mean
+  -3.3679), and the 13 that ran the full 100 iterations average -3.3600, so
+  the old gap came partly from the early stops and partly from the update
+  rule of that code.
+- The decision stands: `keep_best` stays on by default. With these fits it
+  rarely acts, and when it does the gain is small, but a run that ends below
+  its peak (a likelihood decrease near the end, or a Newton fallback) still
+  returns its best iterate.
+
 ## Receipts
 
 - `pamica/torch_impl/amica_torch_ng.py` (`keep_best`, `_snapshot_params`/
   `_restore_params`, `final_ll_`, `_KEEP_BEST_TOL`).
 - `pamica/tests/torch_tests/test_ng_backend.py::test_keep_best_*`.
 - `.context/issue-51/ensemble_ll.py` (Fortran-vs-NG LL ensemble, real data).
+- `.context/issue-351/keep_best_ensemble.py` and `raw/keep_best/` (the
+  re-measurement above), `.context/issue-351/findings.md`.
 - Fortran schedule: `pamica/amica15.f90:1038-1058` (anneal-on-decrease).

.context/decisions/0006-component-orientation.md

@@ -0,0 +1,105 @@
+# ADR 0006: Components are rows of each stored mixing block
+
+**Status:** accepted; the storage part amended by [ADR 0007](0007-component-row-layout.md)
+**Date:** 2026-09-22
+**Owner:** Seyed Yahya Shirazi
+
+Amends [ADR 0001](0001-torch-backend-natural-gradient-em.md),
+which introduced the storage convention without naming its consequence for per-component operations.
+
+## Context
+
+Since issue #24 every array backend (PyTorch, NumPy, MLX) stores the mixing matrix `A` with shape `(n, n_comps)` and indexes each model's block by `comp_list`,
+but each block holds the transpose of the reference's per-model mixing matrix:
+`get_mixing_matrix(h)` returns `A[:, comp_list[:, h]].T`, which is the reference's `A(:, comp_list(:, h))`.
+Source `i` of model `h` is therefore ROW `i` of the stored block `A[:, comp_list[:, h]]`,
+while its density parameters live at `mu[:, comp_list[i, h]]`, `beta[:, comp_list[i, h]]` (the reference's `sbeta`), `alpha` and `rho`.
+A stored COLUMN is not a component: it collects one sphered channel's loadings across the model's components.
+
+The reference's `doscaling` (amica15.f90:1843-1851) divides each component's mixing vector `A(:,k)` by its norm
+and rescales `mu(:,k) *= norm`, `sbeta(:,k) /= norm`.
+That is an exact change of scale: the source is multiplied by the norm, its density is rescaled to match, and the log-likelihood does not change.
+Every backend instead normalized stored columns with that same compensation (issue #333), which is not a change of scale of any component.
+A compensated rescale of stored row 5 changes the log-likelihood by exactly 0.0 on the sample EEG;
+the same rescale of stored column 5 changes it by -3.9e-2.
+Because `doscaling` is on by default, every default fit left the reference's trajectory from the first iteration:
+seeded from the same initialization, the native binary and pamica differed by 7.2e-5 in `A` and 6.6e-5 in `mu` after one iteration,
+while with `doscaling` off they agreed to 1e-15 and 1e-11.
+Scale-blind endpoint metrics hid it (log-likelihood, Hungarian-matched correlation),
+but fitted component norms ranged over [0.94, 1.07] in one-model fits and [0.05, 1.97] in two-model fits,
+where the reference's are exactly 1.
+
+The same orientation error affects the `share_comps` metric and merge, which compare and tie stored columns (issue #334).
+
+## Decision
+
+Keep the storage convention for now and make every per-component operation act on block rows.
+In this change (epic #324 Phase 7) `doscaling` normalizes, for every model `h` and source `i`, the row `A[i, comp_list[:, h]]`,
+multiplies `mu[:, comp_list[i, h]]` by the row's norm and divides `beta[:, comp_list[i, h]]` by it,
+leaving a zero-norm row untouched as the reference does.
+The rule is identical in the three backends (`_rescale_components`) and runs where the reference runs it:
+after the `A` update, before the unmixing matrices are rebuilt, and before a share merge.
+The reference parses `scalestep` but never reads it and rescales every iteration;
+pamica keeps `scalestep` as an extension,
+now counted from 1 like the reference's other cadences (iterations `scalestep`, `2*scalestep`, ...),
+so the default of 1 is the reference (row 14 of `docs/guides/amica-differences.md`).
+
+Epic #324 Phase 8 (issue #334, [ADR 0007](0007-component-row-layout.md)) then changed the storage itself to component rows,
+`A` of shape `(n_comps, n)` with `comp_list` indexing rows,
+so a stored row is a component in every configuration, including merged ones,
+and the rescale acts on each component row once.
+Until then, a column merged by `share_comps` belonged to rows of several blocks, where the per-block rule is not a change of scale;
+it was applied uniformly, model by model.
+
+## Consequences
+
+- Default trajectories change on every backend, toward the reference.
+  Seeded from pamica's initialization, `A`, `mu` and `sbeta` now match the native binary to float64 round-off after one and three natural-gradient iterations
+  (after one: `A` 5e-16, `mu` 8e-11, `sbeta` 1e-14; the column rule was 7e-5 off).
+  A 100-iteration Newton run against the seeded binary keeps its log-likelihood
+  (2.3e-4 to 2.1e-4 apart; the remaining gap is the Newton start, issue #335),
+  and the fitted component norms are exactly 1 instead of [0.94, 1.07].
+  Two-model fits improve most: at 100 iterations the matched correlation with the seeded binary rises from 0.850 to 0.99998.
+- Scale-blind results do not change materially:
+  log-likelihood, component directions, correlation and Amari distance against the bundled `amicaout` fixture move by about 1e-4.
+  Element-wise `A`, `mu` and `sbeta` now match the reference.
+- `doscaling=False` is byte-identical to the previous code on every backend, pinned by an in-process comparison with the pre-fix classes.
+- Saved models need no conversion:
+  the stored format is unchanged, and a model fitted before this change is still a valid model with non-unit component scales.
+  Refit only when comparing parameters element by element with the reference.
+- Several test recipes that relied on a non-monotone trajectory (the `keep_best` overshoot, the `min_dll` stop) no longer overshoot at their old settings:
+  the column rule's perturbation was part of what made them overshoot.
+  They were retuned (see the Phase 7 pull request).
+  ADR 0003's `keep_best` variance figures were measured under the column rule and will be re-measured in the epic's final parity re-measurement.
+- `scalestep > 1` now rescales on iterations `scalestep`, `2*scalestep`, ... instead of 1, `1+scalestep`, ...;
+  the default is unaffected.
+- Former known difference, now resolved: pamica did not normalize its drawn initial `A`
+  (the reference normalizes a drawn one, amica15.f90:818-819, but not a loaded one).
+  The two draws come from different random generators anyway, and the first iteration's rescale normalizes every component.
+  Aligning it changes default trajectories, so it was its own change, issue #341, not part of the layout change of Phase 8.
+  Epic #324 Phase 12 made it: every backend draws its initial `A` through `pamica.initialization.initial_mixing`,
+  which sets each block's diagonal to one and normalizes every component as the reference does, the NumPy restart redraw included,
+  and a supplied or loaded `A` is still used as is.
+  Default fits start from a different point and end close to where they did (the first rescale used to do the normalizing);
+  with `doscaling` off the change reaches the whole fit.
+
+## Alternatives considered
+
+- **Change the storage to component rows in the same change:** rejected for this phase.
+  The layout change touches the update, the accessors, persistence and the share code;
+  the rescale fix alone is small,
+  and landing it first gives Phase 8 a trajectory that is already reference-faithful to stay byte-identical to.
+- **Keep the column rule and document it as a difference:** rejected.
+  It is not a change of scale, so it is not an alternative normalization but a perturbation of the fit,
+  and it made element-wise parity with the reference impossible.
+
+## Receipts
+
+- amica15.f90:1843-1851 (per-iteration `doscaling`),
+  :818-819 and :1039-1040 (normalization of a drawn initial and restart `A`),
+  :793-800 (a loaded `A` is used as is),
+  :3686-3688 (`scalestep` parsed, never read).
+- Issues #333 (this change), #334 (Phase 8, component-row layout), #335 (Newton start), #341 (Phase 12, initial `A`), #24 (the storage convention), epic #324.
+- `pamica/tests/test_doscaling_rows.py` (invariance, cross-backend, byte-identity and the gated native oracle),
+  `pamica/tests/native_oracle.py` (seeded reference runs),
+  `pamica/tests/test_initial_mixing.py` (the initial `A`, with its gated oracles).

.context/decisions/0007-component-row-layout.md

@@ -0,0 +1,94 @@
+# ADR 0007: Store the mixing matrix with one component per row
+
+**Status:** accepted
+**Date:** 2026-09-23
+**Owner:** Seyed Yahya Shirazi
+
+Amends [ADR 0006](0006-component-orientation.md), which named the per-block convention and fixed `doscaling` within it,
+and [ADR 0001](0001-torch-backend-natural-gradient-em.md), which introduced the storage.
+
+## Context
+
+Since issue #24 every array backend stored `A` as `(n, n_comps)` and took model `h`'s block as `A[:, comp_list[:, h]]`,
+the transpose of the reference's per-model mixing matrix, so a component was a ROW of its block
+while the component ids in `comp_list`, and the density columns `mu[:, k]`, `beta[:, k]`, `alpha[:, k]`, `rho[:, k]`, indexed stored COLUMNS.
+A stored column held one sphered channel's loadings across a model's components, not a component.
+ADR 0006 made `doscaling` act on block rows, but `share_comps` still used the column ids for both of its steps (issue #334):
+
+- The metric compared de-sphered stored columns, not scalp maps.
+  On a two-model fit of the sample EEG (150 iterations) the |cosine| between a stored column and the true map `get_sensor_mixing_matrix(h)[:, i]` had median 0.27 and 0.39;
+  a planted exact duplicate component scored 0.74 and was not merged at `comp_thresh=0.99`.
+- The merge tied one channel's loadings across two models instead of two components' mixing vectors,
+  and the `gm`-weighted `dAk/zeta` average then averaged those.
+  Forcing it on a planted identical pair changed the log-likelihood by -0.186, where the reference's fold changes it by exactly 0.
+
+The reference's own scan never merges (its `Spinv2` is declared but never allocated, so every similarity is NaN),
+but its `load_comp_list` seeds a merged `comp_list`, and its update from that state is a bit-level oracle.
+
+## Decision
+
+Every array backend (PyTorch, NumPy, MLX) stores `A` as `(n_comps, n)`:
+row `k` is component `k`'s sphered-space mixing vector, the reference's column `A(:, k)`,
+and model `h`'s block is `A[comp_list[:, h], :]`, the same `n x n` matrix as before, so `W_h = inv(block)` is unchanged.
+Everything that touches `A` follows:
+the natural-gradient and Newton update scatter each model's step into the rows `comp_list` names (the `dAk/zeta` average now averages one component across the models that share it);
+`doscaling` rescales each component row once, including a shared one, and leaves a merged-away row alone;
+the share metric compares `pinv(sphere) @ A.T`, whose column `comp_list[i, h]` is column `i` of `get_sensor_mixing_matrix(h)`;
+the merge re-points `comp_list` and retires the folded row, as the reference does, copying and averaging nothing.
+The per-backend helper `_component_sensor_maps` returns the vectors the metric compares.
+
+Persistence converts or refuses, through one shared function (`pamica.component_layout.rows_from_legacy_columns`):
+the PyTorch `state_dict` is `format_version` 4 and the MLX save format 2;
+a version 3 (PyTorch) or version 1 (MLX) payload is converted without loss when its `comp_list` is unmerged
+(every block carries over element for element),
+and refused with a message to refit when a component id is shared across models,
+because that merge was computed under the column semantics and has no component-row counterpart.
+The `AMICA` wrapper's own save (version 2) wraps the backend payloads and goes through the same path.
+
+The EEGLAB export writes `A` in the reference's layout for every model count:
+`(nw, num_comps)` column-major, the bytes of the component-row `A` in C order.
+`loadmodout15.m` ignores this file; `load_results` reads it back
+and refuses an `A` that does not invert the `W` beside it (a multi-model directory written before this change).
+
+## Consequences
+
+- Every configuration without a merge is byte-identical to the previous code on every backend:
+  one, two and three models, `doscaling` and Newton on and off, `pdftype` 0 and 1 (NumPy supports 0 only), and `do_reject`,
+  pinned against the pre-change package loaded from git.
+  The weight-gradient norm `ndtmpsum` moves by float round-off (2.2e-16 relative in float64, one float32 unit on MLX):
+  it now sums squares per component row, as the reference sums per component.
+- Fits with `share_comps=True` where a merge fires change, toward the reference.
+  From a merged `load_comp_list` state the PyTorch and NumPy updates match the native binary to float64 round-off
+  (after 3 iterations, worst of `doscaling` on and off: `A` 1.5e-12, `mu` 3.7e-9, `sbeta` 1.8e-11, log-likelihood 4.7e-14, at the no-merge noise floor),
+  where the column semantics was off by 0.21 in `A` and 4.3e-4 in log-likelihood (worst of the same two).
+  End to end (2 models, 300 iterations, `comp_thresh=0.95`), the scan now merges three true pairs (map |cos| 0.956-0.971)
+  where the column metric merged pairs whose maps had |cos| 0.06, 0.35 and 0.55.
+- Merging is earlier and more complete when the models are still similar:
+  both models start near the identity, so a scan in the first iterations now merges most components
+  (it compares their true maps; on the sample with seed 42, 32 of 32 at iteration 8 with `comp_thresh=0.95`, 24 with 0.99),
+  and a model that loses its responsibility can collapse.
+  The reference behaves the same way: its similarity on its own early state merges exactly the pairs pamica's does,
+  and its update from the same merged states loses the second model's responsibility in step with pamica's, to NaN where pamica's goes non-finite.
+  The reference's default `share_start` of 100 avoids this; several short test recipes were retuned.
+- Saved models: unmerged saves load as before; saves with merged components must be refit.
+  Multi-model EEGLAB directories written before this change must be written again for `load_results`.
+- Initial-`A` normalization is not part of this change: the reference normalizes a drawn initial `A` and pamica did not,
+  which changes default trajectories and is issue #341 (epic #324 Phase 12, see ADR 0006).
+
+## Alternatives considered
+
+- **Keep the column layout and translate component ids to block rows inside the share code:**
+  rejected. A component shared by two models sits in a row of two different blocks in that layout,
+  so the tied object cannot be stored once, and every future per-component operation would need the same translation.
+- **Convert merged old saves by re-deriving the merge under the new metric:**
+  rejected. The old merge tied parameters that are not components, so there is no fitted state to convert, and a refit is required.
+- **Leave the EEGLAB `A` in the old C-order layout:** rejected. The file then had no meaning outside pamica for several models,
+  and the reference's layout is also what `load_results` should read from a reference run.
+
+## Receipts
+
+- amica15.f90:1749-1761 (`dAk/zeta` over `comp_list`), :1843-1851 (`doscaling` over components),
+  :1916-1963 (`identify_shared_comps`), :825-846 (`load_comp_list`, which reads the file named `c`).
+- Issues #334 (this change), #24 (the storage convention), #333 and ADR 0006 (`doscaling`), #341 (initial `A`), epic #324.
+- `pamica/tests/test_component_rows.py` (byte identity, semantics, export, the gated merged-state oracle),
+  `pamica/tests/test_component_rows_persistence.py` (round trips, conversion, refusal).