From c8116621ca979298fff18dda613cd56b7d578129 Mon Sep 17 00:00:00 2001 From: Seyed Yahya Shirazi Date: Tue, 1 Sep 2026 07:05:38 -0700 Subject: [PATCH 1/2] Surface MLX feature parity in user-facing docs Post-epic-#278 discoverability pass: the AGENTS.md map annotation still said GG-only; the EEGLAB guide never mentioned the MLX export path; the backends table row and README undersold the now-complete MLX surface. Docs only. --- AGENTS.md | 2 +- README.md | 4 +++- docs/guides/backends.md | 2 +- docs/guides/eeglab.md | 12 ++++++++++++ 4 files changed, 17 insertions(+), 3 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 5fcf278c..9543fd83 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -14,7 +14,7 @@ pamica/ │ ├── core.py # Natural-gradient EM port (AMICATorchNG); Fortran-parity, primary backend │ └── utils.py # Preprocessing (sphering, PCA), device selection ├── mlx_impl/ # Optional MLX backend (Apple GPU; AMICAMLXNG, #76/#81) -│ └── core.py # float32 GPU E/M-step + CPU-stream linalg (single- & multi-model GG, NG) +│ └── core.py # float32 GPU E/M-step + CPU-stream linalg; all five pdf families, full torch-equivalent surface (epic #278) ├── numpy_impl/ # Legacy NumPy reference (topic-named modules, issue #34) │ ├── core.py # AMICA_NumPy (incl. inlined Newton); pdf.py, data.py, load.py, viz.py, utils.py, cli.py │ └── ... diff --git a/README.md b/README.md index 765f03b2..c1e98e34 100644 --- a/README.md +++ b/README.md @@ -68,7 +68,9 @@ The wrapper auto-selects a device and computes in float64 for Fortran parity. - float32 (about 7 significant digits, not parity) is required on the Apple GPUs and modestly faster on CPU; it is not a general speedup, since CUDA is overhead-bound (float32 is about as fast as float64). -- On Apple Silicon the MLX backend is the fastest option; import it explicitly. +- On Apple Silicon the MLX backend is the fastest option and carries the full + feature surface (all pdf families, Newton, rejection, EEGLAB export, MIR); + import it explicitly. ```python AMICA(device="cuda").fit(X) # NVIDIA GPU, float64 diff --git a/docs/guides/backends.md b/docs/guides/backends.md index bb38d9ed..e3faed20 100644 --- a/docs/guides/backends.md +++ b/docs/guides/backends.md @@ -8,7 +8,7 @@ interface, plus an optional Apple-GPU backend and a legacy NumPy reference. | Backend | Class | Role | |---|---|---| | PyTorch natural-gradient EM | [`AMICATorchNG`](../api/torch-backend.md) | **Default.** Fortran-parity backend; CUDA / CPU, and float32 on MPS. | -| MLX (Apple GPU) | [`AMICAMLXNG`](../api/mlx-backend.md) (`pamica.mlx_impl`) | Optional Apple-Silicon GPU backend; float32 only. | +| MLX (Apple GPU) | [`AMICAMLXNG`](../api/mlx-backend.md) (`pamica.mlx_impl`) | Optional Apple-Silicon GPU backend; float32 only. Full torch-equivalent feature surface as of epic #278, including EEGLAB export -- see the [backend table](amica-differences.md#backend-differences). | | NumPy reference | [`AMICA_NumPy`](../api/numpy-backend.md) | Legacy oracle + CLI; carries the same parity fixes. | The `AMICA` wrapper uses `AMICATorchNG`. The MLX backend is imported separately diff --git a/docs/guides/eeglab.md b/docs/guides/eeglab.md index cb3d0730..5b298ef8 100644 --- a/docs/guides/eeglab.md +++ b/docs/guides/eeglab.md @@ -34,6 +34,18 @@ This writes the raw binary files EEGLAB's AMICA loader reads: For a single model the bytes are identical to the reference Fortran binary's `amicaout` files, so the directory is interchangeable with a native AMICA run. +The MLX backend writes the same directory: an Apple-Silicon fit exports to +EEGLAB directly, with no torch round trip (epic #278; the export is +byte-compared against a torch twin in the test suite): + +```python +from pamica.mlx_impl import AMICAMLXNG # requires the mlx extra + +model = AMICAMLXNG(n_channels=X.shape[0], n_mix=3) +model.fit(X) +model.write_amica_output("amicaout") +``` + `LLt` is what `loadmodout15.m` turns into `Lht`/`Lt` and the model-probability odds `v`; it is written after a fresh `fit()`, and omitted (with a warning) for a model restored from `load()`, which carries no E-step stash. Under From f4fef32923588a6cfa61d75393daa42e8d9bba2c Mon Sep 17 00:00:00 2001 From: Seyed Yahya Shirazi Date: Tue, 1 Sep 2026 07:13:04 -0700 Subject: [PATCH 2/2] Address docs review: precision claims and naming The export-validation claim now states float32 precision rather than byte comparison (only comp_list is exact), and the test that seeded the overclaim is renamed to match its actual assertion; the backends table row returns to terse form with the parity statement moved to prose that names the wrapper-only from_params_file caveat (#313); MIR is spelled out on first use in the README; the stale GG-only class one-liner in mlx core is corrected. Export suite 12 passed. --- README.md | 4 ++-- docs/guides/backends.md | 9 +++++++-- docs/guides/eeglab.md | 9 ++++++--- pamica/mlx_impl/core.py | 2 +- pamica/tests/mlx_tests/test_mlx_export.py | 13 ++++++++----- 5 files changed, 24 insertions(+), 13 deletions(-) diff --git a/README.md b/README.md index c1e98e34..42b06901 100644 --- a/README.md +++ b/README.md @@ -69,8 +69,8 @@ The wrapper auto-selects a device and computes in float64 for Fortran parity. and modestly faster on CPU; it is not a general speedup, since CUDA is overhead-bound (float32 is about as fast as float64). - On Apple Silicon the MLX backend is the fastest option and carries the full - feature surface (all pdf families, Newton, rejection, EEGLAB export, MIR); - import it explicitly. + feature surface (all pdf families, Newton, rejection, EEGLAB export, and + Mutual Information Reduction (MIR) diagnostics); import it explicitly. ```python AMICA(device="cuda").fit(X) # NVIDIA GPU, float64 diff --git a/docs/guides/backends.md b/docs/guides/backends.md index e3faed20..e89c3265 100644 --- a/docs/guides/backends.md +++ b/docs/guides/backends.md @@ -8,12 +8,17 @@ interface, plus an optional Apple-GPU backend and a legacy NumPy reference. | Backend | Class | Role | |---|---|---| | PyTorch natural-gradient EM | [`AMICATorchNG`](../api/torch-backend.md) | **Default.** Fortran-parity backend; CUDA / CPU, and float32 on MPS. | -| MLX (Apple GPU) | [`AMICAMLXNG`](../api/mlx-backend.md) (`pamica.mlx_impl`) | Optional Apple-Silicon GPU backend; float32 only. Full torch-equivalent feature surface as of epic #278, including EEGLAB export -- see the [backend table](amica-differences.md#backend-differences). | +| MLX (Apple GPU) | [`AMICAMLXNG`](../api/mlx-backend.md) (`pamica.mlx_impl`) | Optional Apple-Silicon GPU backend; float32 only. | | NumPy reference | [`AMICA_NumPy`](../api/numpy-backend.md) | Legacy oracle + CLI; carries the same parity fixes. | The `AMICA` wrapper uses `AMICATorchNG`. The MLX backend is imported separately (`from pamica.mlx_impl import AMICAMLXNG`) so that `import pamica` never -requires MLX. +requires MLX. As of epic #278 the MLX backend carries the full +`AMICATorchNG`-equivalent surface (all pdf families, Newton, rejection, +persistence, EEGLAB export, MIR/PMI); the wrapper conveniences -- among them +the `from_params_file` reader -- remain PyTorch-only (issue #313 tracks +backend selection there). See the +[backend table](amica-differences.md#backend-differences). ## Device selection diff --git a/docs/guides/eeglab.md b/docs/guides/eeglab.md index 5b298ef8..54307bc4 100644 --- a/docs/guides/eeglab.md +++ b/docs/guides/eeglab.md @@ -35,14 +35,17 @@ For a single model the bytes are identical to the reference Fortran binary's `amicaout` files, so the directory is interchangeable with a native AMICA run. The MLX backend writes the same directory: an Apple-Silicon fit exports to -EEGLAB directly, with no torch round trip (epic #278; the export is -byte-compared against a torch twin in the test suite): +EEGLAB directly, with no torch round trip (epic #278). The export is +validated against a torch twin in the test suite to float32 precision +(max abs diff < 1e-5; `comp_list` is the only field compared exactly -- +MLX computes in float32, so its files are not bit-identical to a float64 +export the way a torch fit's are to Fortran's): ```python from pamica.mlx_impl import AMICAMLXNG # requires the mlx extra model = AMICAMLXNG(n_channels=X.shape[0], n_mix=3) -model.fit(X) +model.fit(X) # X is (n_channels, n_samples) model.write_amica_output("amicaout") ``` diff --git a/pamica/mlx_impl/core.py b/pamica/mlx_impl/core.py index cf0299e6..f2c2ff13 100644 --- a/pamica/mlx_impl/core.py +++ b/pamica/mlx_impl/core.py @@ -328,7 +328,7 @@ def _safe_int_cast(name: str, value: np.ndarray, dtype) -> np.ndarray: class AMICAMLXNG: - """MLX natural-gradient EM backend (GG, single- and multi-model; #76/#81). + """MLX natural-gradient EM backend, full torch-equivalent surface (#76/#81, epic #278). Parameters mirror the subset of :class:`AMICATorchNG` that is supported; the same ``seed`` produces the same initial parameters as the PyTorch/NumPy diff --git a/pamica/tests/mlx_tests/test_mlx_export.py b/pamica/tests/mlx_tests/test_mlx_export.py index 62a2f80f..143f0795 100644 --- a/pamica/tests/mlx_tests/test_mlx_export.py +++ b/pamica/tests/mlx_tests/test_mlx_export.py @@ -4,10 +4,11 @@ Round-trips through the shared :func:`pamica.numpy_impl.load.loadmodout` reader (the same reader real EEGLAB output is validated against, issue -#159), plus a torch-twin byte comparison: a float64 ``AMICATorchNG`` built +#159), plus a torch-twin file comparison: a float64 ``AMICATorchNG`` built from the exact fitted MLX state (the ``_torch_twin`` pattern from ``pamica/tests/test_mlx_sharing_cross_backend.py``) exports the same -parameters, so the two on-disk files can be diffed directly. This embeds an +parameters, so the on-disk files agree to float32 precision (only the +integer ``comp_list`` is compared exactly). This embeds an inline torch comparison in an ``mlx_tests`` file, the same precedent ``test_mlx_backend.py``/``test_mlx_sharing.py`` already set for non-drift- guard torch comparisons; the anti-drift AGREEMENT pin for ``do_reject`` @@ -85,7 +86,7 @@ def test_write_amica_output_round_trips_through_loadmodout( # #159), so only order-independent quantities are checked through this # reader -- exact byte-level reproduction (raw files, no reordering) is # pinned separately below, against a torch twin - # (test_export_is_byte_compatible_with_a_torch_twin), mirroring + # (test_export_matches_a_torch_twin_to_float32_precision), mirroring # torch_tests/test_amica_ng_wrapper.py's split between # test_write_amica_output_bytes (raw bytes) and # test_write_amica_output_loadmodout_readable (order-independent, via @@ -277,7 +278,7 @@ def test_write_amica_output_makes_no_extra_forward_pass( assert (tmp_path / "amicaout" / "LLt").exists() -# --- torch-twin byte comparison (W layout, single-model Fortran parity) ---- +# --- torch-twin file comparison (W layout, single-model Fortran parity) ---- def _torch_twin(model, sphere_np: np.ndarray): """A float64 AMICATorchNG holding ``model``'s exact fitted state, for a direct file-level export comparison. Mirrors @@ -334,7 +335,9 @@ def _torch_twin(model, sphere_np: np.ndarray): @pytest.mark.parametrize("n_models", [1, 2]) -def test_export_is_byte_compatible_with_a_torch_twin(real_data, tmp_path, n_models): +def test_export_matches_a_torch_twin_to_float32_precision( + real_data, tmp_path, n_models +): m = _model(n_models=n_models, seed=42, keep_best=False) m.fit(real_data, max_iter=6, verbose=False) assert m._sphere_np is not None