diff --git a/skills/al_configure_search.md b/skills/al_configure_search.md index 60b7cf0..1f6b070 100644 --- a/skills/al_configure_search.md +++ b/skills/al_configure_search.md @@ -51,7 +51,11 @@ Source: `PyAutoFit:autofit/non_linear/search/nest/nautilus/`. Knobs to know: - `n_live` — more = more accurate posterior, slower. Start at 200; go to 400+ only if the posterior looks multi-modal or thin. -- `number_of_cores` — set to your CPU core count for parallel likelihood eval. +- `number_of_cores` — parallel likelihood evaluations via Python multiprocessing, **only when + JAX is off**. JAX disables multiprocessing, so a JAX fit gains nothing from it — leave it + unset there (it defaults to 1; passing `1` explicitly just implies a parallelism that isn't + there). Set it to your core count only for non-JAX CPU fits — see + "Branch — CPU acceleration" below, which decides *which* regime a fit belongs to. - `iterations_per_full_update` / `iterations_per_quick_update` — how often the search writes full output (samples, visualisation) vs quick intermediate updates to disk. **Actively choose `iterations_per_quick_update` so the user always has quick access to @@ -75,6 +79,52 @@ Knobs to know: editing every search, write it to the live config once after `conf.instance.push(...)`: `conf.instance["general"]["updates"]["iterations_per_quick_update"] = N`. +## Branch — CPU acceleration (JAX vs sparse operators) + +On CPU the right accelerator depends on the **source model**. This is the single biggest CPU +runtime lever in lens modelling, and getting it backwards costs days, not minutes — a full SLaM +run misconfigured here spends >12 h stuck in its first search. + +| Fit type | Accelerator | `use_jax` | `number_of_cores` | Dataset | +|---|---|---|---|---| +| **Parametric source** (`source_lp`, most non-pixelised fits) | JAX — vectorises the likelihood, parallelises well **on CPU** | `True` | leave unset (JAX disables multiprocessing) | plain | +| **Pixelised source** (`source_pix`, `light`, `mass` — any fit with a `Pixelization`) | Sparse operator formalism (numba) | `False` | your core count | `dataset.apply_sparse_operator_cpu()` | + +Two rules that are easy to get wrong: + +- **JAX is not GPU-only.** It is the correct accelerator for *parametric* fits on CPU too. Do + not reach for `use_jax=False` + many cores just because there is no GPU. +- **The sparse operator does not support JAX.** It is numba-based, so never combine them: a + pixelised CPU fit is `apply_sparse_operator_cpu()` + `use_jax=False` + `number_of_cores=N`. + `apply_sparse_operator_cpu()` precomputes operator matrices once (seconds to minutes) and + every later pixelised fit reuses them, exploiting the sparsity of the pixelisation linear + algebra for a large CPU speed-up. + +```python +# Parametric source on CPU — JAX, no number_of_cores. +analysis = al.AnalysisImaging(dataset=dataset, use_jax=True) +settings_search = af.SettingsSearch(path_prefix=..., unique_tag=..., session=None) + +# Pixelised source on CPU — sparse operators, JAX off, multiprocessing on. +dataset_pix = dataset.apply_sparse_operator_cpu() +analysis = al.AnalysisImaging(dataset=dataset_pix, use_jax=False) +settings_search = af.SettingsSearch( + path_prefix=..., unique_tag=..., session=None, number_of_cores=8 +) +``` + +`number_of_cores` reaches the search through `af.SettingsSearch` (it is one of `search_dict`'s +fixed keys), so pass it there rather than to the search constructor. + +**A SLaM pipeline spans both regimes.** SOURCE LP is parametric (JAX); every stage after it uses +a pixelised source (sparse/CPU). So build **two** `af.SettingsSearch` objects — one without +`number_of_cores` for SOURCE LP, one with N for the pixelised stages — and hand each stage the +right dataset (`dataset` vs `dataset_pix`). On **GPU**, JAX is used throughout and the sparse +operator is not applied at all. + +Source / worked example: +`autolens_workspace:scripts/imaging/features/pixelization/cpu_fast_modeling.py`. + ## Branch — Dynesty Use Dynesty for problems where you specifically want its dynamic sampling features diff --git a/skills/init-slam.md b/skills/init-slam.md index 499d959..6092c17 100644 --- a/skills/init-slam.md +++ b/skills/init-slam.md @@ -65,6 +65,15 @@ If the script is destined for HPC array runs, preserve the command-line interfac `--number_of_cores`) — the `hpc/batch_cpu/template` and `hpc/batch_gpu/template` submit scripts run `scripts/$SCRIPT` and depend on it. +**If the run is on CPU, wire the two acceleration regimes before submitting** — a SLaM pipeline +spans both, and getting it wrong costs days: SOURCE LP is parametric and wants **JAX** +(`use_jax=True`, no `number_of_cores`), while every pixelised stage after it wants the **sparse +operator formalism** (`dataset.apply_sparse_operator_cpu()` + `use_jax=False` + +`number_of_cores=N`). That means two `af.SettingsSearch` objects and two datasets. The workspace +SLaM scripts do **not** do this by default. See +[`al_configure_search`](./al_configure_search.md) "Branch — CPU acceleration" and +`autolens_workspace:scripts/imaging/features/pixelization/cpu_fast_modeling.py`. + Source paths (relative to `autolens_workspace/scripts/`): ``` diff --git a/skills/start-new-project.md b/skills/start-new-project.md index 54d5ebe..92e0bf3 100644 --- a/skills/start-new-project.md +++ b/skills/start-new-project.md @@ -21,7 +21,7 @@ project's whole lifecycle. There is no second science-project skill — this is > README examples); a real analysis headed for a paper gets its own project. The project copies only what's needed to **reproduce the science** (`config/`, `activate.sh`, -`scripts/`, `data/`, `results/`, `hpc/`) — not the copilot's brain (`skills/`, `wiki/core/`, +`scripts/`, `dataset/`, `results/`, `hpc/`) — not the copilot's brain (`skills/`, `wiki/core/`, `wiki/literature/`, `autoassistant/`, `modes/`), which it refers back to. This keeps the published paper repo clean: a reviewer cloning it sees the analysis, not the whole assistant. @@ -72,10 +72,20 @@ Store as `PROJECT_NAME`. Store as `PROJECT_DESCRIPTION`. ### 3. Datasets -> **Datasets to include?** In a project they live under `data///` -> (`` a grouping dir, e.g. `imaging/`; the assistant clone's equivalent folder is -> `dataset/`); each needs at least `data.fits`, `noise_map.fits`, `info.json` (see -> `wiki/core/operations/dataset.md`). Point me at paths to copy, or skip and add later. +> **Datasets to include?** In a project they live under `dataset///` +> (`` a grouping dir, e.g. `imaging/`); each needs at least `data.fits`, +> `noise_map.fits`, `info.json` (see `wiki/core/operations/dataset.md`). Point me at paths to +> copy, or skip and add later. + +> **The folder is `dataset/` — never `data/`. Be pedantic about this.** It is +> `dataset/` everywhere in the PyAuto workspace convention: this assistant clone, every +> `autolens_workspace` example, `PyAutoReduce`'s scripts/docs, and the `--sample`/`--dataset` +> CLI idiom (`dataset///`). A project that uses `data/` silently +> diverges — scripts, `hpc/sync`'s `DATA_DIRS`, and anything copied from the workspace all +> assume `dataset/`, so the mismatch surfaces later as "why did nothing sync / why can't the +> script find the dataset". Do not "improve" on the name, and do not accept `data/` from an +> existing project without flagging it. (Note `project.yaml`'s `data:` *key* is a different +> thing — a metadata block, not the folder; leave it as `data:`.) ### 4. Modeling scripts > **Modeling scripts?** They live in `scripts/`, normally adapted from `autolens_workspace`. @@ -91,8 +101,9 @@ reproducible-science subset; generate the thin assistant layer; refer back for e - `config/` (PyAutoConf YAML — required: pipelines `conf.instance.push(config, output)`) - `activate.sh` (sourced locally and by HPC batch scripts) - `scripts/` (the chosen pipeline(s), or empty + `/init-slam` later) -- datasets (Step 3) into `data//...` (the project's tracked-by-README data tree — - see the `.gitignore` below and the Publish gate, which audits `git ls-files data/`) +- datasets (Step 3) into `dataset//...` — **`dataset/`, not `data/`** (workspace + convention; see the pedantic note in Step 3). The project's tracked-by-README dataset tree — + see the `.gitignore` below and the Publish gate, which audits `git ls-files dataset/` **Generate the lean project tree:** ``` @@ -104,7 +115,7 @@ reproducible-science subset; generate the thin assistant layer; refer back for e .claude/settings.json # PyAuto* API code-gate via refer-back (below) project.yaml # minimal manifest incl. assistant_ref (below) config/ activate.sh scripts/ # copied above - data/ (datasets) + dataset/ (datasets — `dataset/`, NEVER `data/`; workspace convention) results/{manifests,figures,tables}/.gitkeep # manifests/figures/tables TRACKED paper/{figures,tables}/.gitkeep wiki/project/ # journal — copy _profile_template.md + _template.md + README; @@ -242,11 +253,12 @@ assistant clone — promotion is deliberate, never the default. **`.gitignore`** (exclude data/output/secrets/cloned-assistant; **keep** manifests/figures/journal): ``` -data/raw/* -data/reduced/* -data/external/* -!data/**/README.md -!data/**/.gitkeep +dataset/raw/* +dataset/reduced/* +dataset/external/* +!dataset/**/README.md +!dataset/**/.gitkeep +!dataset/**/info.json output/ results/runs/ scripts/scratch/* @@ -302,7 +314,7 @@ on two things only — **no transcript/hash machinery**: "environment_file": "environment.yml", "python_version": "3.11.x", "package_versions": { "autolens": "", "autofit": "", "numpy": "", "jax": "" }, "seed": 42, - "inputs": [{ "path": "data/reduced/slacs0946/data.fits", "sha256": "" }], + "inputs": [{ "path": "dataset/reduced/slacs0946/data.fits", "sha256": "" }], "outputs": [{ "path": "results/figures/fit.png", "sha256": "" }], "started": "", "finished": "", "notes": "smooth SLaM baseline" } @@ -351,7 +363,8 @@ section now speaks to them too). Gate — confirm **every** item before the repo goes public (`visibility_stage: public`): -- [ ] **No raw/restricted data tracked** (`git ls-files data/` shows only READMEs/`.gitkeep`); +- [ ] **No raw/restricted data tracked** (`git ls-files dataset/` shows only + READMEs/`.gitkeep`/`info.json`); `data.publish_raw` still `false` unless the user explicitly cleared it. - [ ] **No full transcripts / scratch / secrets** in history (`.env`, keys, `scripts/scratch/`). - [ ] **LICENSE** chosen and added (e.g. MIT code; CC-BY-4.0 for shared figures/data); set diff --git a/wiki/core/operations/hpc.md b/wiki/core/operations/hpc.md index c4e12ef..8a127c2 100644 --- a/wiki/core/operations/hpc.md +++ b/wiki/core/operations/hpc.md @@ -5,8 +5,9 @@ sources: paths: - scripts/guides/hpc/README.md - scripts/guides/hpc/example_cpu.py + - scripts/imaging/features/pixelization/cpu_fast_modeling.py pinned_commit: main -last_updated: 2026-05-22 +last_updated: 2026-07-15 --- # HPC and cluster runs @@ -20,27 +21,54 @@ This page is the high-level guide. The canonical worked example is `autolens_workspace:scripts/guides/hpc/example_cpu.py` and the `batch/` subfolder beside it. -## CPU parallelism +## CPU acceleration — two regimes -PyAutoFit's non-linear searches accept a `number_of_cores` argument: +On CPU there are **two different accelerators**, and which one is correct depends on the +**source model**. This is the biggest CPU runtime lever; choosing wrong costs days. + +| Fit type | Accelerator | `use_jax` | `number_of_cores` | Dataset | +|---|---|---|---|---| +| Parametric source (e.g. SOURCE LP) | JAX | `True` | leave unset | plain | +| Pixelised source (SOURCE PIX, LIGHT, MASS) | Sparse operators (numba) | `False` | node core count | `apply_sparse_operator_cpu()` | + +**JAX is not GPU-only.** For *parametric* fits it vectorises the likelihood and parallelises +efficiently on CPU, and it is the right choice there even with no GPU present. Because JAX +disables Python multiprocessing, such a fit gains nothing from `number_of_cores` — leave it +unset (it defaults to 1). + +**Pixelised sources use sparse operators instead.** Pixelised reconstruction leans on sparse +linear algebra; `dataset.apply_sparse_operator_cpu()` precomputes operator matrices once +(seconds to minutes) which every later pixelised fit reuses, for a large CPU speed-up. The +implementation is numba-based and does **not** support JAX, so pair it with `use_jax=False` +and multiprocessing: ```python -search = af.Nautilus( - path_prefix="...", - name="...", - n_live=200, - number_of_cores=16, # parallel likelihood evaluations +# Parametric source on CPU — JAX, no number_of_cores. +analysis = al.AnalysisImaging(dataset=dataset, use_jax=True) +settings_search = af.SettingsSearch(path_prefix="...", unique_tag="...", session=None) + +# Pixelised source on CPU — sparse operators, JAX off, multiprocessing on. +dataset = dataset.apply_sparse_operator_cpu() +analysis = al.AnalysisImaging(dataset=dataset, use_jax=False) +settings_search = af.SettingsSearch( + path_prefix="...", unique_tag="...", session=None, number_of_cores=16 ) ``` -Set this to the number of CPU cores on your node. PyAutoFit dispatches likelihood -evaluations across them via multiprocessing. +`number_of_cores` reaches the search via `af.SettingsSearch` (one of `search_dict`'s fixed +keys). PyAutoFit then dispatches likelihood evaluations across those cores. + +A **SLaM pipeline spans both regimes**: SOURCE LP is parametric (JAX), every stage after it is +pixelised (sparse/CPU) — so build two `SettingsSearch` objects and two datasets. Worked example: +`autolens_workspace:scripts/imaging/features/pixelization/cpu_fast_modeling.py`. ## JAX on GPU With JAX installed against CUDA, the likelihood evaluations and array ops will run on the GPU automatically. A modern consumer GPU (24 GB VRAM) handles galaxy-scale fits comfortably. Cluster-scale or pixelised fits with large meshes need 40+ GB. +On GPU, JAX is used throughout the pipeline and the sparse operator is not applied +(it is CPU/numba-only). Set `XLA_PYTHON_CLIENT_PREALLOCATE=false` if running multiple JAX processes on the same node (e.g. one per lens in a batch) to prevent each one from pre-allocating