test(inversion): certified active-set solver numerics, fallback, jit/vmap, gradients, dispatch, settings

Refs #566

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

Author: Jammy2211

test_autoarray/inversion/inversion/test_inversion_util.py

@@ -384,3 +384,160 @@ def test__curvature_matrix_via_mapping_matrix_from__mixed_precision__weights_are
     )
 
     assert jax_curvature == pytest.approx(expected, rel=1.0e-14)
+
+
+# ----------------------------------------------------------------------------
+# The `solver` switch of `reconstruction_positive_only_from` (PyAutoArray#566).
+# ----------------------------------------------------------------------------
+
+
+def _positive_only_system(n=30, seed=7):
+    """An SPD system whose unconstrained solution has negative entries, so the constraint binds."""
+    rng = np.random.default_rng(seed)
+    A = rng.normal(size=(3 * n, n))
+    b = rng.normal(size=3 * n)
+    curvature_reg_matrix = A.T @ A
+    data_vector = A.T @ b
+    assert (np.linalg.solve(curvature_reg_matrix, data_vector) < 0.0).any()
+    return data_vector, curvature_reg_matrix
+
+
+@requires_jax
+@pytest.mark.parametrize("use_jacobi", [True, False])
+def test__reconstruction_positive_only_from__certified_matches_pdip_on_jax(
+    use_jacobi, monkeypatch
+):
+    import jax
+
+    jax.config.update("jax_enable_x64", True)
+    import jax.numpy as jnp
+    from autonerves import conf
+
+    if not use_jacobi:
+        # Exercise the unpreconditioned branch too: the `solver` switch lives in both. The dict stands in for
+        # the config; the Settings' certified_* keys resolve through their KeyError defaults.
+        monkeypatch.setattr(
+            conf,
+            "instance",
+            {
+                "general": {
+                    "inversion": {
+                        "nnls_jacobi_preconditioning": False,
+                        "nnls_target_kappa": 1.0e-11,
+                    }
+                }
+            },
+        )
+
+    for data_vector, curvature_reg_matrix in [
+        (
+            np.array([1.0, 1.0, 2.0]),
+            np.array([[2.0, 1.0, 0.0], [1.0, 3.0, 1.0], [0.0, 1.0, 1.0]]),
+        ),
+        _positive_only_system(),
+    ]:
+        pdip = np.asarray(
+            aa.util.inversion.reconstruction_positive_only_from(
+                data_vector=jnp.asarray(data_vector),
+                curvature_reg_matrix=jnp.asarray(curvature_reg_matrix),
+                settings=aa.Settings(),
+                xp=jnp,
+            )
+        )
+
+        stats = {}
+        certified = np.asarray(
+            aa.util.inversion.reconstruction_positive_only_from(
+                data_vector=jnp.asarray(data_vector),
+                curvature_reg_matrix=jnp.asarray(curvature_reg_matrix),
+                settings=aa.Settings(),
+                xp=jnp,
+                solver="certified",
+                stats=stats,
+            )
+        )
+
+        assert stats["solver"] == "certified"
+        assert bool(stats["certified"])
+        assert int(stats["passes"]) >= 1
+
+        assert (certified >= 0.0).all()
+        assert certified == pytest.approx(
+            pdip, rel=1.0e-8, abs=1.0e-8 * np.max(np.abs(certified))
+        )
+
+
+@requires_jax
+def test__reconstruction_positive_only_from__certified_budget_exhausted_falls_back_to_pdip():
+    import jax
+
+    jax.config.update("jax_enable_x64", True)
+    import jax.numpy as jnp
+
+    data_vector, curvature_reg_matrix = _positive_only_system()
+
+    kwargs = dict(
+        data_vector=jnp.asarray(data_vector),
+        curvature_reg_matrix=jnp.asarray(curvature_reg_matrix),
+        xp=jnp,
+    )
+
+    pdip = np.asarray(
+        aa.util.inversion.reconstruction_positive_only_from(
+            settings=aa.Settings(), **kwargs
+        )
+    )
+
+    stats = {}
+    fallback = np.asarray(
+        aa.util.inversion.reconstruction_positive_only_from(
+            settings=aa.Settings(certified_pass_budget=0),
+            solver="certified",
+            stats=stats,
+            **kwargs,
+        )
+    )
+
+    assert not bool(stats["certified"])
+    assert np.array_equal(fallback, pdip)
+
+
+def test__reconstruction_positive_only_from__numpy_path_ignores_solver(monkeypatch):
+    # The NumPy path must run fnnls whatever `solver` says: spy on it.
+    from autoarray.util import fnnls
+
+    calls = []
+    original = fnnls.fnnls_cholesky
+
+    def spy(*args, **kwargs):
+        calls.append(1)
+        return original(*args, **kwargs)
+
+    monkeypatch.setattr(fnnls, "fnnls_cholesky", spy)
+
+    data_vector = np.array([1.0, 1.0, 2.0])
+    curvature_reg_matrix = np.array([[2.0, 1.0, 0.0], [1.0, 3.0, 1.0], [0.0, 1.0, 1.0]])
+
+    results = []
+    for solver in ["pdip", "certified"]:
+        calls.clear()
+        results.append(
+            aa.util.inversion.reconstruction_positive_only_from(
+                data_vector=data_vector,
+                curvature_reg_matrix=curvature_reg_matrix,
+                solver=solver,
+            )
+        )
+        assert len(calls) == 1
+
+    assert np.array_equal(results[0], results[1])
+    assert results[0] == pytest.approx(np.array([0.5, 0.0, 2.0]), 1.0e-4)
+
+
+def test__reconstruction_positive_only_from__invalid_solver_raises():
+    with pytest.raises(ValueError, match="solver"):
+        aa.util.inversion.reconstruction_positive_only_from(
+            data_vector=np.array([1.0]),
+            curvature_reg_matrix=np.array([[1.0]]),
+            solver="bogus",
+        )

test_autoarray/inversion/inversion/test_positive_only_dispatch.py

@@ -0,0 +1,211 @@
+"""
+Which positive-only solver `AbstractInversion.reconstruction` dispatches to (PyAutoArray#566).
+
+`Settings.positive_only_solver = "certified"` selects the certified active-set solve only on the JAX
+backend for mapper-only inversions; inversions with linear light-profile / MGE coefficients
+(`AbstractLinearObjFuncList`) and every NumPy inversion keep today's solver.
+"""
+
+import importlib.util
+
+import numpy as np
+import pytest
+
+import autoarray as aa
+
+
+requires_jax = pytest.mark.skipif(
+    importlib.util.find_spec("jax") is None,
+    reason="requires jax (installed via the [optional] extras; absent on the NumPy-only matrix env)",
+)
+
+
+def _system(n, seed=3):
+    """An SPD system whose unconstrained solution has negative entries, so positivity binds."""
+    rng = np.random.default_rng(seed)
+    A = rng.normal(size=(3 * n, n))
+    b = rng.normal(size=3 * n)
+    curvature_reg_matrix = A.T @ A
+    data_vector = A.T @ b
+    assert (np.linalg.solve(curvature_reg_matrix, data_vector) < 0.0).any()
+    return data_vector, curvature_reg_matrix
+
+
+def _mapper(pixels=16):
+    return aa.m.MockMapper(
+        mesh=aa.mesh.RectangularUniform(shape=(4, 4)),
+        parameters=pixels,
+        regularization=aa.reg.Constant(),
+        # Only its shape is read, by the NumPy warm-start memo's fingerprint.
+        source_plane_mesh_grid=np.zeros((pixels, 2)),
+    )
+
+
+def _inversion(
+    linear_obj_list,
+    data_vector,
+    curvature_reg_matrix,
+    use_jax,
+    positive_only_solver="certified",
+    use_edge_zeroed_pixels=False,
+):
+    if use_jax:
+        import jax
+
+        jax.config.update("jax_enable_x64", True)
+        import jax.numpy as jnp
+
+        data_vector = jnp.asarray(data_vector)
+        curvature_reg_matrix = jnp.asarray(curvature_reg_matrix)
+
+    inversion = aa.m.MockInversion(
+        linear_obj_list=linear_obj_list,
+        data_vector=data_vector,
+        curvature_reg_matrix=curvature_reg_matrix,
+        settings=aa.Settings(
+            use_positive_only_solver=True,
+            use_edge_zeroed_pixels=use_edge_zeroed_pixels,
+            positive_only_solver=positive_only_solver,
+            # Off so two NumPy solves of the same system are bit-comparable: with the memo on, the second
+            # solve starts from the first's passive set and can differ in the last ulp.
+            nnls_warm_start_memo=False,
+        ),
+    )
+    # `MockInversion` does not take `xp`; the backend flag is what `AbstractInversion(xp=jnp)` sets.
+    inversion.use_jax = use_jax
+    return inversion
+
+
+def test__positive_only_solver_used__default_is_pdip():
+    data_vector, curvature_reg_matrix = _system(16)
+
+    inversion = _inversion(
+        [_mapper()],
+        data_vector,
+        curvature_reg_matrix,
+        use_jax=False,
+        positive_only_solver=None,
+    )
+
+    assert inversion.settings.positive_only_solver == "pdip"
+    assert inversion.positive_only_solver_used == "pdip"
+
+
+def test__positive_only_solver_used__numpy_backend_keeps_pdip_and_fnnls():
+    data_vector, curvature_reg_matrix = _system(16)
+
+    certified = _inversion(
+        [_mapper()], data_vector, curvature_reg_matrix, use_jax=False
+    )
+    default = _inversion(
+        [_mapper()],
+        data_vector,
+        curvature_reg_matrix,
+        use_jax=False,
+        positive_only_solver="pdip",
+    )
+
+    assert certified.positive_only_solver_used == "pdip"
+    assert np.array_equal(certified.reconstruction, default.reconstruction)
+
+
+@requires_jax
+def test__positive_only_solver_used__mapper_only_jax_is_certified_and_matches_pdip():
+    data_vector, curvature_reg_matrix = _system(16)
+
+    certified = _inversion([_mapper()], data_vector, curvature_reg_matrix, use_jax=True)
+    pdip = _inversion(
+        [_mapper()],
+        data_vector,
+        curvature_reg_matrix,
+        use_jax=True,
+        positive_only_solver="pdip",
+    )
+
+    assert certified.positive_only_solver_used == "certified"
+    assert pdip.positive_only_solver_used == "pdip"
+
+    x_certified = np.asarray(certified.reconstruction)
+    x_pdip = np.asarray(pdip.reconstruction)
+
+    assert (x_certified >= 0.0).all()
+    assert (x_certified == 0.0).any()
+    assert x_certified == pytest.approx(
+        x_pdip, rel=1.0e-8, abs=1.0e-8 * np.max(np.abs(x_certified))
+    )
+
+
+@requires_jax
+def test__positive_only_solver_used__linear_func_list_present_keeps_pdip():
+    data_vector, curvature_reg_matrix = _system(17)
+
+    func_list = aa.m.MockLinearObjFuncList(parameters=1)
+
+    certified_setting = _inversion(
+        [func_list, _mapper()], data_vector, curvature_reg_matrix, use_jax=True
+    )
+    pdip = _inversion(
+        [func_list, _mapper()],
+        data_vector,
+        curvature_reg_matrix,
+        use_jax=True,
+        positive_only_solver="pdip",
+    )
+
+    assert certified_setting.positive_only_solver_used == "pdip"
+    assert np.array_equal(
+        np.asarray(certified_setting.reconstruction), np.asarray(pdip.reconstruction)
+    )
+
+
+@requires_jax
+def test__positive_only_solver_used__no_mapper_keeps_pdip():
+    data_vector, curvature_reg_matrix = _system(4)
+
+    inversion = _inversion(
+        [aa.m.MockLinearObjFuncList(parameters=4)],
+        data_vector,
+        curvature_reg_matrix,
+        use_jax=True,
+    )
+
+    assert inversion.positive_only_solver_used == "pdip"
+
+
+@requires_jax
+def test__positive_only_solver_used__edge_zeroed_subset_is_preserved():
+    # The 4x4 rectangular mesh's 12 edge pixels are zeroed and the 4 interior pixels [5, 6, 9, 10] solved.
+    data_vector, curvature_reg_matrix = _system(16)
+
+    certified = _inversion(
+        [_mapper()],
+        data_vector,
+        curvature_reg_matrix,
+        use_jax=True,
+        use_edge_zeroed_pixels=True,
+    )
+    pdip = _inversion(
+        [_mapper()],
+        data_vector,
+        curvature_reg_matrix,
+        use_jax=True,
+        positive_only_solver="pdip",
+        use_edge_zeroed_pixels=True,
+    )
+
+    assert certified.positive_only_solver_used == "certified"
+    assert np.asarray(certified.solve_ids_to_keep) == pytest.approx(
+        np.array([5, 6, 9, 10])
+    )
+
+    x_certified = np.asarray(certified.reconstruction)
+    x_pdip = np.asarray(pdip.reconstruction)
+
+    keep = np.array([5, 6, 9, 10])
+    edge = np.setdiff1d(np.arange(16), keep)
+
+    assert (x_certified[edge] == 0.0).all()
+    assert (x_pdip[edge] == 0.0).all()
+    assert x_certified == pytest.approx(
+        x_pdip, rel=1.0e-8, abs=1.0e-8 * np.max(np.abs(x_certified))
+    )