refactor: mass profiles — MGE/CSE fallback, JAX port, testing & docs

[Jammy2211]

Overview

The mass profiles module (autogalaxy/profiles/mass/) needs a phased refactoring to address three problems: 16 profiles return zeros for convergence_2d_from or potential_2d_from (placeholders that should compute the quantity via MGE/CSE decomposition), the CSE module is not JAX-compatible, and docstrings lack mathematical definitions and paper references.

This is an umbrella issue tracking the full effort. Individual sub-tasks will be filed as separate issues and linked here as they're started.

Plan

• Phase 1 — Self-consistency test suite in autolens_workspace_test: per-category scripts verifying lensing relations (∇·α = 2κ, ∇ψ = α, ∇²ψ = 2κ) for every mass profile using numerical differentiation independent of the source code. Safety net for all subsequent changes.
• Phase 2 — CSE JAX port: thread xp=np through all MassProfileCSE methods, replace scipy.linalg.lstsq with JAX-compatible path, cache decomposition results.
• Phase 3 — MGE/CSE fallback mechanism: replace zero-returning methods with automatic fallback to MGE or CSE decomposition. Analytic formulas for PointMass and MassSheet potentials.
• Phase 4 — Spring clean: xp threading audit across utility files, dead code removal, decorator consistency, parameter naming standardization.
• Phase 5 — LaTeX documentation: first-class docstrings with mathematical definitions, paper references (Oguri 2021, Shajib 2019, Tessore 2015/2016), parameter units, and updated docs/api/mass.rst.

Profiles returning zeros (Phase 3 targets)

Profile	convergence_2d_from	potential_2d_from
ExternalShear	zeros (physically correct)	—
PowerLawMultipole	zeros (physically correct)	—
cNFW / cNFWSph	zeros	zeros
NFWTruncatedSph	—	zeros
PointMass	—	zeros
MassSheet	—	zeros
Chameleon	—	zeros
Gaussian	—	zeros
Sersic (AbstractSersic)	—	zeros
dPIEPotential / Sph	—	zeros
PowerLawBroken	—	zeros
dPIEMass	—	zeros

Dependency chain

Phase 1 (tests) → Phase 2 (CSE JAX) → Phase 3 (fallback) → Phase 4 (clean) → Phase 5 (docs)

Strictly serial — each phase's test suite or output is a prerequisite for the next.

Original Prompt

Click to expand starting prompt

We are going to refactor the mass profile module to work with the following assumptions:

1. All mass profiles must have a deflections_yx_2d_from method defined, as they do now.
2. The best method they can use is an analytic function, however this could instead be an MGE decomposition or CSE decomposition.
3. The potential_2d_from and convergence_2d_from methods can, optionally, be defined inside mass profiles. This is common if they have analytic solutions which are quick to compute.

However, a potential or convergence does not need to be defined analytically. In this case, we will "fall back to an MGE or CSE method", which computes the potential via those functions. These methods still need a dedicated convergence function for the decomposition though, which not all functions have. Therefore, we will draw up a list of mass profiles which currently need this MGE / CSE treatment applied to them. The red flag to catch this is those where the potential_2d_from or convergence_2d_from method returns a numpy array of zeros.

Do deep research on autogalaxy/profiles/mass to fully understand this design. Think hard about the transformer and data structure decorators. In autolens_workspace_test, design a test set in scripts/mass which runs on every mass profile individually, using methods / algebraic equations independent of the source code implementation, verifies their convergence, potential, deflections, shear and any other key quantity are fully self consistent INTERNALLY WITH THEMSELVES.

The MGE mass profile module is mature and allows us to compute deflection angles using a multi-Gaussian expansion. We will look at this module in more detail for any refactor / design choices, but most importantly then redesign the CSE module to also be fully JAX'd.

Check that the CSE and MGE have fast and robust potential_2d_from methods.

The first milestone is to have the really extensive testing structure in autolens_workspace_test up and running. Many methods in the source code are bypassed due laziness, often having an np.zeros for their convergence or potential. The MGE / CSE trick will sort these but look out for it.

A final task will be first class documentation on all functions, including latex readable code in docstrings. Also, a feature should be a "spring clean" of dodgy code and other improvements.

Break this down into phased tasks in z_features.

[Jammy2211]

Epic Complete — 2026-05-26

All 5 phases shipped in a single session:

Phase	PR	Summary
1 — Self-consistency tests	autolens_workspace_test#124	7 scripts testing 42 profiles via numerical differentiation of lensing relations
2 — CSE JAX port	#447	Threaded xp=np through all MassProfileCSE methods
3 — MGE/CSE fallback	#449	Replaced 16 zero-returning methods with proper implementations (12 files, 271 insertions)
4 — Spring clean	#450 (audit-only)	Codebase already clean after Phases 1-3; no changes needed
5 — LaTeX docs	#453	Sphinx docstrings for all 26 mass profile files (1283 lines)

Key outcomes

• Every mass profile now returns physically meaningful values for convergence and potential (no more zeros)
• MGE potential uses E1 exponential integral formula (Shajib 2019), self-consistent to 0.3%
• CSE module is JAX-ready
• All profiles have LaTeX mathematical definitions and paper references
• Phase 1 test suite found a genuine NFWSph potential-deflection mismatch (11% median) — future investigation