Skip to content

docs: consolidate agent instructions into canonical AGENTS.md #346

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Jammy2211 merged 1 commit into from
Jun 22, 2026
Merged

docs: consolidate agent instructions into canonical AGENTS.md #346

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
153 changes: 103 additions & 50 deletions AGENTS.md
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,70 +1,123 @@
# PyAutoArray — Agent Instructions

**PyAutoArray** is the low-level data structures and numerical utilities package for the PyAuto ecosystem. It provides grids, masks, arrays, datasets, inversions, and the decorator system used throughout PyAutoGalaxy and PyAutoLens.
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]"
```
**PyAutoArray** (package `autoarray`) is the low-level data-structure and
numerical-utility layer: masks, arrays, (y,x) grids, imaging/interferometer
datasets, inversions/pixelizations, convolution/over-sampling operators, and
the grid decorators used throughout PyAutoGalaxy and PyAutoLens.

Dependency direction: autoarray depends on **autoconf** only. It does **not**
import `autofit`, `autogalaxy`, or `autolens` — never add such an import.
Shared utilities (e.g. `test_mode`, `jax_wrapper`) belong in autoconf.

## Related repos

## Running Tests
- **Source siblings:** PyAutoConf (upstream). PyAutoGalaxy / PyAutoLens build
directly on autoarray.
- No `_workspace`, `_workspace_test`, or HowTo of its own. The JAX/`xp` path is
exercised by the parity scripts in **autogalaxy_workspace_test** and
**autolens_workspace_test**.
- **docs/** — Sphinx source; published to ReadTheDocs.

## Quick commands

```bash
python -m pytest test_autoarray/
python -m pytest test_autoarray/structures/test_arrays.py
python -m pytest test_autoarray/structures/test_arrays.py -s
pip install -e ".[dev]" # install with dev/test extras
python -m pytest test_autoarray/ # full test suite
python -m pytest test_autoarray/structures/test_arrays.py # one focused test (add -s for output)
black autoarray/ # 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_autoarray/
```

## 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`.)

## Configuration & defaults

- **Data structures**: `Array2D`, `Grid2D`, `Grid2DIrregular`, `VectorYX2D` — all inherit from `AbstractNDArray`
- **Decorator system** (`structures/decorators/`): `@to_array`, `@to_grid`, `@to_vector_yx`, `@transform` — ensures output type matches input grid type
- **Datasets**: `Imaging`, `Interferometer` — containers for observational data
- **Inversions** (`inversion/`): sparse linear algebra for source reconstruction via pixelizations
- **Operators**: `Convolver` (PSF convolution), over-sampling utilities
autoconf supplies the packaged defaults under `autoarray/config/`. Workspaces
override them via their own `config/` directory; the test suite pushes a local
config dir via `conf.instance.push(...)` in `test_autoarray/conftest.py`. When
a change adds a new config key, mirror it into the packaged defaults so
downstream workspaces inherit it.

## Key Rules
## JAX & `xp`

- The `xp` parameter pattern controls NumPy vs JAX: `xp=np` (default) or `xp=jnp`
- Autoarray types are **not** JAX pytrees — they cannot be returned from `jax.jit` functions
- Decorated functions must return **raw arrays**, not autoarray wrappers
- All files must use Unix line endings (LF)
- Format with `black autoarray/`
NumPy is the default everywhere; JAX is opt-in and never imported at module
level. The `xp` parameter is the single point of control:

## Working on Issues
- `xp=np` (default) — pure NumPy path.
- `xp=jnp` — JAX path; `jax` / `jax.numpy` imported locally inside the function.

Thread `xp` through **every** nested call (`self.X()`, helpers, properties) — 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
functions that return a raw `jax.Array`, and **pytree registration**
(`abstract_ndarray._register_as_pytree` / `register_instance_pytree`) for
functions that return a real wrapper or structured object.

**Unit tests are NumPy-only.** A JAX/`xp` change is validated only by the
parity scripts in `autogalaxy_workspace_test` / `autolens_workspace_test`
(`jax.jit` round-trip + `fitness._vmap` batch eval), not by `test_autoarray/`.

Full detail (decorator internals, `xp`-threading hazards, both JIT patterns):
**[`docs/agents/jax_and_decorators.md`](docs/agents/jax_and_decorators.md)**.

## Public API

The public surface is defined authoritatively in `autoarray/__init__.py` — read
it rather than trusting a hand-maintained list. Canonical import:

```python
import autoarray as aa
```

Core types (`Array2D`, `Grid2D`, `Grid2DIrregular`, `VectorYX2D`, …) inherit
from `AbstractNDArray`; `.array` returns the raw `numpy.ndarray` / `jax.Array`.

## Key rules / footguns

- Import direction: autoconf only — never `autofit` / `autogalaxy` / `autolens`.
- Grid-consuming functions decorated with `@aa.decorators.to_array` / `to_grid`
/ `to_vector_yx` must return a **raw array** — the decorator wraps it. (Write
`aa.decorators.*`; `aa.grid_dec` is a deprecated alias.)
- Access grid coordinates via `grid.array[:, 0]`, not `grid[:, 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_autoarray/`
4. Ensure all tests pass before opening a PR.
5. If changing public API, note the change in your PR description — downstream packages (PyAutoGalaxy, PyAutoLens) and 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_autoarray/`.
4. If you changed public API, say so explicitly — downstream packages
(PyAutoGalaxy, PyAutoLens) and the workspaces may need updates.
5. Ensure all tests pass before opening a PR.

## Deep dives

- [`docs/agents/jax_and_decorators.md`](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
```
209 changes: 5 additions & 204 deletions CLAUDE.md
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,204 +1,5 @@
# CLAUDE.md

This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository.

## Dependency Graph

PyAutoArray depends on **autoconf** (shared configuration and utilities).
PyAutoArray does **NOT** depend on PyAutoFit, PyAutoGalaxy, or PyAutoLens.
Never import from `autofit`, `autogalaxy`, or `autolens` in this repo.
Shared utilities (e.g. `test_mode`, `jax_wrapper`) belong in autoconf.

## Commands

### Install
```bash
pip install -e ".[dev]"
```

### Run Tests
```bash
# All tests
python -m pytest test_autoarray/

# Single test file
python -m pytest test_autoarray/structures/test_arrays.py

# With output
python -m pytest test_autoarray/structures/test_arrays.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_autoarray/
```

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 autoarray/
```

### Plot Output Mode

Set `PYAUTO_OUTPUT_MODE=1` to capture every figure produced by a script into numbered PNG files in `./output_mode/<script_name>/`. 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

**PyAutoArray** is the low-level data structures and numerical utilities package for the PyAuto ecosystem. It provides:
- **Grid and array structures** — uniform and irregular 2D grids, arrays, vector fields
- **Masks** — 1D and 2D masks that define which pixels are active
- **Datasets** — imaging and interferometer dataset containers
- **Inversions / pixelizations** — sparse linear algebra for source reconstruction
- **Decorators** — input/output homogenisation for grid-consuming functions

## Core Data Structures

All data structures inherit from `AbstractNDArray` (`abstract_ndarray.py`). Key subclasses:

| Class | Description |
|---|---|
| `Array2D` | Uniform 2D array tied to a `Mask2D` |
| `ArrayIrregular` | Unmasked 1D collection of values |
| `Grid2D` | Uniform (y,x) coordinate grid tied to a `Mask2D` |
| `Grid2DIrregular` | Irregular (y,x) coordinate collection |
| `VectorYX2D` | Uniform 2D vector field |
| `VectorYX2DIrregular` | Irregular vector field |

`AbstractNDArray` provides arithmetic operators (`__add__`, `__sub__`, `__rsub__`, etc.), all decorated with `@to_new_array` and `@unwrap_array` so that operations between autoarray objects and raw scalars/arrays work naturally and return a new autoarray of the same type.

The `.array` property returns the raw underlying `numpy.ndarray` or `jax.Array`:
```python
arr = aa.ArrayIrregular(values=[1.0, 2.0])
arr.array # raw numpy array
arr._array # same, internal attribute
```

The constructor unwraps nested autoarray objects automatically:
```python
# while isinstance(array, AbstractNDArray): array = array.array
```

## Decorator System

`autoarray/structures/decorators/` contains three output-wrapping decorators used on all grid-consuming functions. They ensure that the **type of the output structure matches the type of the input grid**:

| 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` |

### How the decorators work

All three share `AbstractMaker` (`decorators/abstract.py`). The decorator:
1. Wraps the function in a `wrapper(obj, grid, xp=np, *args, **kwargs)` signature
2. Instantiates the relevant `*Maker` class with the function, object, grid, and `xp`
3. `AbstractMaker.result` checks the grid type and calls the appropriate `via_grid_2d` / `via_grid_2d_irr` method to wrap the raw result

The function body receives the grid as-is and **must return a raw array** (not an autoarray wrapper). The decorator does the wrapping:

```python
@aa.grid_dec.to_array
def convergence_2d_from(self, grid, xp=np, **kwargs):
# grid is Grid2D or Grid2DIrregular — access raw values via grid.array[:,0]
y = grid.array[:, 0]
x = grid.array[:, 1]
return xp.sqrt(y**2 + x**2) # return raw array; decorator wraps it
```

`AbstractMaker` also stores `use_jax = xp is not np` and exposes `_xp` (either `jnp` or `np`), but the wrapping step always runs regardless of `xp`. Autoarray types are **not registered as JAX pytrees**, so they cannot be directly returned from inside a `jax.jit` trace (see JAX section below).

### Accessing grid coordinates inside a decorated function

Inside a decorated function body, access the raw underlying array with `.array`:

```python
# Correct — works for both numpy and jax backends
y = grid.array[:, 0]
x = grid.array[:, 1]

# Also correct for simple slicing (returns raw array via __getitem__)
y = grid[:, 0]
x = grid[:, 1]
```

The `@transform` decorator (also in `decorators/`) shifts and rotates the input grid to the profile's reference frame before passing it to the function. It calls `obj.transformed_to_reference_frame_grid_from(grid, xp)` (decorated with `@to_grid`) and passes the result as the `grid` argument. After transformation the grid is still an autoarray object; `.array` still works.

### Decorator stacking order

Decorators are applied bottom-up (innermost first). The canonical order for mass/light profile methods is:

```python
@aa.grid_dec.to_array # outermost: wraps output
@aa.grid_dec.transform # innermost: transforms grid input
def convergence_2d_from(self, grid, xp=np, **kwargs):
...
```

## JAX Support

The `xp` parameter pattern is the single point of control:
- `xp=np` (default) — pure NumPy path
- `xp=jnp` — JAX path; `jax` / `jax.numpy` are only imported locally

### Why autoarray types cannot be returned from `jax.jit`

`AbstractNDArray` subclasses (`Array2D`, `ArrayIrregular`, `VectorYX2DIrregular`, etc.) are **not registered as JAX pytrees**. The `instance_flatten` / `instance_unflatten` class methods are defined on `AbstractNDArray` but are never passed to `jax.tree_util.register_pytree_node`. As a result:

- Constructing an autoarray wrapper **inside** a JIT trace is fine (Python-level code runs normally during tracing)
- **Returning** an autoarray wrapper as the output of a `jax.jit`-compiled function **fails** with `TypeError: ... is not a valid JAX type`

### The `if xp is np:` guard pattern

Functions that are called directly inside `jax.jit` (i.e., as the outermost call in the lambda) must not return autoarray wrappers on the JAX path. The correct pattern is:

```python
def convergence_2d_via_hessian_from(self, grid, xp=np):
hessian_yy, hessian_xx = ...
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 used in `autogalaxy/operate/lens_calc.py` for all `LensCalc` methods that are called inside `jax.jit`. It does **not** affect decorated helper functions (like `deflections_yx_2d_from`) because those are called as intermediate steps — their autoarray wrappers are consumed by downstream Python code, never returned as JIT outputs.

## 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.
# PyAutoArray — 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
Loading
Loading