diff --git a/AGENTS.md b/AGENTS.md index f94c255ef..4fb200db7 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -1,75 +1,124 @@ # PyAutoLens — Agent Instructions -**PyAutoLens** is a Python library for strong gravitational lens modeling, built on PyAutoGalaxy. It adds multi-plane ray-tracing via the `Tracer` object and lensing-specific fit/analysis classes. +Canonical, agent-agnostic instructions for this repo. `CLAUDE.md` imports this +file; any tool that does not process `@`-imports should read this directly. -## Setup +## What this repo is -```bash -pip install -e ".[dev]" -``` +**PyAutoLens** (package `autolens`) is the strong gravitational-lensing layer +built on PyAutoGalaxy. It adds multi-plane ray-tracing via the `Tracer`, and +lensing-specific `Fit*` / `Analysis*` classes for imaging, interferometer, and +point-source datasets. + +Dependency direction: autolens sits at the top of the stack and may import all +four layers below it — **autogalaxy**, **autoarray**, **autofit**, and +**autoconf**. Nothing in the ecosystem imports autolens. -## Running Tests +## Related repos + +- **Source siblings (all upstream):** PyAutoConf, PyAutoArray, PyAutoFit, + PyAutoGalaxy. +- **autolens_workspace** — runnable tutorials/examples (`../autolens_workspace`). +- **autolens_workspace_test** — integration + JAX/likelihood parity scripts. +- **autolens_profiling** — performance/profiling harness (`../autolens_profiling`). +- **HowToLens** — the lecture-style tutorial series (`../HowToLens`). +- **docs/** — Sphinx source; published to ReadTheDocs. +- **Science context:** the strong-lensing knowledge wiki at + `autolens_assistant/wiki/literature/` (concepts, entities, sources) — mass + models, source reconstruction, degeneracies, substructure, surveys. + +## Quick commands ```bash -python -m pytest test_autolens/ -python -m pytest test_autolens/lens/test_tracer.py -python -m pytest test_autolens/imaging/test_fit_imaging.py -s +pip install -e ".[dev]" # install with dev/test extras +python -m pytest test_autolens/ # full test suite +python -m pytest test_autolens/lens/test_tracer.py # one focused test (add -s for output) +black autolens/ # formatter (advisory — not gated) ``` -### Sandboxed / Codex runs +In a sandboxed / restricted environment, point numba and matplotlib at +writable caches: ```bash NUMBA_CACHE_DIR=/tmp/numba_cache MPLCONFIGDIR=/tmp/matplotlib python -m pytest test_autolens/ ``` -## Key Architecture +## CI / definition of green + +PRs must pass `pytest --cov` on the CI matrix (Python 3.12 **and** 3.13). There +is no black/ruff/flake8 gate — formatting is advisory. (`requires-python` in +`pyproject.toml` is `>=3.9`.) -- **Tracer** (`lens/tracer.py`): groups galaxies by redshift plane, performs multi-plane ray-tracing -- **Fit classes**: `FitImaging`, `FitInterferometer`, `FitPointDataset` — extend autogalaxy equivalents with lensing -- **Analysis classes**: `AnalysisImaging`, `AnalysisInterferometer`, `AnalysisPoint` -- **Namespace**: `al.mp.*` (mass), `al.lp.*` (light), `al.Galaxy`, `al.Tracer` +## Configuration & defaults -## Dependencies +autoconf supplies the packaged defaults under `autolens/config/`. Workspaces +override them via their own `config/` directory; the test suite pushes a local +config dir via `conf.instance.push(...)` in `test_autolens/conftest.py`. When a +change adds a new config key, mirror it into the packaged defaults so +downstream workspaces inherit it. -- `autogalaxy` — galaxy morphology, profiles, single-plane fitting -- `autoarray` — data structures, grids, masks, inversions -- `autofit` — non-linear search and model-fitting framework +## JAX & `xp` -## Key Rules +NumPy is the default everywhere; JAX is opt-in and never imported at module +level. `xp=np` (default) selects NumPy; `xp=jnp` selects JAX (imported locally). +Thread `xp` through **every** nested call — a missed site silently defaults to +`xp=np` and fails when a tracer hits an `np.*` op. Two patterns cross the +`jax.jit` boundary: the `if xp is np:` **guard** for raw `jax.Array` returns +(the `LensCalc` hessian methods), and **pytree registration** for functions +returning real wrappers/structured objects — `FitImaging`, `Tracer`, and +`DatasetModel` register via `register_instance_pytree`, so +`jax.jit(analysis.fit_from)(instance)` returns a real `FitImaging` with +`jax.Array` leaves. -- The `xp` parameter controls NumPy vs JAX: `xp=np` (default) or `xp=jnp` -- Functions inside `jax.jit` must guard autoarray wrapping with `if xp is np:` -- Decorated functions return **raw arrays** — the decorator wraps them -- All files must use Unix line endings (LF) -- Format with `black autolens/` +**Unit tests are NumPy-only.** A JAX/`xp` change is validated only by the +parity scripts in `autolens_workspace_test` (`jax.jit` round-trip + +`fitness._vmap` batch eval) — never by `test_autolens/`. + +Full detail lives in PyAutoArray: +**[`PyAutoArray/docs/agents/jax_and_decorators.md`](../PyAutoArray/docs/agents/jax_and_decorators.md)**. + +## Public API + +The public surface is defined authoritatively in `autolens/__init__.py` — read +it rather than trusting a hand-maintained namespace table. Canonical import: + +```python +import autolens as al +``` -## Working on Issues +Profiles re-export from autogalaxy (`al.mp.*`, `al.lp.*`) alongside `al.Galaxy`, +`al.Tracer`, `al.FitImaging`/`al.AnalysisImaging`, and the point-source classes. + +## Key rules / footguns + +- Import direction: autolens may use all four upstream packages; nothing + imports autolens. +- Grid-decorated profile methods return a **raw array** (the decorator wraps + it); write `aa.decorators.*` and read coordinates via `grid.array[:, 0]`. +- All files use Unix line endings (LF, `\n`) — never `\r\n`. + +## Working on issues 1. Read the issue description and any linked plan. -2. Identify affected files and write your changes. -3. Run the full test suite: `python -m pytest test_autolens/` -4. Ensure all tests pass before opening a PR. -5. If changing public API, note the change in your PR description — downstream workspaces may need updates. -## Never rewrite history - -NEVER perform these operations on any repo with a remote: - -- `git init` in a directory already tracked by git -- `rm -rf .git && git init` -- Commit with subject "Initial commit", "Fresh start", "Start fresh", "Reset - for AI workflow", or any equivalent message on a branch with a remote -- `git push --force` to `main` (or any branch tracked as `origin/HEAD`) -- `git filter-repo` / `git filter-branch` on shared branches -- `git rebase -i` rewriting commits already pushed to a shared branch - -If the working tree needs a clean state, the **only** correct sequence is: - - git fetch origin - git reset --hard origin/main - git clean -fd - -This applies equally to humans, local Claude Code, cloud Claude agents, Codex, -and any other agent. The "Initial commit — fresh start for AI workflow" pattern -that appeared independently on origin and local for three workspace repos is -exactly what this rule prevents — it costs ~40 commits of redundant local work -every time it happens. +2. Identify affected files and make the change. +3. Run the full suite: `python -m pytest test_autolens/`. +4. If you changed public API, say so explicitly — the workspaces and + downstream pipelines may need updates. +5. Ensure all tests pass before opening a PR. + +## Deep dives + +- [`PyAutoArray/docs/agents/jax_and_decorators.md`](../PyAutoArray/docs/agents/jax_and_decorators.md) + — decorator system, `xp` backend pattern, and the `jax.jit` boundary. + +## Clean state + +Never rewrite history on a repo with a remote (no `git init` over a tracked +tree, no force-push to `main`, no rebasing pushed shared branches). To reset a +dirty tree the only correct sequence is: + +```bash +git fetch origin +git reset --hard origin/main +git clean -fd +``` diff --git a/CLAUDE.md b/CLAUDE.md index 758ea6bbc..211b06330 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -1,180 +1,5 @@ -# CLAUDE.md - -This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository. - -## Commands - -### Install -```bash -pip install -e ".[dev]" -``` - -### Run Tests -```bash -# All tests -python -m pytest test_autolens/ - -# Single test file -python -m pytest test_autolens/lens/test_tracer.py - -# With output -python -m pytest test_autolens/imaging/test_fit_imaging.py -s -``` - -### Codex / sandboxed runs - -When running Python from Codex or any restricted environment, set writable cache directories so `numba` and `matplotlib` do not fail on unwritable home or source-tree paths: - -```bash -NUMBA_CACHE_DIR=/tmp/numba_cache MPLCONFIGDIR=/tmp/matplotlib python -m pytest test_autolens/ -``` - -This workspace is often imported from `/mnt/c/...` and Codex may not be able to write to module `__pycache__` directories or `/home/jammy/.cache`, which can cause import-time `numba` caching failures without this override. - -### Formatting -```bash -black autolens/ -``` - -### Plot Output Mode - -Set `PYAUTO_OUTPUT_MODE=1` to capture every figure produced by a script into numbered PNG files in `./output_mode//`. This is useful for visually inspecting all plots from an integration test without needing a display. - -```bash -PYAUTO_OUTPUT_MODE=1 python scripts/my_script.py -# -> ./output_mode/my_script/0_fit.png, 1_tracer.png, ... -``` - -When this env var is set, all `save_figure`, `subplot_save`, and `_save_subplot` calls are intercepted — the normal output path is bypassed and figures are written sequentially to the output_mode directory instead. - -## Architecture - -**PyAutoLens** is the gravitational lensing layer built on top of PyAutoGalaxy. It adds multi-plane ray-tracing, the `Tracer` object, and lensing-specific fit classes. It depends on: -- **`autogalaxy`** — galaxy morphology, mass/light profiles, single-plane fitting -- **`autoarray`** — low-level data structures (grids, masks, arrays, datasets, inversions) -- **`autofit`** — non-linear search / model-fitting framework - -### Core Class Hierarchy - -``` -Tracer (lens/tracer.py) - └── List[List[Galaxy]] — galaxies grouped by redshift plane - ├── ray-traces from source to lens to observer - ├── delegates to autogalaxy Galaxy/Galaxies for per-plane operations - └── returns lensed images, deflection maps, convergence, magnification -``` - -### Dataset Types and Fit Classes - -| Dataset | Fit class | Analysis class | -|---|---|---| -| `aa.Imaging` | `FitImaging` | `AnalysisImaging` | -| `aa.Interferometer` | `FitInterferometer` | `AnalysisInterferometer` | -| Point source | `FitPointDataset` | `AnalysisPoint` | - -All inherit from the corresponding `autogalaxy` base classes (`ag.FitImaging`, etc.) and extend them with multi-plane lensing via the `Tracer`. - -### Key Directories - -``` -autolens/ - lens/ Tracer, ray-tracing, multi-plane deflection logic - imaging/ FitImaging, AnalysisImaging - interferometer/ FitInterferometer, AnalysisInterferometer - point/ Point-source datasets, fits, and analysis - analysis/ Shared analysis base classes, adapt images - aggregator/ Scraping results from autofit output directories - plot/ Visualisation for all data types -``` - -## Decorator System (from autoarray) - -PyAutoLens inherits the same decorator conventions as PyAutoGalaxy. Mass and light profile methods that take a grid and return an array/grid/vector are decorated with: - -| Decorator | `Grid2D` → | `Grid2DIrregular` → | -|---|---|---| -| `@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 innermost) transforms the grid to the profile's reference frame. Standard stacking: - -```python -@aa.grid_dec.to_array -@aa.grid_dec.transform -def convergence_2d_from(self, grid, xp=np, **kwargs): - y = grid.array[:, 0] # .array extracts raw numpy/jax array - x = grid.array[:, 1] - return ... # raw array — decorator wraps it -``` - -The function body must return a **raw array**. Use `grid.array[:, 0]` (not `grid[:, 0]`) to access coordinates safely for both numpy and jax backends. - -See PyAutoArray's `CLAUDE.md` for full decorator internals. - -## JAX Support - -The `xp` parameter pattern controls the backend: -- `xp=np` (default) — pure NumPy, no JAX dependency -- `xp=jnp` — JAX path; `jax`/`jax.numpy` imported locally inside the function only - -### JAX and the `jax.jit` boundary - -Two patterns coexist for crossing the JIT boundary: - -**Pattern 1: `if xp is np:` guard (raw `jax.Array` return).** Functions intended to be called directly inside `jax.jit` as the outermost op — where no wrapper is needed on the JAX path — guard their autoarray wrapping: - -```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 -``` - -All `LensCalc` hessian-derived methods use this pattern. Intermediate helpers (e.g. `deflections_yx_2d_from`) don't need the guard — they're consumed by downstream Python before the JIT boundary. - -**Pattern 2: pytree-registered wrapper return.** Functions that must return a real autoarray wrapper (or a structured object built from them) opt in to JAX pytree registration. `AbstractNDArray` auto-registers its subclass with `jax.tree_util` the first time an instance is built with `xp=jnp` (via `autoarray.abstract_ndarray._register_as_pytree`). Higher-level types (`FitImaging`, `Tracer`, `DatasetModel`) use `autoarray.abstract_ndarray.register_instance_pytree(cls, no_flatten=...)`, which flattens `__dict__` and carries `no_flatten` names through `aux_data` for per-analysis constants (dataset, settings, cosmology). `AnalysisImaging._register_fit_imaging_pytrees` wires these up when `use_jax=True`, so `jax.jit(analysis.fit_from)(instance)` returns a real `FitImaging` with `jax.Array` leaves. - -### `LensCalc` (autogalaxy) - -The hessian-derived lensing quantities (`convergence_2d_via_hessian_from`, `shear_yx_2d_via_hessian_from`, `magnification_2d_via_hessian_from`, `magnification_2d_from`, `tangential_eigen_value_from`, `radial_eigen_value_from`) all implement the `if xp is np:` guard in `autogalaxy/operate/lens_calc.py` and return raw `jax.Array` on the JAX path, making them safe to call inside `jax.jit`. - -## Namespace Conventions - -When importing `autolens as al`: -- `al.mp.*` — mass profiles (re-exported from autogalaxy) -- `al.lp.*` — light profiles (re-exported from autogalaxy) -- `al.Galaxy`, `al.Galaxies` -- `al.Tracer` -- `al.FitImaging`, `al.AnalysisImaging`, `al.SimulatorImaging` -- `al.FitInterferometer`, `al.AnalysisInterferometer` -- `al.FitPointDataset`, `al.AnalysisPoint` - -## Line Endings — Always Unix (LF) - -All files **must use Unix line endings (LF, `\n`)**. Never write `\r\n` line endings. -## Never rewrite history - -NEVER perform these operations on any repo with a remote: - -- `git init` in a directory already tracked by git -- `rm -rf .git && git init` -- Commit with subject "Initial commit", "Fresh start", "Start fresh", "Reset - for AI workflow", or any equivalent message on a branch with a remote -- `git push --force` to `main` (or any branch tracked as `origin/HEAD`) -- `git filter-repo` / `git filter-branch` on shared branches -- `git rebase -i` rewriting commits already pushed to a shared branch - -If the working tree needs a clean state, the **only** correct sequence is: - - git fetch origin - git reset --hard origin/main - git clean -fd - -This applies equally to humans, local Claude Code, cloud Claude agents, Codex, -and any other agent. The "Initial commit — fresh start for AI workflow" pattern -that appeared independently on origin and local for three workspace repos is -exactly what this rule prevents — it costs ~40 commits of redundant local work -every time it happens. +# PyAutoLens — agent instructions +The canonical, agent-agnostic instructions live in `AGENTS.md`. Claude Code loads them +via the import below; if your tool does not process `@`-imports, open `AGENTS.md` in +this directory and read it directly. +@AGENTS.md