Merge branch 'main' into feature/ellipse_utils

Author: samlange04

CLAUDE.md

@@ -78,9 +78,75 @@ 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
 
-JAX is integrated via the `xp` parameter pattern throughout the codebase. Fit classes accept `xp=np` (NumPy, default) or `xp=jnp` (JAX). The `AbstractFitInversion.use_jax` property tracks which backend is active. The `AnalysisImaging.__init__` has `use_jax: bool = True`. The conftest.py forces JAX backend initialization before tests run.
+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.
+
+The `xp` parameter pattern is the single point of control:
+- `xp=np` (default throughout) — pure NumPy path, no JAX dependency at runtime
+- `xp=jnp` — JAX path, imports `jax` / `jax.numpy` locally inside the function
+
+This means:
+- **Unit tests** (`test_autogalaxy/`) always run on the NumPy path. No test should import JAX or pass `xp=jnp` unless it is explicitly testing the JAX path.
+- **Integration tests** (in `autogalaxy_workspace_test/`) are where the JAX path is exercised, typically wrapped in `jax.jit` to test both correctness and compilation.
+- `conftest.py` forces JAX backend initialisation before the test suite runs, but this only ensures JAX is available — it does not switch the default backend.
+
+`AbstractFitInversion.use_jax` tracks whether a fit was constructed with JAX. `AnalysisImaging` has `use_jax: bool = True` to opt into the JAX path for model-fitting.
+
+When adding a new function that should support JAX:
+1. Default the parameter to `xp=np`
+2. Guard any JAX imports with `if xp is not np:` and import `jax` / `jax.numpy` locally inside that branch
+3. Add the NumPy implementation as the default path (finite-difference, `np.*` calls, etc.)
+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
 
@@ -101,6 +167,10 @@ Default priors, visualization settings, and general config live in `autogalaxy/c
 
 Both are mixin classes inherited by `LightProfile`, `MassProfile`, `Galaxy`, and `Galaxies`.
 
+### Workspace Script Style
+
+Scripts in `autogalaxy_workspace` and `autogalaxy_workspace_test` use `"""..."""` docstring blocks as prose commentary throughout — **not** `#` comments. Every script opens with a module-level docstring (title + underline + description), and each logical section of code is preceded by a `"""..."""` block with a `__Section Name__` header explaining what follows. See any script in `autogalaxy_workspace/scripts/` for examples of this style.
+
 ### Workspace (Examples & Notebooks)
 
 The `autogalaxy_workspace` at `/mnt/c/Users/Jammy/Code/PyAutoJAX/autogalaxy_workspace` contains runnable examples and tutorials. Key locations:
@@ -125,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/__init__.py

@@ -69,7 +69,7 @@
 from .operate.image import OperateImage
 from .operate.image import OperateImageList
 from .operate.image import OperateImageGalaxies
-from .operate.deflections import OperateDeflections
+from .operate.lens_calc import LensCalc
 from .gui.scribbler import Scribbler
 from .imaging.fit_imaging import FitImaging
 from .imaging.model.analysis import AnalysisImaging

autogalaxy/analysis/model_util.py

@@ -11,6 +11,7 @@ def mge_model_from(
     centre_prior_is_uniform: bool = True,
     centre: Tuple[float, float] = (0.0, 0.0),
     centre_fixed: Optional[Tuple[float, float]] = None,
+    centre_sigma: float = 0.3,
     use_spherical: bool = False,
 ) -> af.Collection:
     """
@@ -89,8 +90,8 @@ def mge_model_from(
             lower_limit=centre[1] - 0.1, upper_limit=centre[1] + 0.1
         )
     else:
-        centre_0 = af.GaussianPrior(mean=centre[0], sigma=0.3)
-        centre_1 = af.GaussianPrior(mean=centre[1], sigma=0.3)
+        centre_0 = af.GaussianPrior(mean=centre[0], sigma=centre_sigma)
+        centre_1 = af.GaussianPrior(mean=centre[1], sigma=centre_sigma)
 
     if use_spherical:
         model_cls = GaussianSph
@@ -129,6 +130,106 @@ def mge_model_from(
     )
 
 
+def mge_point_model_from(
+    pixel_scales: float,
+    total_gaussians: int = 10,
+    centre: Tuple[float, float] = (0.0, 0.0),
+) -> af.Model:
+    """
+    Construct a Multi-Gaussian Expansion (MGE) model for a compact or unresolved
+    point-like component (e.g. a nuclear starburst, AGN, or unresolved bulge).
+
+    The model is composed of ``total_gaussians`` linear Gaussians whose sigma values
+    are logarithmically spaced between 0.01 arcseconds and twice the pixel scale.
+    All Gaussians share the same centre and ellipticity components, keeping the
+    parameter count low while capturing a realistic PSF-convolved point source.
+
+    Parameters
+    ----------
+    pixel_scales
+        The pixel scale of the image in arcseconds per pixel.  The maximum Gaussian
+        width is set to ``2 * pixel_scales`` so that the model is compact relative to
+        the resolution of the data.
+    total_gaussians
+        Number of Gaussian components in the basis.
+    centre
+        (y, x) centre of the point source in arc-seconds.  A ±0.1 arcsecond uniform
+        prior is placed on each coordinate.
+
+    Returns
+    -------
+    af.Model
+        An ``autofit.Model`` wrapping a ``Basis`` of linear Gaussians.
+    """
+
+    from autogalaxy.profiles.light.linear import Gaussian
+    from autogalaxy.profiles.basis import Basis
+
+    if total_gaussians < 1:
+        raise ValueError(
+            f"mge_point_model_from requires total_gaussians >= 1, got {total_gaussians}."
+        )
+
+    if pixel_scales <= 0:
+        raise ValueError(
+            f"mge_point_model_from requires pixel_scales > 0, got {pixel_scales}."
+        )
+
+    # Sigma values are logarithmically spaced between 0.01 arcsec (10**-2)
+    # and twice the pixel scale, with a floor to avoid taking log10 of
+    # very small or non-positive values.
+    min_log10_sigma = -2.0  # corresponds to 0.01 arcsec
+    max_sigma = max(2.0 * pixel_scales, 10**min_log10_sigma)
+    max_log10_sigma = np.log10(max_sigma)
+
+    log10_sigma_list = np.linspace(min_log10_sigma, max_log10_sigma, total_gaussians)
+    centre_0 = af.UniformPrior(lower_limit=centre[0] - 0.1, upper_limit=centre[0] + 0.1)
+    centre_1 = af.UniformPrior(lower_limit=centre[1] - 0.1, upper_limit=centre[1] + 0.1)
+
+    gaussian_list = af.Collection(af.Model(Gaussian) for _ in range(total_gaussians))
+
+    for i, gaussian in enumerate(gaussian_list):
+        gaussian.centre.centre_0 = centre_0
+        gaussian.centre.centre_1 = centre_1
+        gaussian.ell_comps = gaussian_list[0].ell_comps
+        gaussian.sigma = 10 ** log10_sigma_list[i]
+
+    return af.Model(Basis, profile_list=gaussian_list)
+
+
+def hilbert_pixels_from_pixel_scale(pixel_scale: float) -> int:
+    """
+    Return the number of Hilbert-curve pixels appropriate for a given image pixel scale.
+
+    The Hilbert pixel count controls the resolution of the Hilbert-curve ordering used
+    in adaptive source-plane pixelizations. Finer pixel scales resolve smaller angular
+    features and therefore benefit from a higher Hilbert resolution.
+
+    Parameters
+    ----------
+    pixel_scale
+        The pixel scale of the image in arcseconds per pixel.
+
+    Returns
+    -------
+    int
+        The recommended number of Hilbert pixels.
+    """
+    if not np.isfinite(pixel_scale) or pixel_scale <= 0:
+        raise ValueError(
+            f"hilbert_pixels_from_pixel_scale requires pixel_scale to be finite and > 0, got {pixel_scale}."
+        )
+
+    if pixel_scale > 0.06:
+        return 1000
+    elif pixel_scale > 0.04:
+        return 1250
+    elif pixel_scale >= 0.03:
+        return 1500
+    else:
+        return 1750
+
+
 def simulator_start_here_model_from():
 
     from autogalaxy.profiles.light.snr import Sersic

autogalaxy/config/general.yaml

@@ -1,5 +1,5 @@
 fits:
-  flip_for_ds9: true
+  flip_for_ds9: false
 psf:
   use_fft_default: true              # If True, PSFs are convolved using FFTs by default, which is faster and uses less memory in all cases except for very small PSFs, False uses direct convolution.
 updates:

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