[2610] Agent context docs

AGENT-README.md

@@ -0,0 +1,104 @@
+# WeatherGenerator
+
+A machine-learning Earth system model: a hierarchical transformer trained on diverse
+data streams (ERA5 reanalysis, satellite and in-situ observations, ocean/climate model
+output) with self-supervised student-teacher losses (JEPA/DINO-style) alongside
+physical prediction losses. Data from all streams is tokenized onto the HEALPix grid,
+assimilated into a latent state, optionally rolled out in time autoregressively, and
+decoded per stream.
+
+This file holds always-relevant rules for LLM tools and an index into the reference
+docs; keep it lean and put detail in `agent_docs/` and the `DOCS-*.md`
+reference files.
+
+Every change must keep context in sync: update the `agent_docs/` files, the `DOCS-*.md`
+reference files, and the lines in this file that describe what you changed, in the same
+change — stale context misleads the next agent. The index below says which docs cover what.
+The code may be newer than the docs — where code and docs disagree, trust the code and
+update the docs.
+
+## Repositories & branches
+
+- Upstream is `github.com/ecmwf/WeatherGenerator`; development happens on personal
+  forks with PRs back to upstream. Branches: `develop` (default; fast-moving, breaking
+  changes), `main` (stable, for experiments), `develop-ssl` (fast-moving pretraining
+  experiments).
+- `../WeatherGenerator-private` (sibling repo): HPC-specific paths, private configs,
+  SLURM launch scripts. May contain credentials — never read files that look like
+  secrets. Split rule: anything that defines a model or affects training results goes
+  in this repo; machine/team-specific paths and settings go in the private repo.
+  Never hardcode HPC paths in this repo. Details: `agent_docs/infrastructure.md`.
+- Some users also have a `weathergen-research` sibling (experiment planning,
+  documentation, literature) with its own instruction files.
+
+## Environment & tooling
+
+- Python 3.12, managed with uv. Deps declared in `pyproject.toml`, locked in `uv.lock`.
+- `uv sync` to set up, `uv run <cmd>` to run. Never `pip install` into the env — it bypasses the lock.
+- Entry points (`pyproject.toml [project.scripts]`, run as `uv run <cmd>`): `train`,
+  `train_continue`, `inference`, `evaluate`, `export`, `plot_train`. CLI flags are
+  dash-separated (`--run-id`, not `--run_id`); full CLI in `src/weathergen/utils/cli.py`.
+  Invocations, run-id mechanics, and HPC launching: `agent_docs/infrastructure.md`.
+- Dev tasks: `scripts/actions.sh {lint|lint-check|type-check|unit-test|toml-check}`; `integration-test*` targets need a GPU.
+
+## Checks & style
+
+- Run `./scripts/actions.sh lint` before considering a change done. CI
+  (`.github/workflows/ci.yml`) runs lint-check, toml-check, type-check, unit-test, and
+  requires the PR to be linked to a GitHub issue (`scripts/check_gh_issue.py`) — all
+  must pass.
+- `unit-test` runs pytest on the `*_test.py` files colocated in `src/` (what CI runs);
+  the standalone tests in `tests/` run via `uv run --extra cpu pytest tests/` and are
+  not currently in CI.
+- Line length 100, ruff formatting; type hints required (pyrefly); use `logging`,
+  never `print`; match existing naming conventions. Do not commit large binary blobs
+  (>2MB commits are rejected — TODO: check this holds on upstream) or the contents of
+  runtime output dirs.
+
+## Layout
+
+- `src/weathergen/` — core model + training code. Entry points in `run_train.py`;
+  `datasets/` (multi-stream loading, readers, tokenizer, masking), `model/`
+  (encoder, engines, attention, EMA), `train/` (trainer, losses, LR schedule, SSL
+  teacher), `utils/` (CLI, distributed, logging, plotting)
+- `packages/` — uv-workspace libraries (common, evaluate, metrics, readers_extra).
+  Also here but not workspace members: `dashboard/` (own lockfile —
+  `agent_docs/decisions/dashboard-not-in-workspace.md`) and `science/` (standalone
+  analysis scripts)
+- `config/` — YAML run configs; `config/streams/` — per-stream config sets
+- `tests/` — standalone unit tests; `integration_tests/` — GPU integration tests
+- `agent_docs/` — agent-oriented docs indexed below: systems dataflows, `recipes/`
+  (procedures), `decisions/` (rationale); how the setup works: `agent_docs/agentic-setup.md`
+- `docs/` — human-facing reference (e.g. `docs/evaluate_config_reference.md`)
+- `ci/`, `.github/workflows/` — CI definitions
+- `logs/`, `models/`, `plots/`, `results/` — runtime output, gitignored; on HPC these are symlinks into shared storage. Never commit contents; details in `agent_docs/infrastructure.md`.
+
+## Documentation (read when relevant)
+
+Full index of `agent_docs/`. Every new doc gets a line here — procedure:
+`agent_docs/recipes/add-documentation.md`.
+
+Systems (runtime dataflows):
+- `agent_docs/training-step.md` — the end-to-end training step: trainer, model forward, losses. Read before changing any of those; update after.
+- `agent_docs/data-pipeline.md` — stream configs → readers → tokenizer → ModelBatch. Read before touching data loading or stream configs; update after.
+- `agent_docs/ssl-training.md` — SSL/student-teacher delta: masking, teachers, latent losses. Read before touching SSL or masking code; update after.
+- `agent_docs/config-system.md` — config sources, merge precedence, stage configs, runtime mutation. Read before adding/renaming config options; update after changing the merge logic.
+- `agent_docs/infrastructure.md` — software stack (cluster base env + uv), SLURM/GH200 hardware, run commands and run IDs, runtime output dirs; local runs not supported yet. Read before setting up an env, launching or continuing runs, or touching run outputs; update after changing tooling or workflow.
+- `agent_docs/agentic-setup.md` — how these instruction files work (opt-in loading) and what belongs in AGENT-README.md vs agent_docs/. Read before editing this file or adding documentation.
+
+Recipes (procedures):
+- `agent_docs/recipes/add-data-reader.md` — add a reader for a new data source.
+- `agent_docs/recipes/add-documentation.md` — add or change agent documentation.
+
+Decisions (rationale):
+- `agent_docs/decisions/dashboard-not-in-workspace.md` — why packages/dashboard has its own lockfile.
+
+Directory-scoped reference (`DOCS-*.md`, next to the code): file-by-file detail of one
+subsystem and how its scripts function — vs `agent_docs/`, which holds cross-directory
+workflows, coupling, and rationale (split defined in `agent_docs/agentic-setup.md`):
+- `config/DOCS-Config.md` — config merging, default_config.yml options, variants.
+- `config/streams/DOCS-Streams.md` — stream YAML schema, readers, adding a stream.
+- `src/weathergen/model/DOCS-model.md` — architecture: encoder, engines.py classes, attention, EMA.
+- `src/weathergen/datasets/DOCS-Datasets.md` — data pipeline: samplers, readers, tokenizer, masking.
+- `src/weathergen/train/DOCS-Train.md` — trainer, loss system, SSL teacher, checkpointing.
+- `packages/DOCS-Packages.md` — workspace packages: common, evaluate, metrics, readers_extra.

README.md

@@ -35,3 +35,8 @@ features. However, it is currently evolving at a fast pace. It should not be exp
 This software is licensed under the terms of the Apache Licence Version 2.0 which can be obtained at [http://www.apache.org/licenses/LICENSE-2.0](http://www.apache.org/licenses/LICENSE-2.0).
 
 In applying this licence, ECMWF does not waive the privileges and immunities granted to it by virtue of its status as an intergovernmental organisation nor does it submit to any jurisdiction.
+
+
+# Agentic coding
+
+Context for coding agents is captured in `AGENT-README.md`. To use it, mention `AGENT-README.md` in your context file (e.g. `WeatherGenerator/AGENTS.md` or `WeatherGenerator/CLAUDE.md`).

agent_docs/agentic-setup.md

@@ -0,0 +1,130 @@
+# Agentic setup
+
+How LLM coding tools get their instructions in this repo.
+
+## Single source of truth, loaded by opt-in
+
+`AGENT-README.md` at the repo root is the canonical instruction file, and it is
+**opt-in**: the repo tracks no `CLAUDE.md`/`AGENTS.md`/`GEMINI.md`, so no LLM tool
+loads context from this repo by default. Each user points their personal, untracked
+instruction file at it — for Claude Code, an `@AGENT-README.md` import line in
+`CLAUDE.local.md` (repo-specific, gitignored) or in `~/.claude/CLAUDE.md`; other
+tools use their equivalent personal-context mechanism. This keeps one tracked file
+to edit (no per-provider copies to drift) while letting users combine it with
+personal instructions or opt out entirely.
+
+Consequence: nothing here is auto-loaded for every agent. Write `AGENT-README.md`
+assuming it is the only instruction file in context, and keep `agent_docs/`
+discoverable from it by explicit root-relative pointers.
+
+Planned, not yet implemented: scoped instruction files inside `src/`, `packages/`,
+`config/`, `tests/` for directory-local rules. Until they exist, scope-local rules
+live in `AGENT-README.md` or the relevant `agent_docs/` doc.
+
+## What goes where
+
+The common agent task — assess or change a code section — has three phases that fail
+differently: routing (find the right files; failure = wasted tokens), understanding
+(grasp enough to change it correctly; failure = a plausible wrong edit), and
+verification (run the right checks). Effectiveness is decided in the understanding
+phase, by knowledge the code itself cannot state. Documentation priority: effectiveness
+first; token efficiency second, and only where it doesn't compromise effectiveness.
+
+**AGENT-README.md — in context every session (once opted in). Route, rule, verify.
+Nothing else.**
+
+- Orientation map: enough for the first search to land right.
+- Hard rules and invariants that apply to most edits.
+- Verification commands.
+- Pointers to `agent_docs/` with a trigger phrased in task vocabulary, bidirectional
+  where the doc tracks code: "read before touching X; update after changing X". The
+  update half is what keeps central docs from going stale — the agent editing code
+  has the reminder in context. Always reference docs by root-relative path
+  (`agent_docs/<name>.md`).
+- Inclusion test per line: would an agent lacking it take a wrong turn or make a wrong
+  edit? No → cut it, or demote it to `agent_docs/`.
+- State each fact in exactly one place — `AGENT-README.md` and the docs load
+  together, so duplication is paid twice. Pointers are the one exception: redundant
+  pointers are cheap and aid discovery; facts are not exempt.
+
+**agent_docs/ — loaded on demand. Everything a correct edit of a given kind needs.**
+
+- Three genres, split by the question they answer; no fact belongs to two:
+  - `agent_docs/*.md` — systems: what happens at runtime, one doc per dataflow (not per
+    directory). Template: ~5-line summary, H2 per stage with `file:symbol` anchors
+    and data shapes at boundaries, then a "Coupling & invariants" section ("change
+    X → check Y") — that section is the payload.
+  - `agent_docs/recipes/` — procedures: how to make a change of kind X. Trigger line,
+    numbered steps with file anchors, verification command, pitfalls.
+  - `agent_docs/decisions/` — rationale: why something is deliberately this way
+    (context → decision → consequences). Prevents "fixing" intentional choices.
+  - Mechanism goes in systems, procedure in recipes, rationale in decisions; the
+    latter two link into systems docs instead of restating them.
+
+- Document the delta between what the code says and what is true: dataflow across
+  files, coupling ("if you change X, also update Y"), invariants, which of two
+  similar mechanisms is canonical, why something is deliberately unusual. A file
+  inventory is grep-replaceable — low value; coupling and rationale are not — high value.
+- Structure for partial reads: conclusion/summary first, clear H2 sections, `file:symbol`
+  anchors. Agents often read only the top of a file.
+- Every doc must be listed in the `AGENT-README.md` documentation index; an unlisted
+  doc isn't discovered reliably. Procedure: `agent_docs/recipes/add-documentation.md`.
+- All agent docs live in top-level `agent_docs/` — the highest-value content is
+  cross-cutting and has no home directory, and one tree keeps grep-fallback discovery
+  reliable. Scoping comes from the pointer's trigger phrase, not from the doc's
+  location (one file, N pointers — this is also what avoids repeating shared content).
+  Central files never move when coupling grows, so pointers stay valid. Flat until it
+  hurts; then group by subsystem or task, never by directory. Name files in task
+  vocabulary (`config-system.md`, `masking.md`) — filenames are matched by searches.
+  Exceptions: human-facing reference stays in `docs/` (e.g.
+  `docs/evaluate_config_reference.md`), directory-scoped reference stays next to the
+  code as `DOCS-*.md` (next block), package `README.md`s stay in their package, and a
+  quasi-independent package (own lockfile, deployable on its own) may keep real docs
+  in-package so they travel if it is extracted.
+- Scope-local rules and coupling one-liners are not docs: they go in
+  `AGENT-README.md` (or, once scoped instruction files exist, the narrowest one that
+  covers all files involved).
+
+**DOCS-*.md — directory-scoped reference, living next to the code it describes.**
+
+- One per subsystem directory (`config/`, `config/streams/`,
+  `src/weathergen/{model,datasets,train}/`, `packages/`): file-by-file inventories,
+  class/function detail with line anchors, schemas, option lists — what is in that
+  directory and how its scripts function. This is deliberately the grep-replaceable
+  layer: faster than reading the code, but derivable from it.
+- The split against `agent_docs/` is scope-by-task vs scope-by-directory, not level
+  of detail (both are detailed): a fact that spans directories — a runtime dataflow,
+  coupling ("change X → also update Y"), a workflow, a design choice — goes in
+  `agent_docs/`; detail local to one directory's files goes in its `DOCS-*.md`. When
+  both need a fact, one states it and the other links to it.
+
+**The pointer/content asymmetry:** a pointer costs ~15 tokens, a missed invariant costs
+a wrong edit. When in doubt, the pointer goes in `AGENT-README.md` and the content
+goes in `agent_docs/`. Be generous with pointers, strict with content.
+
+## Editing rules
+
+- Context maintenance is part of every change, not a documentation task: if an edit
+  alters behavior described in a systems doc, a `DOCS-*.md` reference file, a coupling
+  line, or an `AGENT-README.md` rule/inventory, update that description in the same
+  commit.
+  Agents take statements literally, so stale context turns into wrong edits later.
+  The bidirectional pointer triggers ("update after changing X") mark which docs a
+  change touches; keeping descriptions true is on the author.
+- Facts only — document what the setup is, not what it should be; agents take
+  statements literally. Mark unverified claims as such or leave them out.
+- Plain imperatives over formatting: bold and callouts don't raise compliance. Keep
+  the *why* on non-obvious rules ("never X — it breaks Y") so agents know when a rule
+  generalizes.
+- Personal instructions never go in the tracked files. Use `CLAUDE.local.md`
+  (repo-specific, gitignored) or `~/.claude/CLAUDE.md` (user-global) — the same files
+  that carry the opt-in pointer to `AGENT-README.md`.
+
+## Enforcement
+
+Not yet implemented: a `check-agent-docs` CI check asserting that every
+`agent_docs/...`/`docs/...` path referenced in `AGENT-README.md` exists and that
+every `agent_docs/` file is listed in its documentation index. Until it exists,
+verify links manually (grep) when adding, renaming, or removing docs.
+
+When adding or changing docs, follow `agent_docs/recipes/add-documentation.md`.

agent_docs/config-system.md

@@ -0,0 +1,64 @@
+# Config system — sources, precedence, runtime shape
+
+Summary: configs are OmegaConf objects (`Config` in
+`packages/common/src/weathergen/common/config.py`). A run's config is merged from
+base → private → overwrites (ascending precedence) by `load_merge_configs`, gets a
+run id, is mutated during the run (step counter, streams), and is saved with every
+checkpoint so runs can be continued or inspected.
+
+## Sources and precedence (`config.py:load_merge_configs`)
+
+Ascending precedence (later wins):
+
+1. Base: `config/default_config.yml` (`_DEFAULT_CONFIG_PTH`) — or, when continuing
+   (`--from-run-id`), the saved config of that run (`load_run_config`).
+2. Private config: platform-dependent paths and secrets; the `secrets` section is
+   discarded from the merged result. Found in order (`config.py:_load_private_conf`): explicit
+   `--private-config` path → `WEATHERGEN_PRIVATE_CONF` env var → auto-detection by
+   running `../WeatherGenerator-private/hpc/platform-env.py`.
+3. Overwrites, in order: each `--config` file (paths may be `:`-joined, split
+   automatically), then CLI `--options` as an OmegaConf dot-list
+   (`from_cli_arglist`).
+
+Special case: an overwrite containing `streams_directory` loads those streams and
+replaces any inherited streams (`base_config.streams = None` before merge).
+
+## CLI (`src/weathergen/utils/cli.py`, entry `run_train.py`)
+
+- Stages: `train`, `train_continue`, `inference` (`Stage` enum; positional or
+  `WEATHERGEN_STAGE` env var).
+- Key args: `--config` (repeatable overwrites), `--base-config`, `--private-config`,
+  `--options key=value ...`, `--run-id`/`--reuse-run-id`, `--from-run-id` +
+  `--mini-epoch` for continuation.
+
+## Stage configs (`trainer.py:Trainer.init`)
+
+- `training_config` is the source of truth. `validation_config` and `test_config` are
+  *deltas* applied on top of it (`get_active_stage_config`), in that chain:
+  training → validation → test.
+- Entries in `losses`, `model_input`, `target_input` (`cfg_keys_to_filter`) that are
+  disabled by an overwrite are removed by `filter_config_by_enabled` — that's the
+  mechanism for turning loss terms off per stage.
+
+## Runtime lifecycle
+
+- `set_run_id` assigns/reuses the run id; `config.save(cf, mini_epoch)` writes the
+  config into the run directory at init and with every checkpoint.
+- The config is mutated during the run: `general.istep` (global step),
+  `general.run_history` (continuation lineage), `streams`, `world_size*`. Continuation
+  correctness depends on these saved values.
+
+## Coupling & invariants
+
+- New options belong in `config/default_config.yml` with a sensible default — partial
+  overwrite configs must stay valid against it.
+- Config keys are consumed by name at many sites with `.get(..., default)`; a rename
+  silently falls back to the default rather than failing. Grep for the key before
+  renaming.
+- `losses` / `model_input` / `target_input` are special-cased (`cfg_keys_to_filter`):
+  their entries can be disabled per stage; other sections merge wholesale.
+- Loss-term names in `losses` are cross-referenced by the trainer, the target-aux
+  calculators, and (for SSL) model latent-head names — see `agent_docs/training-step.md`
+  and `agent_docs/ssl-training.md`.
+- `validate_forecast_policy_and_steps` runs per stage config at init; forecast
+  settings must satisfy it for all three stages, not just training.