Skip to content

feat: AdaptPower regularization — power input + factor-2 scatter fix #511

Description

@Jammy2211

Overview

The Adapt regularization family puts its coefficient into the regularization matrix at the fourth power (adapt.py:45-47 squares, then adapt.py:84 / regularization_util.py:257,334 square again), while Constant squares exactly once. Both carry the identical LogUniform(1e-6, 1e6) prior, so the adaptive schemes explore a λ⁴ smoothing range and go numerically non-positive-definite from c ≈ 1e4 instead of Constant's c ≈ 1e6. That is the mechanism behind the likelihood-overflow flood diagnosed on RAL pilot 341908_5 (slam_source_pix_nn): fp64 Cholesky returned finite garbage, Nautilus accepted log_l up to 3e+303, and a 6-hour GPU run never terminated.

A second asymmetry sits in the same builder: weighted_regularization_matrix_from scatters every mesh edge in both directions, so Adapt(inner=outer=c) is exactly 2 × Constant(c) — contradicting its own docstring. (Verified on main: Adapt diag 4.00000001 vs Constant diag 2.00000001, ratio 2.0 on every entry, 4-connected 3×3 mesh.)

This task adds corrected sibling classes and leaves the legacy classes byte-for-byte, so stored identifiers, output/ directories and aggregator reloads all stay valid.

Plan

  • Give adapt_regularization_weights_from a power keyword defaulting to 2.0 — the legacy numerics, untouched.
  • Add a single-scatter matrix builder beside the existing one (numpy + JAX), which is the weighted graph Laplacian and reduces exactly to Constant for uniform weights.
  • Add four new sibling classes — AdaptPower, AdaptSplitPower, AdaptSplitZerothPower, MaternAdaptPowerKernel — each taking power: float = 1.0 (effective coefficient exponent 2 · power), so the default matches Constant's λ² convention and power=2.0 reproduces the legacy classes exactly.
  • Export them through aa.reg → ag.reg / al.reg.
  • Copy the PyAutoGalaxy prior yamls for the new classes, declaring power as a Constant prior so af.Model never samples it; mirror into the PyAutoLens test config.
  • Test the exact equalities (AdaptPower(c) == Constant(c), AdaptSplitPower(c) == ConstantSplit(c), power=2.0 == legacy), numpy/JAX parity, and that the af.Model identifier differs from the legacy class.
  • Document the λ⁴ legacy and the factor-2 scatter on the old docstrings; document the λ² convention and the c_new = c_old ** 2 migration on the new ones.
Detailed implementation plan

Affected Repositories

  • PyAutoArray (primary) — new classes, new util, tests, docstrings
  • PyAutoGalaxy — prior yamls, docs/api/pixelization.rst, one composition test
  • PyAutoLens — test_autolens/config/priors/regularization.yaml mirror, docs/api/pixelization.rst, one composition test

Branch Survey

Repository Current Branch Dirty?
./PyAutoArray main clean
./PyAutoGalaxy main clean
./PyAutoLens main clean

worktree_check_conflict reports PyAutoGalaxy + PyAutoLens claimed by lenscalc-adaptive-hessian-step. Its file sets (autogalaxy/operate/lens_calc.py, test_autogalaxy/operate/test_deflections.py, test_autolens/lens/test_multi_plane_cross_validation.py) are disjoint from this task's, so a parallel worktree is taken (own index, own branch) rather than a fold.

Suggested branch: feature/adapt-linear-regularization

Implementation Steps

  1. autoarray/inversion/regularization/adapt.py

    • adapt_regularization_weights_from(inner_coefficient, outer_coefficient, pixel_signals, power=2.0) → (...) ** power. Default 2.0 = existing behaviour, bit-identical.
    • New weighted_regularization_matrix_single_scatter_from(regularization_weights, neighbors, xp=np): builds the weighted graph Laplacian with edge weight 0.5 * (w_i**2 + w_j**2), scattering each ordered neighbour pair once (mat[i,i] += w_ij, mat[i,j] -= w_ij), plus the same 1e-8 diagonal floor. Symmetric, PSD (row sums zero), and exactly Constant's matrix when the weights are uniform. numpy and JAX branches, as the existing builder.
    • Docstring notes on Adapt: λ⁴ legacy convention; the factor-2 scatter (Adapt(inner=outer=1) is 2 × Constant(1), not equal to it); prefer AdaptPower.
  2. New modules (one class each, so the prior yaml filename matches the module — JSONPriorConfig keys config by module path):

    • adapt_power.py → class AdaptPower(Adapt) — power: float = 1.0; overrides regularization_weights_from (passes power) and regularization_matrix_from (single-scatter builder).
    • adapt_split_power.py → class AdaptSplitPower(AdaptSplit) — power; overrides regularization_weights_from only (the split builder is shared with ConstantSplit and already scatters once).
    • adapt_split_zeroth_power.py → class AdaptSplitZerothPower(AdaptSplitZeroth) — same, zeroth leg unchanged (BrightnessZeroth is already single-square).
    • matern_adapt_power_kernel.py → class MaternAdaptPowerKernel(MaternAdaptKernel) — same.
  3. autoarray/inversion/regularization/__init__.py: export all four. aa.reg is the module itself, and ag.reg / al.reg are from autoarray.inversion import regularization as reg, so the chain propagates with no downstream edit.

  4. PyAutoGalaxy autogalaxy/config/priors/regularization/{adapt_power,adapt_split_power,adapt_split_zeroth_power,matern_adapt_power_kernel}.yaml — copies of the legacy files with the class key renamed and power: {type: Constant, value: 1.0} added. PyAutoLens test_autolens/config/priors/regularization.yaml — mirror entries.

  5. Tests

    • test_autoarray/inversion/regularizations/test_adapt_power.py: AdaptPower(inner=outer=c) matrix == Constant(c) matrix; AdaptPower(power=2.0) == Adapt; weights unsquared at power=1.0; numpy/JAX parity of the new builder.
    • test_adapt_split_power.py / test_adapt_split_zeroth_power.py / test_matern_adapt_power_kernel.py: AdaptSplitPower(inner=outer=c) == ConstantSplit(c); power=2.0 reproduces the legacy class; numpy/JAX parity of the split path.
    • af.Model composition + identifier divergence from the legacy class.
    • PyAutoGalaxy + PyAutoLens: one af.Model(al.reg.AdaptSplitPower) composition test each (proves the prior config resolves).
  6. Docs: docs/api/pixelization.rst autosummary in PyAutoGalaxy and PyAutoLens gains AdaptPower / AdaptSplitPower.

Key Files

  • autoarray/inversion/regularization/adapt.py — power keyword + the single-scatter builder
  • autoarray/inversion/regularization/{adapt_split,adapt_split_zeroth,matern_adapt_kernel}.py — docstring notes only
  • autoarray/inversion/regularization/__init__.py — exports
  • autogalaxy/config/priors/regularization/*.yaml, test_autolens/config/priors/regularization.yaml

Trade-offs

  • New modules, not same-module siblings. autonerves.json_prior.config keys prior config by cls.__module__ + cls.__name__ and matches the config dict path built from the yaml filename, so a adapt_power.yaml only resolves for a class living in adapt_power.py. Same-module classes would silently fall back to the legacy Adapt config via the family() base-class walk and leave the new yaml an orphan.
  • The factor-2 fix is only in the new class. Halving weighted_regularization_matrix_from in place would silently change the effective regularization of every Adapt fit ever run. Filed as documentation-only: PyAutoMind/draft/bug/autoarray/adapt_scatter_factor_two.md.
  • power is a Constant prior, not a free parameter. It is a convention switch, not something to sample; leaving it free would let a search wander between λ² and λ⁴ mid-fit.

Original Prompt

Click to expand starting prompt

do the 5 things above listed and make sure we have some SMC runs going soon

Rung 2 of the approved plan "Overflow-flood fix wave + SMC on the A100" (task A3), as amended 2026-08-29 by the human: "maximal changes, power as an input". Full prompt: PyAutoMind/active/adapt_linear_regularization.md.

Generated by the PyAutoLabs agent workflow.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions