Merge branch 'main' into feature/two_galaxies_lens

Author: NiekWielders

CLAUDE.md

@@ -78,6 +78,32 @@ Each dataset type has an `Analysis*` class that implements `log_likelihood_funct
 
 These inherit from `AnalysisDataset` → `Analysis` (in `analysis/analysis/`), which inherits `af.Analysis`. The `log_likelihood_function` builds a `Fit*` object from the `af.ModelInstance` and returns its `figure_of_merit`.
 
+### Decorator System (from autoarray)
+
+Profile methods that consume a grid and return an array, grid, or vector use decorators from `autoarray.structures.decorators`. These ensure the **output type matches the input grid type**:
+
+| Decorator | `Grid2D` input | `Grid2DIrregular` input |
+|---|---|---|
+| `@aa.grid_dec.to_array` | `Array2D` | `ArrayIrregular` |
+| `@aa.grid_dec.to_grid` | `Grid2D` | `Grid2DIrregular` |
+| `@aa.grid_dec.to_vector_yx` | `VectorYX2D` | `VectorYX2DIrregular` |
+
+The `@aa.grid_dec.transform` decorator (always stacked below the output decorator) shifts and rotates the grid to the profile's reference frame before passing it to the function body.
+
+The canonical stacking order is:
+```python
+@aa.grid_dec.to_array      # outermost: wraps output
+@aa.grid_dec.transform     # innermost: transforms grid
+def convergence_2d_from(self, grid, xp=np, **kwargs):
+    y = grid.array[:, 0]   # use .array to get raw numpy/jax array
+    x = grid.array[:, 1]
+    return ...             # return raw array; decorator wraps it
+```
+
+**Key rule**: the function body must return a **raw array** (not an autoarray). The decorator handles wrapping. Access grid coordinates via `grid.array[:, 0]` / `grid.array[:, 1]` (not `grid[:, 0]`), because after `@transform` the grid is still an autoarray object and `.array` is the safe way to extract the underlying data for both numpy and jax backends.
+
+See PyAutoArray's `CLAUDE.md` for full details on the decorator internals.
+
 ### JAX Support
 
 The codebase is designed so that **NumPy is the default everywhere and JAX is opt-in**. JAX is never imported at module level — it is only imported locally inside functions when explicitly requested.
@@ -100,6 +126,28 @@ When adding a new function that should support JAX:
 4. Add a JAX implementation in the guarded branch (e.g. `jax.jacfwd`, `jnp.vectorize`)
 5. Verify correctness by comparing both paths in `autogalaxy_workspace_test/scripts/`
 
+### JAX and autoarray wrappers at the `jax.jit` boundary
+
+Autoarray types (`Array2D`, `ArrayIrregular`, `VectorYX2DIrregular`, etc.) are **not registered as JAX pytrees**. This means:
+
+- Constructing them **inside** a JIT trace is fine (Python code runs normally during tracing)
+- **Returning** them as the output of a `jax.jit`-compiled function **fails** with `TypeError: ... is not a valid JAX type`
+
+Functions decorated with `@aa.grid_dec.to_array` / `@to_vector_yx` wrap their return value in an autoarray type. This wrapping is safe for intermediate calls (the autoarray object is consumed by downstream Python code). However, if such a function is the **outermost call** inside a `jax.jit` lambda, its return value will fail at the JIT boundary.
+
+The solution is the **`if xp is np:` guard** in the function body:
+
+```python
+def convergence_2d_via_hessian_from(self, grid, xp=np):
+    convergence = 0.5 * (hessian_yy + hessian_xx)
+
+    if xp is np:
+        return aa.ArrayIrregular(values=convergence)  # numpy: wrapped
+    return convergence                                  # jax: raw jax.Array
+```
+
+This pattern is applied throughout `autogalaxy/operate/lens_calc.py`. Functions that are only ever called as intermediate steps (e.g. `deflections_yx_2d_from`) do NOT need this guard — their autoarray wrappers are never the JIT output.
+
 ### Linear Light Profiles & Inversions
 
 `LightProfileLinear` subclasses do not take an `intensity` parameter—it is solved via a linear inversion (provided by `autoarray`). The `GalaxiesToInversion` class (`galaxy/to_inversion.py`) handles converting galaxies with linear profiles or pixelizations into the inversion objects needed by `autoarray`.
@@ -147,3 +195,22 @@ When importing `autogalaxy as ag`:
 - `ag.ps.*` – point sources
 - `ag.Galaxy`, `ag.Galaxies`
 - `ag.FitImaging`, `ag.AnalysisImaging`, `ag.SimulatorImaging`
+
+## Line Endings — Always Unix (LF)
+
+All files in this project **must use Unix line endings (LF, `\n`)**. Windows/DOS line endings (CRLF, `\r\n`) will break Python files on HPC systems.
+
+**When writing or editing any file**, always produce Unix line endings. Never write `\r\n` line endings.
+
+After creating or copying files, verify and convert if needed:
+
+```bash
+# Check for DOS line endings
+file autogalaxy/galaxy/galaxy.py   # should say "ASCII text", not "CRLF"
+
+# Convert all Python files in the project
+find . -type f -name "*.py" | xargs dos2unix
+```
+
+Prefer simple shell commands.
+Avoid chaining with && or pipes.

autogalaxy/config/priors/mass/dark/cnfw.yaml

@@ -1,3 +1,78 @@
+cNFW:
+  centre_0:
+    type: Gaussian
+    mean: 0.0
+    sigma: 0.1
+    width_modifier:
+      type: Absolute
+      value: 0.05
+    limits:
+      lower: -inf
+      upper: inf
+  centre_1:
+    type: Gaussian
+    mean: 0.0
+    sigma: 0.1
+    width_modifier:
+      type: Absolute
+      value: 0.05
+    limits:
+      lower: -inf
+      upper: inf
+  ell_comps_0:
+    type: TruncatedGaussian
+    mean: 0.0
+    sigma: 0.3
+    lower_limit: -1.0
+    upper_limit: 1.0
+    width_modifier:
+      type: Absolute
+      value: 0.2
+    limits:
+      lower: -1.0
+      upper: 1.0
+  ell_comps_1:
+    type: TruncatedGaussian
+    mean: 0.0
+    sigma: 0.3
+    lower_limit: -1.0
+    upper_limit: 1.0
+    width_modifier:
+      type: Absolute
+      value: 0.2
+    limits:
+      lower: -1.0
+      upper: 1.0
+  kappa_s:
+    type: Uniform
+    lower_limit: 0.0
+    upper_limit: 1.0
+    width_modifier:
+      type: Relative
+      value: 0.2
+    limits:
+      lower: 0.0
+      upper: inf
+  scale_radius:
+    type: Uniform
+    lower_limit: 0.0
+    upper_limit: 30.0
+    width_modifier:
+      type: Relative
+      value: 0.2
+    limits:
+      lower: 0.0
+      upper: inf
+  core_radius:
+    type: Uniform
+    lower_limit: 0.0
+    upper_limit: 15.0
+    width_modifier:
+      type: Relative
+      value: 0.2
+    limits:
+      lower: 0.0
+      upper: inf
 cNFWSph:
   centre_0:
     type: Gaussian

autogalaxy/config/priors/mass/dark/cnfw_mcr.yaml

@@ -1,3 +1,88 @@
+cNFWMCRLudlow:
+  centre_0:
+    type: Gaussian
+    mean: 0.0
+    sigma: 0.1
+    width_modifier:
+      type: Absolute
+      value: 0.05
+    limits:
+      lower: -inf
+      upper: inf
+  centre_1:
+    type: Gaussian
+    mean: 0.0
+    sigma: 0.1
+    width_modifier:
+      type: Absolute
+      value: 0.05
+    limits:
+      lower: -inf
+      upper: inf
+  ell_comps_0:
+    type: TruncatedGaussian
+    mean: 0.0
+    sigma: 0.3
+    lower_limit: -1.0
+    upper_limit: 1.0
+    width_modifier:
+      type: Absolute
+      value: 0.2
+    limits:
+      lower: -1.0
+      upper: 1.0
+  ell_comps_1:
+    type: TruncatedGaussian
+    mean: 0.0
+    sigma: 0.3
+    lower_limit: -1.0
+    upper_limit: 1.0
+    width_modifier:
+      type: Absolute
+      value: 0.2
+    limits:
+      lower: -1.0
+      upper: 1.0
+  mass_at_200:
+    type: LogUniform
+    lower_limit: 100000000.0
+    upper_limit: 1000000000000000.0
+    width_modifier:
+      type: Relative
+      value: 0.5
+    limits:
+      lower: 0.0
+      upper: inf
+  f_c:
+    type: Uniform
+    lower_limit: 0.0001
+    upper_limit: 0.5
+    width_modifier:
+      type: Relative
+      value: 0.2
+    limits:
+      lower: 0.0001
+      upper: inf
+  redshift_object:
+    type: Uniform
+    lower_limit: 0.0
+    upper_limit: 1.0
+    width_modifier:
+      type: Relative
+      value: 0.5
+    limits:
+      lower: 0.0
+      upper: inf
+  redshift_source:
+    type: Uniform
+    lower_limit: 0.0
+    upper_limit: 1.0
+    width_modifier:
+      type: Relative
+      value: 0.5
+    limits:
+      lower: 0.0
+      upper: inf
 cNFWMCRLudlowSph:
   centre_0:
     type: Gaussian

autogalaxy/config/priors/mass/dark/cnfw_mcr_scatter.yaml

@@ -1,3 +1,98 @@
+cNFWMCRScatterLudlow:
+  centre_0:
+    type: Gaussian
+    mean: 0.0
+    sigma: 0.1
+    width_modifier:
+      type: Absolute
+      value: 0.05
+    limits:
+      lower: -inf
+      upper: inf
+  centre_1:
+    type: Gaussian
+    mean: 0.0
+    sigma: 0.1
+    width_modifier:
+      type: Absolute
+      value: 0.05
+    limits:
+      lower: -inf
+      upper: inf
+  ell_comps_0:
+    type: TruncatedGaussian
+    mean: 0.0
+    sigma: 0.3
+    lower_limit: -1.0
+    upper_limit: 1.0
+    width_modifier:
+      type: Absolute
+      value: 0.2
+    limits:
+      lower: -1.0
+      upper: 1.0
+  ell_comps_1:
+    type: TruncatedGaussian
+    mean: 0.0
+    sigma: 0.3
+    lower_limit: -1.0
+    upper_limit: 1.0
+    width_modifier:
+      type: Absolute
+      value: 0.2
+    limits:
+      lower: -1.0
+      upper: 1.0
+  mass_at_200:
+    type: LogUniform
+    lower_limit: 100000000.0
+    upper_limit: 1000000000000000.0
+    width_modifier:
+      type: Relative
+      value: 0.5
+    limits:
+      lower: 0.0
+      upper: inf
+  f_c:
+    type: Uniform
+    lower_limit: 0.0001
+    upper_limit: 0.5
+    width_modifier:
+      type: Relative
+      value: 0.2
+    limits:
+      lower: 0.0001
+      upper: inf
+  redshift_object:
+    type: Uniform
+    lower_limit: 0.0
+    upper_limit: 1.0
+    width_modifier:
+      type: Relative
+      value: 0.5
+    limits:
+      lower: 0.0
+      upper: inf
+  redshift_source:
+    type: Uniform
+    lower_limit: 0.0
+    upper_limit: 1.0
+    width_modifier:
+      type: Relative
+      value: 0.5
+    limits:
+      lower: 0.0
+      upper: inf
+  scatter_sigma:
+    type: Gaussian
+    mean: 0.0
+    sigma: 3.0
+    width_modifier:
+      type: Absolute
+      value: 1.0
+    limits:
+      lower: -inf
+      upper: inf
 cNFWMCRScatterLudlowSph:
   centre_0:
     type: Gaussian

autogalaxy/cosmology/model.py

@@ -62,6 +62,27 @@ def kpc_per_arcsec_from(self, redshift: float, xp=np) -> float:
         """
         return self.kpc_proper_per_arcsec(z=redshift, xp=xp)
 
+    def luminosity_distance(self, z: float, xp=np) -> float:
+        """
+        Luminosity distance to redshift z in Mpc.
+
+        For a flat universe:
+
+            D_L(z) = (1 + z)^2 * D_A(0, z)
+
+        where D_A(0, z) is the angular diameter distance from Earth.
+
+        Returns Mpc, matching the convention of astropy.cosmology.FlatLambdaCDM.luminosity_distance(z).value.
+
+        Parameters
+        ----------
+        z
+            Redshift at which the luminosity distance is calculated.
+        """
+        D_A_kpc = self.angular_diameter_distance_to_earth_in_kpc_from(redshift=z, xp=xp)
+        D_A_Mpc = D_A_kpc / xp.asarray(1.0e3)
+        return (xp.asarray(1.0) + xp.asarray(z)) ** 2 * D_A_Mpc
+
     def angular_diameter_distance_to_earth_in_kpc_from(
         self, redshift: float, xp=np
     ) -> float: