diff --git a/.claude/notes/architecture.md b/.claude/notes/architecture.md index fd75ab8..a473236 100644 --- a/.claude/notes/architecture.md +++ b/.claude/notes/architecture.md @@ -42,9 +42,13 @@ Four packages under `src/`: **molix** (infra) ← **molrep** (representation) **molrep** — pure representation blocks (no energy/force orchestration) - `src/molrep/` — package root re-exports embedding + ScalarHead + ProductHead + pooling - `src/molrep/embedding/` — JointEmbedding, RBF, cutoffs, SphericalHarmonics -- `src/molrep/interaction/` — ConvTP, SymmetricContraction, RadialWeightMLP, ElementUpdate, … + - `embedding/mace.py` — MACE-only `EmbeddingBlock` / `EmbeddingSpec` +- `src/molrep/interaction/` — SymmetricContraction, RadialWeightMLP, ElementUpdate, ResidualInteraction, EquivariantProductBasis, … + - `interaction/mace/` — MACE-only blocks (`conv.ConvTP`, `block.InteractionBlock`, `density.{DensityInteraction,DensityResidualInteraction,SKIP_TP_METHOD}`) - `interaction/pinet/` — pure GC blocks (FF/message/residual/blocks) -- `src/molrep/readout/` — ProductHead, BasisProjection, masked pooling, NonLinearBiasReadout +- `src/molrep/readout/` — BasisProjection, masked pooling + - `readout/mace.py` — MACE-only `ProductHead`, `LinearReadout`, `NonLinearReadout`, `NonLinearBiasReadout` + - deprecated re-export shims pending removal: `interaction/density.py`, `readout/scalar.py`, `readout/product.py` - `src/molrep/heads/` — ScalarHead, TypeHead, Labeler / ProxyLabeler - `src/molrep/utils/` — geometry helpers + equivariance test utils @@ -55,17 +59,24 @@ Four packages under `src/`: **molix** (infra) ← **molrep** (representation) - `nonbonded.py`, `mixing.py`, `polarization.py` - `elec/` — Ewald/PME/P3M stack (calculators / potentials / lib / kernels / tuning) + multipole - `src/molpot/heads/` — Energy, edge, multipole, charge, electrostatics, rescale, type -- `src/molpot/composition/` — PotentialComposer, MultiHead, Sonata (+ `build_sonata`), parameter heads +- `src/molpot/composition/` — PotentialComposer, MultiHead, Sonata (`from_encoder`), parameter heads - `src/molpot/derivation/` — EnergyAggregation, ForceDerivation (explicit `functorch` | `autograd`), StressDerivation - `src/molpot/pooling/` — Layer / EdgeToNode / Sum / Mean / Max - `src/molpot/graph/radius.py` — `radius_graph` helper (no package `__init__`) **molzoo** — encoder recipes (+ temporary PiNet potential façade) - `src/molzoo/allegro.py` — Allegro encoder + AllegroSpec -- `src/molzoo/mace.py` — MACE encoder + MACESpec -- `src/molzoo/mace_omol.py` — MACEOMol + `load_omol_state_dict` (lazy via PEP 562) +- `src/molzoo/mace/` — industrial split (mace-subpackage-restructure, `1ddd5ff..e825a51`): + - `spec.py` — torch-free `MACESpec` / `MACEMatpesSpec` / `MACEOMolSpec` + - `geometry.py` — `edge_vectors` / `edge_lengths` (PBC shifts) + - `encoder.py` — `MACEEncoder` block graph (no energy, no forces) + - `potential.py` — `MACEPotential` (`energy_core`, `from_checkpoint`, forces) + - `checkpoint.py` — `CheckpointRemap` + `MATPES_REMAP` / `OMOL_REMAP` + - `variants.py` — `MACEMatpes` / `MACEOMol` thin aliases + `load_{matpes,omol}_state_dict` + - `research.py` — encoder-only `MACE` + `MACEResearchSpec` - `src/molzoo/pinet/` — industrial split: spec, geometry, encoder, **potential**, properties -- `src/molzoo/specs/` — `allegro.md`, `mace_omol.md`, `pinet2.md` (**no `mace.md`**) +- `src/molzoo/specs/` — `allegro.md`, `mace.md`, `mace_matpes.md`, `mace_omol.md`, `pinet2.md` + (`mace_omol.md` and `allegro.md` are mirrored under `docs/molzoo/specs/`) **tests** (mirror + regression; not a library package) - `tests/test_molix/…`, `tests/test_molrep/…`, `tests/test_molpot/…`, `tests/test_molzoo/…` @@ -90,7 +101,7 @@ Four packages under `src/`: **molix** (infra) ← **molrep** (representation) | `molix.nn` | `KeyedMLP`, `KeyedMLPSpec`, `NeighborList`, `ScatterSum`, `BatchAggregation` | | `molix.F` | `get_neighbor_pairs`, `pme_direct`, `pme_reciprocal`, `scatter_sum`, `batch_add` | | `molix.io` | `JournalReader`, `JournalWriter` | -| `molix.md` | ForceField family, Langevin integrator, `MDRunner`, trajectory helpers, `make_pinet_calculator` | +| `molix.md` | ForceField family, Langevin integrator, `MDRunner`, trajectory helpers | | `molix.engine` | `EngineAdapter`, `EngineForward`, `FlatTensorAdapter`, `MolnexTensorDictAdapter`, `StaticForward`, `export_for_lammps`, `LAMMPS_META_SCHEMA` | | `molix.compile` / `export` / `quant` | `Compiler` / `Exporter` / quant schemes + `Quantizer` | @@ -105,7 +116,7 @@ Four packages under `src/`: **molix** (infra) ← **molrep** (representation) | `molrep.readout` / `heads` | ProductHead, BasisProjection, pooling; ScalarHead, TypeHead, Labeler | **molpot** (`src/molpot/__init__.py`): -Classical potentials + `BasePotential`; derivation (`ForceDerivation` dual backend); pooling; composition (`PotentialComposer`, `Sonata`, `build_sonata`); thin head re-exports +Classical potentials + `BasePotential`; derivation (`ForceDerivation` dual backend); pooling; composition (`PotentialComposer`, `Sonata`, `Sonata.from_encoder`); thin head re-exports | Subpackage | Exports | |---|---| @@ -114,13 +125,14 @@ Classical potentials + `BasePotential`; derivation (`ForceDerivation` dual backe | `molpot.heads` | Energy/edge/multipole/charge/electrostatics/rescale heads | | `molpot.derivation` | `EnergyAggregation`, `ForceDerivation`, `StressDerivation`, `autograd_forces`, `functorch_forces` | -**molzoo** (`src/molzoo/__init__.py`): -`Allegro`, `AllegroSpec`, `MACE`, `MACESpec`, `MACEOMol` (lazy), `PiNet`, `PiNetSpec`, `load_omol_state_dict` (lazy) +**molzoo** (`src/molzoo/__init__.py`) — **all-lazy** PEP 562 `__getattr__`, no eager import of any model: +`Allegro`, `AllegroSpec`, `MACE` (← `molzoo.mace.research`), `MACESpec` (← `molzoo.mace.spec`), `MACEMatpes`, `MACEOMol` (both ← `molzoo.mace.variants`), `PiNet`, `PiNetSpec`, `load_matpes_state_dict`, `load_omol_state_dict` (← `molzoo.mace.variants`) | Subpackage | Exports | |---|---| +| `molzoo.mace` | eager (torch-free): `MACESpec`, `MACEMatpesSpec`, `MACEOMolSpec`. Lazy (PEP 562): `MACEPotential`, `MACEMatpes`, `MACEOMol`, `MACE`, `MACEResearchSpec`, `CheckpointRemap`, `MATPES_REMAP`/`MATPES_KEY_REMAP`, `OMOL_REMAP`/`OMOL_KEY_REMAP`, `load_matpes_state_dict`, `load_omol_state_dict`, `EmbeddingBlock`/`EmbeddingSpec`, `InteractionBlock`/`InteractionSpec`. `MACEEncoder` and `molzoo.mace.geometry` are reached by module path, not re-exported | | `molzoo.pinet` | `PiNet`, `PiNetSpec`, `PiNetPotential`, `PiNetDipole`, `PiNetPolarizability`, geometry helpers | -| `molzoo.specs/` | markdown only: `allegro.md`, `mace_omol.md`, `pinet2.md`; **`mace.md` missing** | +| `molzoo.specs/` | markdown only: `allegro.md`, `mace.md`, `mace_matpes.md`, `mace_omol.md`, `pinet2.md` | --- @@ -129,7 +141,7 @@ Classical potentials + `BasePotential`; derivation (`ForceDerivation` dual backe - **molix**: PascalCase types + `Hook` suffix; protocols for Step/Hook/DataSource; `TrainState` rejects slash/tuple **writes**; flat-dict pre-collate → nested plain `TensorDict` post-`collate_molecules`; `PackedCache` single-file mmap; arch-tagged native op load; soft-optional `molrs` for MolRec only - **molrep**: pure `nn.Module` + Pydantic `*Spec`; cuEquivariance for TP; PiNet GC blocks have **no** energy/force - **molpot**: `BasePotential` + explicit force backends; `ForceDerivation(method="autograd"|"functorch")` is the shared contract (default `autograd` for cuEq; `functorch` for pure-torch e.g. PiNet); elec is multi-layer calculator/lib/kernel/tuning -- **molzoo**: encoder recipes prefer `TensorDictModuleBase` writing `atoms.node_features` `(N, layers, features)`; paper refs in module docstring; **PiNet potential temporarily co-located under `molzoo.pinet.potential`** (long-term home molpot); MACE-OMOL is a full energy/force model (lazy-loaded) +- **molzoo**: encoder recipes prefer `TensorDictModuleBase` writing `atoms.node_features` `(N, layers, features)`; paper refs in module docstring; **PiNet potential temporarily co-located under `molzoo.pinet.potential`** (long-term home molpot); the MACE foundation variants are full energy/force models — one `molzoo.mace.potential.MACEPotential` with `MACEMatpes` / `MACEOMol` as thin `variants.py` aliases, all reached lazily - **tests**: industrial path mirror; unit tests under `tests/`; numerical parity in `tests/regression/` --- @@ -147,7 +159,7 @@ Classical potentials + `BasePotential`; derivation (`ForceDerivation` dual backe | `molix.md` / `export` / `compile` / `quant` / `engine` | leaf execution utilities → `interface/` C++ | | `molrep.embedding` → `interaction` → `readout`/`heads` | representation pipeline | | `molpot.heads` / `potentials` / `derivation` / `pooling` / `composition` | physics + composition | -| `molzoo.*` | encoder recipes (+ temporary full models: PiNetPotential, MACEOMol) | +| `molzoo.*` | encoder recipes (+ temporary full models: PiNetPotential, `molzoo.mace.MACEPotential` / MACEMatpes / MACEOMol) | | `molzoo.specs` | paper↔code contracts (not runtime) | | `tests/` | unit mirror + regression oracles | @@ -157,6 +169,6 @@ Classical potentials + `BasePotential`; derivation (`ForceDerivation` dual backe 3. Edge: `edge_index[:,0]=source`, `edge_diff = pos[target]-pos[source]`; `bond_index` is `(2,N)` COO. 4. Cache: `PackedCache` only — never `TensorDict.memmap_()`. 5. Forces: single entry `ForceDerivation`; default `autograd` (cuEq-safe); `functorch` only for pure-torch energy graphs. -6. Known gaps: **PiNet energy/force still subclasses `molpot.composition.EnergyForceModel` but lives under `molzoo.pinet` for import stability** (physics path unified); `molpot.composition.pooling` coexists with `molpot.pooling`; full-repo test-mirror is incremental (PiNet spine gated); neighbor kernel is O(N²) pair enum (auto buffer sizing + overflow assert; cell-list still TODO). +6. Known gaps: `molpot.composition.pooling` coexists with `molpot.pooling`; full-repo test-mirror is incremental (PiNet spine gated); neighbor kernel is O(N²) pair enum (auto buffer sizing + overflow assert; cell-list still TODO). *(Closed 2026-08-09 by `mace-subpackage-restructure-07-cleanup`: the zero-subclass energy/force wrapper in `molpot.composition` was deleted — PiNet goes through `molpot.derivation.protocol` helpers, the MACE variants through `molpot.derivation.kernels.grad_force_pass`.)* diff --git a/.claude/notes/learnable-classical-ff.md b/.claude/notes/learnable-classical-ff.md new file mode 100644 index 0000000..5c75dff --- /dev/null +++ b/.claude/notes/learnable-classical-ff.md @@ -0,0 +1,136 @@ + +# Learnable classical FF — placement & reuse (2026-08-10) + +## Why + +The architecture must sit **between** continuous ML chemical perception and +classical MM, without re-implementing molpy/molrs force-field infrastructure. + +## Rule (binding for `learnable-classical-ff-*` specs and all new code) + +### 1. Prefer existing modules + +Before adding a type in molnex, search molpy (≥0.13) / molrs (via molpy only) / +in-tree molpot/molrep/molix. Prefer **reuse** or **generalize**; invent only +when no existing owner fits. + +Known homes (do not fork): + +| Concern | Owner | +|---------|--------| +| Force-field model (styles, types, params) | `molpy.core.forcefield.ForceField` / `molpy.potential.*` (molrs-backed) | +| Classical non-torch E/F | `forcefield.to_potentials().calc_energy/forces(frame)` | +| Topology **enumerate** (optional offline) | molrs `Topology` via molpy — produces index **columns**, not a second store | +| Batch topology **storage** | **molix TensorDict only** — never park angles/propers in molpy Frame for the ML path | +| SMARTS match / typifier base | `molpy.typifier.smarts.SmartsTypifier`, `molrs.perceive.SmartsPattern` | +| Torch classical terms for **differentiable** training | `molpot.potentials.*` (align names/params with molpy styles) | +| Batch collate / rebase | `molix.data.collate` | +| Chem perception (learned continuous) | `molrep` + `molzoo` recipes | +| Learnable param heads / IR torch bags for training | `molpot.composition` / thin `molpot.ir` **only as torch-facing view** of Class-I params | + +### TensorDict topology contract (molix — not molpy) + +Post-collate (and flat sample) connectivity for classical MM uses **column +keys under a namespace**, same spirit as molpy/molrs Frame blocks +(`atomi` / `atomj` / …), nested in TensorDict: + +```python +batch["bonds", "atomi"] # (N_b,) long +batch["bonds", "atomj"] # (N_b,) +batch["angles", "atomi"] # (N_a,) — user-facing: td["angles"]["atomi"] +batch["angles", "atomj"] +batch["angles", "atomk"] # central atom for angles is atomj (i-j-k) +batch["propers", "atomi"] # (N_p,) +batch["propers", "atomj"] +batch["propers", "atomk"] +batch["propers", "atoml"] +batch["impropers", "atomi"] # molrs center-first: atomi = center +batch["impropers", "atomj"] +batch["impropers", "atomk"] +batch["impropers", "atoml"] +# optional type columns: +batch["angles", "type"] # (N_a,) or "angle_types" — pick one, document +``` + +**Not** a packed `angle_index [3, N]` as the primary batch schema (that may +exist only as a kernel-local stack at the potential call site: +`torch.stack([atomi, atomj, atomk], dim=0)`). + +**Not** “store the batch in molpy”. molpy may *emit* columns when building a +sample; the live training/MD batch is TensorDict under molix. + +Rebase on collate: each of `atomi`/`atomj`/`atomk`/`atoml` is an atom-index +1-D vector; add `atom_offset` to every present column under the valence +namespaces (register keys in `INDEX_KEYS` or a sibling column registry). + +### 2. Generalize without multi-method switch + +If two modules do similar work, **promote a single more general type** with one +clear responsibility — never a switch: + +```python +# ❌ forbidden in new APIs +Foo(method="a" | "b") +Bar(mode="x") # when mode selects unrelated implementations + +# ✅ required +MoreGeneralFoo(...) # one implementation, broader domain +# or two peer types with distinct names if both must exist: +AutogradForces / FunctorchForces +``` + +Pre-existing `ForceDerivation(method=…)` is **legacy**; do **not** copy this +pattern into new classical-FF surfaces. New force entry points for Class-I +training use one path (prefer `BasePotential.calc_forces` / autograd) unless a +second named type is justified. + +### 3. Non-diff sinks to molpy / molrs + +Anything that does **not** require PyTorch differentiation **must not** be +reimplemented in molnex: + +- valence enumeration, ring/aromatic perception, SMARTS matching +- discrete FF table models, style registries, non-torch energy +- unit conversion tables for engine export (prefer molpy IO / conventions) +- graph chemical feature extraction from molpy `Atomistic` / Frame + +Molnex owns: + +- continuous chem encoder (diff) +- continuous → MM parameter heads (diff) +- torch classical energy for training + autograd forces +- collate of **already-built** valence index tensors into TensorDict batches +- thin adapters: molpy ForceField / Topology / SMARTS hits → torch IR bags + +**Import hard rule (unchanged):** `from molpy import …` only under `src/` / +`tests/`; never bare `import molrs`. + +### Improper index convention + +**Source of truth:** molrs `Topology` impropers = `[center, i, j, k]` (center at +**row 0**). Torch `improper_index` matches that layout. OpenFF trefoil reordering +is an **export/import adapter**, not a second internal convention. + +### Dependency pin + +`molcrafts-molpy>=0.13.0` (molpy 0.13.x line; molrs major.minor paired by molpy). + +## Supersedes + +- Spec drafts that invented parallel SMARTS engines, valence enumerators, or + OpenMM-only IR without molpy ForceField alignment +- Improper “central at row 1” chain-wide default (replaced by molrs center-first) + +## Spec impact (chain) + +| Sub-spec | Must change | +|----------|-------------| +| 01 | IR bags align to molpy Class-I styles; improper center row 0; no second FF model | +| 02 | Collate TensorDict namespaces `bonds`/`angles`/`propers`/`impropers` with **atomi… columns** (`td["angles"]["atomi"]`); enum optional upstream; never molpy as batch store | +| 03–05 | Stay molnex (diff path); inject topology/features from molpy-built tensors | +| 06 | Perception-side merge; physical_eval may call torch kernels or molpy for residuals | +| 07 | **Reuse** SmartsTypifier / SmartsPattern; no second matcher engine | +| 08 | Prefer molpy forcefield IO / conventions; case matrix still molnex if missing | +| 09 | Thin surfaces only | + +**Status.** active (binding for learnable-classical-ff chain). diff --git a/.claude/notes/notes.md b/.claude/notes/notes.md index c020d20..08cbff2 100644 --- a/.claude/notes/notes.md +++ b/.claude/notes/notes.md @@ -101,7 +101,8 @@ path must work for every potential). Models must not invent a third force path. **Status.** active (promoted into CLAUDE.md Key Design Patterns; docs/gradients aligned). Landed 2026-07-29: OMOL/`energy_forces`, Sonata forces, PiNet via -`EnergyForceModel`, `BasePotential.calc_forces` all go through `ForceDerivation`. +`molpot.derivation.protocol` helpers, `BasePotential.calc_forces` all go +through `ForceDerivation`. `has_aux` supported on functorch backend for single-pass eval. --- @@ -153,14 +154,118 @@ nodes need both x86_64 and aarch64 builds in the same tree. --- -## cuEq `use_fallback` split by force backend (2026-07-29) + +## cuEq `use_fallback` split by force backend (2026-08-08) -**Context.** Hardcoding `use_fallback=True` on every TP / SymmetricContraction -abandoned fused kernels even when forces used autograd (OMOL). +**Context.** Hardcoding the pure-torch path abandoned fused kernels even when +forces used autograd; conversely `MACEMatpes` shipped `use_fallback=True` as +its default — a measured **35.7x** per-step regression with no correctness +upside (its forces are always autograd, so the functorch reason never applies). -**Decision.** Constructors take `use_fallback: bool = True` (functorch-safe -default). Autograd-only full models (e.g. `MACEOMol`) pass `use_fallback=False`. -Encoder-only MACE keeps default True so composed functorch force paths stay -traceable. +**Rule**: Block constructors (`ConvTP`, `SymmetricContraction`, +`DensityInteraction`, `ProductHead`, …) default `use_fallback=True` +(functorch-safe). **Autograd-only full models default `False`** +(`MACEMatpes`, `MACEOMol`); CPU/test call sites pass `True` explicitly. +Composable encoders (`MACE`) expose the knob and inherit the safe default. + +**Supersedes**: the 2026-07-29 entry (which left MACEMatpes on the slow +default and MACE with no knob at all). + +**Status.** active. + +--- + + +## [2026-08-08] cuEq fused-kernel capability is probed, never assumed + +Without the `cuequivariance-ops-torch` wheel, cuEq honours +`use_fallback=False` by silently degrading ~30x (one UserWarning). + +**Rule**: Report fused-kernel status from an actual +`import cuequivariance_ops_torch` probe, never from the `use_fallback` +request flag. GPU installs use the `cueq-cu12` / `cueq-cu13` extras +(pyproject); a degraded run must warn, not self-report "fused". + +--- + + +## [2026-08-08] Headline-number encoders need a benchmark guard + +The `use_fallback` regression shipped because MACE had no benchmark while +carrying the repo's headline GH200 compile numbers. + +**Rule**: Every molzoo encoder whose docs/specs cite performance numbers has a +`benchmarks/bench_.py` guard (see `bench_mace_matpes.py`, which +asserts the fused/fallback ratio). + +--- + + +## [2026-08-08] scatter one-hot GEMM is explicit opt-in + +The one-hot matmul (bit-exact under inductor) measured 2.4–3.3x slower than +`index_add_` at every profiled molecular-graph shape, including inside the +size window that used to auto-select it. + +**Rule**: `scatter_sum_compile_safe` defaults to `index_add_`. The one-hot +GEMM is chosen only by `MOLNEX_SCATTER_ONEHOT=1` (bit-exactness as a +deliberate, global choice) — never by a size heuristic. + + +## [2026-08-09] build.check covers tests/scripts/regressions + advisory ty + +The gate used to lint `src/` only, which let `scripts/` and test-side debt +accumulate invisibly (found during the mace-restructure follow-up sweep). +`benchmarks/` is deliberately excluded while the PiNet bench scripts are +mid-experiment; fold it in once that work lands. + +**Rule**: `mol_project.build.check` runs ruff (check + format) over +`src/ tests/ scripts/ regressions/` plus `ty check src/ +--exit-zero-on-warning` (warnings stay advisory per `[tool.ty.rules]`; +real type errors block). Do not narrow it back to `src/` alone. + + +## [2026-08-09] Construct at config.ftype; init consumes the global RNG + +A 48-site sweep found `nn.Linear`/`nn.Embedding`/`cuet.Linear`/buffer +constructors omitting `dtype=config.ftype`, yielding silent fp32 params +under the fp64 config (hidden by post-hoc `.double()` casts in fixtures). +`cuequivariance_torch.Linear` honours `dtype=` — the old claim that it +ignores `config.ftype` is false. + +**Rule**: every parameter/buffer constructor under `src/` passes +`dtype=config.ftype` explicitly (fp64-contract tests pin the pattern +per module). Weight inits draw from the **global** torch RNG (e3nn +convention; `_ScalarO3Linear` weight ~ N(0,1), bias zero) — never a +private Generator, never zero-init for trainable readout weights. + + +## [2026-08-09] __all__ stays alphabetized + +ruff's isort rule does not cover `__all__` literals, so ordering drifts +silently (caught in molzoo/mace/__init__.py during review). + +**Rule**: keep `__all__` alphabetically sorted in every package +`__init__.py`; re-sort when inserting a name. + + +## [2026-08-09] Type-suppression pragmas: ty syntax only + +The repo's checker is ty (no mypy/pyright config exists). mypy-style +`# type: ignore[code]` is inert under ty — 67 dead pragmas were swept +2026-08-09. Most suppressions are unnecessary anyway: the known +TensorDict/torch-stub false-positive classes are already downgraded to +"warn" via `[tool.ty.rules]` in pyproject.toml. + +**Rule**: never write `# type: ignore[...]`. If a per-line suppression +is truly needed, use `# ty: ignore[rule]`; prefer relying on the +`[tool.ty.rules]` downgrades over per-line pragmas. + +--- + +## Learnable classical FF placement (2026-08-10) + +**Rule.** Reuse molpy≥0.13; no new `Foo(method=…)`; non-diff work sinks to molpy/molrs. +Full text: [learnable-classical-ff.md](learnable-classical-ff.md). **Status.** active. diff --git a/.claude/scheduled_tasks.lock b/.claude/scheduled_tasks.lock deleted file mode 100644 index 5f0cad0..0000000 --- a/.claude/scheduled_tasks.lock +++ /dev/null @@ -1 +0,0 @@ -{"sessionId":"5f5a2a03-31a1-4dd1-88e2-f575a7d0a72b","pid":95863,"procStart":"Mon May 18 12:16:17 2026","acquiredAt":1779107329016} \ No newline at end of file diff --git a/.claude/specs/INDEX.md b/.claude/specs/INDEX.md index fd808d5..73aa6ce 100644 --- a/.claude/specs/INDEX.md +++ b/.claude/specs/INDEX.md @@ -2,10 +2,36 @@ One row per spec generated by `/mn-spec` or `/mol:spec`. Specs are append-only; once ratified, they are the authoritative description of a feature. -- [pinet-quantization-thermal-noise-02-md-driver](pinet-quantization-thermal-noise-02-md-driver.md) — in-process velocity-Verlet+Langevin MD driver + ASE-Calculator shim; paired quantized-vs-fp64 trajectory artifact [approved] -- [md-component-engine-01-compilable-redesign](md-component-engine-01-compilable-redesign.md) — component-based, statically-typed, fullgraph-compilable redesign of molix.md: ForceField(nn.Module)≠Potential (molpy taxonomy), typed MDState/ForceOutput, torch.compile(rollout) incl. PiNet functorch forces 0-break; LJ-cluster NVE 5ns energy-conservation gate; preserves all post-review correctness fixes [approved] +- [pinet-quantization-thermal-noise-02-md-driver](pinet-quantization-thermal-noise-02-md-driver.md) — in-process velocity-Verlet+Langevin MD driver + ASE-Calculator shim; paired quantized-vs-fp64 trajectory artifact [done] - [dynamic-batching-packed-collate-02-collate](dynamic-batching-packed-collate-02-collate.md) — packed-aware collate fast path: slice PackedCache tensors directly into the batched TensorDict, equivalence-oracle tested, full fallback [done] -- [cuet-force-doublebackward](cuet-force-doublebackward.md) — minimal-repro-first fix for cuet equivariant Linear (l>0) severing force-supervised double-backward; root-cause vs working mace teacher, pure-cuet fix, preserve 1e-9 port consistency [approved] +- [cuet-force-doublebackward](cuet-force-doublebackward.md) — minimal-repro-first fix for cuet equivariant Linear (l>0) severing force-supervised double-backward; root-cause vs working mace teacher, pure-cuet fix, preserve 1e-9 port consistency [done] - [mace-omol-port-01-native-port](mace-omol-port-01-native-port.md) — native molnex reimplementation of official MACE-omol-0 (charge/spin, non-linear residual interactions) that loads official weights; every block bit-exact vs official on CPU, full model E/F 7e-7 eV / 4e-6 eV·Å [done] - [mace-omol-port-02-pipeline-integration](mace-omol-port-02-pipeline-integration.md) — wire MACEOMol into molnex runtime: TensorDict adapter + molpot/ForceDerivation + lazy molzoo export + RadialMLP dtype fix; all criteria verified (E/F 7e-7 eV / 4e-6 eV·Å ≤ 1e-4 accepted bar; O3_e3nn bit-exact dropped) [done] - [lammps-pair-molnex-01-aoti-force-export-wall](lammps-pair-molnex-01-aoti-force-export-wall.md) — AOTI exports MACE-OMol energy bit-exact but forces=0 (cuet fused ops have no export-traceable backward); 5-wall investigation-of-record + TorchScript+libtorch pair_style route forward (run backward() in C++) [investigation-complete] +- [learnable-classical-ff-01-ir-kernels](learnable-classical-ff-01-ir-kernels.md) — Potential IR + Class-I kernels (periodic proper, impropers, 1-4 scaling) [done] +- [learnable-classical-ff-02-valence-topology](learnable-classical-ff-02-valence-topology.md) — molix collate namespaces for angles/propers/impropers [done] +- [learnable-classical-ff-03-mm-heads](learnable-classical-ff-03-mm-heads.md) — continuous canonical MM parameter heads + ClassicalMMComposer [done] +- [learnable-classical-ff-04-chem-encoder](learnable-classical-ff-04-chem-encoder.md) — continuous chemical perception (atom/bond embeddings + contexts) [done] +- [learnable-classical-ff-05-neural-parameterizer](learnable-classical-ff-05-neural-parameterizer.md) — encoder→heads→IR→classical E/F composition path [done] +- [learnable-classical-ff-06-condensation](learnable-classical-ff-06-condensation.md) — physics-aware multi-system chemical class condensation [done] +- [learnable-classical-ff-07-smarts](learnable-classical-ff-07-smarts.md) — SMARTS/SMIRKS symbolic interface + SymbolicForceField [done] +- [learnable-classical-ff-08-ff-export](learnable-classical-ff-08-ff-export.md) — Potential IR → OpenMM force-spec compiler (4 translation cases) [done] +- [learnable-classical-ff-09-provenance](learnable-classical-ff-09-provenance.md) — confidence, chemical-space coverage, provenance surfaces [done] +- [mace-neighbor-graph-correctness](mace-neighbor-graph-correctness.md) — independent multi-image oracle vs NeighborList (≤ cutoff); E/F wrap-invariance slow; metrics.jsonl/molplot [approved, grilled] + +## Chain: mm-param-val (MM Parameter Learning Baseline — milestone 1) + +Espaloma-capability validation before project-specific condensation/IR claims. Datasets via **MolHub only**; workspace/workflows via **molexp**; hard-to-find surfaces via **molmcp** route hints. Architecture primitives already landed in `learnable-classical-ff-01..09`. + +Later milestones (not yet specified): B3 energy+force, C torsion/minimize, D1–D4 QM, OPLS teacher, physics-aware condensation. + +| # | Spec | Owner surface | One-line | +|---:|---|---|---| +| 01 | mm-param-val-01-molhub-contract | molhub dataset | molecule_id, units, TargetSchema families, MoleculeSplit **[done]** | +| 02 | mm-param-val-02-zinc-typing | molhub + registry | `dataset:espaloma/zinc-typing@1` **[done]** | +| 03 | mm-param-val-03-phalkethoh-mm | molhub + registry | `dataset:espaloma/phalkethoh-mm-small@1` **[done]** | +| 04 | mm-param-val-04-workspace | molexp + molmcp + molnex scripts | workspace scaffold + route hints **[done]** | +| 05 | mm-param-val-05-potential-parity | molnex molpot | Validation B0 IR/kernel parity **[done]** | +| 06 | mm-param-val-06-typing-recovery | molnex molrep | Validation A GAFF typing probe **[done]** | +| 07 | mm-param-val-07-mm-energy | molnex molix | Validation B1/B2 centered MM energy **[done]** | +| 08 | mm-param-val-08-latent-analysis | molnex molrep.analysis | Validation D/E latent purity + artifacts **[done]** | diff --git a/.claude/specs/mace-neighbor-graph-correctness.acceptance.md b/.claude/specs/mace-neighbor-graph-correctness.acceptance.md new file mode 100644 index 0000000..d9896f2 --- /dev/null +++ b/.claude/specs/mace-neighbor-graph-correctness.acceptance.md @@ -0,0 +1,97 @@ +# Acceptance — mace-neighbor-graph-correctness + +Binding criteria after grill supersede 2026-08-10. Types follow evaluator-protocol. + +## Criteria + +### ac-001 — Independent multi-image brute-force oracle + +- **type:** unit +- **verify:** `tests/test_molix/test_md/oracle_bruteforce_neighbors.py` (or equivalent) implements multi-image directed edges `(i,j,sx,sy,sz)` with **`0 < |dr| ≤ cutoff`**, excludes only true self; **no** imports from `molix.md.neighbors`, `molix.op`/`get_neighbor_pairs`, matscipy, ASE, freud, LAMMPS; **does not** call or copy `test_neighbors._reference_pairs` (MIC-only). +- **status:** verified +- **last_checked:** 2026-08-10 + +### ac-002 — Open random system + +- **type:** unit +- **verify:** N∈[10,30], open; NeighborList fresh build vs oracle → missing=0, extra=0. +- **status:** verified +- **last_checked:** 2026-08-10 + +### ac-003 — Cutoff boundary (≤ convention) + +- **type:** unit +- **verify:** distances `r_c±1e-3` and `r_c±1e-6`; **edge present at `r=r_c`**; absent for `r>r_c` (consistent with `0 < r ≤ r_c`). Failure messages state the `≤` convention. +- **status:** verified +- **last_checked:** 2026-08-10 + +### ac-004 — Periodic wrap + +- **type:** unit +- **verify:** L=10 Å, atoms x=0.1 and 9.9, r_c=1 Å; shared edge keys; physical |dr|≈0.2 Å. +- **status:** verified +- **last_checked:** 2026-08-10 + +### ac-005a — Graph wrap / lattice reparametrization + +- **type:** unit +- **verify:** wrap one atom / translate by lattice / translate whole structure → physical `dr` multisets match (graph-level); missing=extra=0 for each representation after fresh rebuild. +- **status:** verified +- **last_checked:** 2026-08-10 + +### ac-005b — E/F invariance under PBC-equivalent reparametrization + +- **type:** unit +- **verify:** same representations as ac-005a; MACE `energy_core` total E and forces agree within a fixed tol when lists are rebuilt. Marked **`@pytest.mark.slow`** and **skipped when MatPES weights are unavailable**. Not SO(3) equivariance. +- **status:** pending + +### ac-006 — Triclinic + +- **type:** unit +- **verify:** small triclinic cell; missing=extra=0. +- **status:** verified +- **last_checked:** 2026-08-10 + +### ac-007 — Multi-image small cell + +- **type:** unit +- **verify:** box edge ≲ 2 r_c; multi-image edges allowed; missing=extra=0 (not MIC-only). +- **status:** verified +- **last_checked:** 2026-08-10 + +### ac-008 — Cutoff-crossing frames + +- **type:** unit +- **verify:** two-atom trajectory r>r_c → r≤r_c → r>r_c; **per-frame rebuild**; no missing/extra; no hysteresis. +- **status:** verified +- **last_checked:** 2026-08-10 + +### ac-009 — Trajectory CLI + molplot metrics + +- **type:** unit +- **verify:** `scripts/matpes_port/neighbor_graph_audit.py` (or path named in Files) over multi-frame input: table `frame n_ref n_mace missing extra`; mismatch dump with i,j,S,distances; writes `metrics/metrics.jsonl` with scalar keys `n_ref`, `n_mace`, `missing`, `extra` (and `max_dr_mismatch`); non-zero exit on any mismatch. +- **status:** verified +- **last_checked:** 2026-08-10 + +### ac-010 — Knowledge + experiment scaffold in mace-nve + +- **type:** manual +- **verify:** Note under `mace-nve` project `mace-r2san` holds the design; experiment `neighbor-graph-correctness` exists. (Seeded 2026-08-10; re-sync Note body after this supersede.) +- **status:** verified +- **last_checked:** 2026-08-10 +- **verified_by:** agent-auto (Note + experiment seeded) + +### ac-011 — LAMMPS dump (optional) + +- **type:** manual +- **verify:** non-invasive dump of LAMMPS→molnex neighbor payload **or** documented seam + “not in this suite” in the knowledge Note. +- **status:** verified +- **last_checked:** 2026-08-10 +- **verified_by:** agent-auto (seam documented in CLI; no invasive dump) + +### ac-012 — Dual-backend light coverage + +- **type:** unit +- **verify:** at least one synthetic case runs against both default/`bin=` paths when both backends are available (or skip with reason if one backend missing). +- **status:** verified +- **last_checked:** 2026-08-10 diff --git a/.claude/specs/mace-neighbor-graph-correctness.md b/.claude/specs/mace-neighbor-graph-correctness.md new file mode 100644 index 0000000..9e5b209 --- /dev/null +++ b/.claude/specs/mace-neighbor-graph-correctness.md @@ -0,0 +1,151 @@ +--- +title: MACE neighbor-graph correctness vs independent brute-force oracle +status: code-complete +slug: mace-neighbor-graph-correctness +created: 2026-08-10 +revised: 2026-08-10 +grilled: true +knowledge_home: mace-nve/projects/mace-r2san (Note + experiment neighbor-graph-correctness) +supersede: grill-2026-08-10 (cutoff ≤; E/F slow gate; no MIC oracle reuse; no O3 equivariance) +--- + +# MACE neighbor-graph correctness vs independent brute-force oracle + +## Summary + +Implement a **self-contained neighbor-graph correctness suite** for debugging NVE energy drift: compare an **independent brute-force O(N²) oracle** against the neighbor graph that **molnex MACE actually consumes** (`molix.md.NeighborList` → `edge_index` `(E,2)` + integer/image `shifts`), with optional instrumentation of the **LAMMPS/ML-IAP → molnex** path when available without invasive production changes. + +This is **not** a re-test of matscipy/ASE/LAMMPS against themselves. The oracle enumerates atoms × periodic images and returns directed edges `(i, j, sx, sy, sz)` with **`0 < |dr| ≤ cutoff`** (matching production filter parity). + +**Plotting / observability:** diagnostics are written as **molrec-compatible `metrics/metrics.jsonl`** (keys `n_ref`, `n_mace`, `missing`, `extra`, `max_dr_mismatch`) so the **molexp molplot plugin** can render curves from the run; human tables go to `artifacts/*.txt`. No ad-hoc matplotlib as the primary chart path. + +**Primary knowledge home:** design also lives as a molexp **Note** under the `mace-nve` workspace project `mace-r2san`, experiment `neighbor-graph-correctness`. + +## Domain basis + +- Edge convention (molnex hard rule): `edge_index[:,0]=source`, `edge_index[:,1]=target`, + `edge_diff = pos[target] - pos[source] (+ shift for PBC)`, + `dr = pos[j] - pos[i] + S @ cell` for integer image vector `S=(sx,sy,sz)`. +- Exclude only the true self-edge `i==j && S==(0,0,0)`. Periodic self-images inside cutoff **remain valid**. +- **Cutoff convention (pinned to code):** production pair filter is **`0 < r ≤ r_c`** (equivalently `distance2 > cutoff2 || distance2 == 0` drops; **`r == r_c` is kept**). Documented in `neighbors.py` filter parity and `get_neighbor_pairs.cu`. Oracle and tests **must** use the same inequality. +- **Multiple images:** when any box edge ≲ `2 r_c`, minimum-image alone is insufficient; the oracle uses an image range large enough that all images with `0 < |dr| ≤ r_c` are found. +- Verlet skin / every / delay affect **when** the list rebuilds in MD, not completeness of a **fresh** build. Fresh-build tests use construction / `rebuild`; cutoff-crossing evaluates **each frame independently** (no stale list). +- **PBC wrap / lattice reparametrization:** same physical configuration → **E invariant**, **F numerically the same** on corresponding atoms (invariance under reparametrization — **not** SO(3) equivariance). +- **O(3) equivariance** (E invariant + F transforms as vectors under rotation) is **out of scope** for this suite. +- Implementation target: **molnex `NeighborList` + edges fed to MACE**. Not `mace.data.get_neighborhood` / mace-torch (forbidden under `src/` / `tests/`). In-tree there is **no matscipy** neighbor path. + +## Design + +### Discovered neighbor surface (molnex) + +| Layer | Role | +|-------|------| +| `molix.md.NeighborList` | Stateful fixed-capacity list; `rebuild`/`update`/`build(td)`; skin/every/delay/check; `edge_index` `(capacity,2)` + `shifts` `(capacity,3)` Å | +| `get_neighbor_pairs` (C++/CUDA) | Pair search kernel; filter `distance2 > cutoff2 \|\| distance2 == 0` | +| `NeighborList._build_binned` | Pure-torch cell-list (`bin=`); same `0 < r ≤ r_build` filter | +| `molzoo.mace.MACEPotential.energy_core` | Consumes caller `edge_index` + `shifts` — **does not** build neighbors | +| MD bind | Force fields call `neighbors.update` then read live buffers | + +The MACE energy graph is only as correct as the **list supplied to it**; this suite validates that list. + +### Graph representation & comparison + +```text +EdgeKey = (i, j, sx, sy, sz) # ints; S @ cell → Cartesian +dr = pos[j] - pos[i] + (S @ cell) +``` + +- Live edges: `[0, num_edges)`. Convert continuous `shifts` (Å) to integer `S` via `S ≈ shifts @ inv(cell)`, nearest-int; fail if residual large. +- Compare sets: missing = ref \ sut, extra = sut \ ref. **Hard-fail** if either non-empty. +- For matching keys report max ‖dr_ref − dr_sut‖. + +### Brute-force oracle (independent) + +Dedicated test-only module (e.g. `tests/test_molix/test_md/oracle_bruteforce_neighbors.py`): + +- Inputs: `pos (N,3)`, `cell (3,3)|None`, `cutoff`, `pbc (3,) bool`. +- Image range from perpendicular cell widths vs cutoff (not MIC-only). +- Emit directed pairs with **`0 < |dr| ≤ cutoff`**, exclude true self only. +- **Forbidden imports:** `molix.md.neighbors`, `molix.op` / `get_neighbor_pairs`, matscipy, ASE, freud, LAMMPS. +- **Must not** reuse `test_neighbors._reference_pairs` (that helper is **MIC-only** and would fail multi-image cases). + +### SUT coverage + +- **Default:** production `NeighborList(...)` construction / `rebuild` (whatever backend the environment selects). +- **Light dual-backend:** parametrize **at least one** synthetic case with `bin=` vs non-bin/kernel path when both are available — not a full 2× matrix of every test. + +### Required synthetic tests (pytest) + +1. **Open random** — N∈[10,30], no PBC: missing=extra=0 (**unit**). +2. **Cutoff boundary** — `r_c±1e-3`, `r_c±1e-6`; **`r=r_c` has edge** (because ≤). Failures must state the `≤` convention (**unit**). +3. **PBC wrap** — cubic L=10 Å, x=0.1 vs 9.9, r_c=1 Å → |dr|≈0.2 Å; correct `S` (**unit**). +4. **Translation / wrap** + - **Graph (unit):** physical `dr` multisets equal across wrap / lattice translate / whole-structure lattice translate. + - **E/F (slow):** same representations → MACE `energy_core` E and forces agree within fixed tol when list is **rebuilt** each time; mark `@pytest.mark.slow` and **skip if MatPES weights unavailable**. +5. **Triclinic** — small tilted cell (**unit**). +6. **Small cell / multi-image** — box edge ≲ 2 r_c; multi-image oracle (**unit**). +7. **Cutoff-crossing frames** — r>r_c → rr_c; **per-frame rebuild**; no hysteresis (**unit**). +8. **Trajectory CLI** — multi-frame file; table + mismatch dump + **metrics.jsonl** for molplot. + +### Diagnostics & molplot contract + +Per frame / comparison, append to run `metrics/metrics.jsonl`: + +```json +{"t":"scalar","k":"n_ref","s":,"v":...} +{"t":"scalar","k":"n_mace","s":,"v":...} +{"t":"scalar","k":"missing","s":,"v":...} +{"t":"scalar","k":"extra","s":,"v":...} +{"t":"scalar","k":"max_dr_mismatch","s":,"v":...} +``` + +Human table: `artifacts/neighbor_compare.txt`. +molexp molplot reads **metrics.jsonl** (SoT), not PNG. + +### Optional LAMMPS / ML-IAP + +Non-invasive dump of neighbor payload into molnex if feasible; else document the seam under the knowledge Note and mark ac-011 done-with-doc. No production behavior change by default. + +### Reuse decision + +| Candidate | Decision | +|-----------|----------| +| `NeighborList` + MD tests | **reuse** as SUT | +| Skin/ndanger tests | **pattern** only (policy ≠ fresh-build geometry) | +| `_reference_pairs` in `test_neighbors.py` | **do not reuse** as multi-image oracle | +| New multi-image brute force | **new** under tests | +| matscipy / freud / ASE / mace-torch | **forbid** | + +## Files + +- `tests/test_molix/test_md/oracle_bruteforce_neighbors.py` — independent multi-image oracle + compare helpers +- `tests/test_molix/test_md/test_neighbor_graph_oracle.py` — synthetic cases 1–7 (E/F under slow) +- `scripts/matpes_port/neighbor_graph_audit.py` — trajectory CLI + metrics.jsonl (+ optional experiment wiring notes) +- Knowledge: `mace-nve/projects/mace-r2san/mace-neighbor-graph-correctness/` +- Experiment: `…/experiments/neighbor-graph-correctness/` + +## Tasks + +- [x] Pin cutoff comments: production filter is `0 < r ≤ r_c`; note MACE does not build its own list +- [x] Implement independent multi-image brute-force oracle + integer-shift key helpers (no MIC-only reuse) +- [x] Pytest unit: open, boundary (≤), PBC wrap, graph wrap-invariance, triclinic, multi-image, cutoff-crossing +- [x] Pytest slow (or skip-no-weights): E/F invariance under PBC-equivalent reparametrization +- [x] One dual-backend smoke (`bin=` vs default) on at least one case +- [x] CLI trajectory audit + `neighbor_compare.txt` + metrics.jsonl +- [x] Optional LAMMPS dump or documented seam +- [x] `ruff` + targeted pytest green + +## Testing + +- Synthetic **unit** cases hard-fail on any missing/extra edge; print diagnostics on failure. +- E/F path is **not** required for default CI green without weights. +- CLI non-zero exit if any frame mismatches. +- No third-party neighbor library as oracle; no mace-torch / ASE / e3nn. + +## Out of scope + +- Changing production cutoff convention or NeighborList algorithm (separate decision). +- **O(3) energy invariance / force equivariance** under rotation (separate suite). +- Full LAMMPS CI matrix if dump needs invasive patches (document only). +- Oracle performance. +- Training accuracy / general equivariance ports. diff --git a/.claude/specs/mace-omol-port-02-pipeline-integration.md b/.claude/specs/mace-omol-port-02-pipeline-integration.md index 4ba26f6..25087f9 100644 --- a/.claude/specs/mace-omol-port-02-pipeline-integration.md +++ b/.claude/specs/mace-omol-port-02-pipeline-integration.md @@ -78,7 +78,7 @@ the CG-basis convention swap. The post-collate batch schema and edge convention - [x] Backfill `src/molzoo/specs/mace_omol.md` per the molzoo-spec workflow (ac-002, docs) — status `partial`, mirrored to docs + zensical - [x] Model-level E/F accuracy ≤ 1e-4 vs official (ac-003) — met at 7e-7 eV / 4.3e-6 eV·Å with default cue O3 (01 ac-006); 1e-4 is the accepted bar (operator decision) - [x] Drop O3_e3nn pursuit + remove dead `MACEOMol(group=)` hook (CG diff only 1.4e-8/op; 7e-7 is reimplementation accumulation, O3_e3nn would not help) -- [ ] Edge sourcing via `molpot.graph` / `NeighborList` for standalone use (optional; pipeline path sources edges at collate) +- [ ] Edge sourcing via `molpot.graph` / `NeighborList` for standalone use (optional; pipeline path sources edges at collate) (obsolete — molpot.graph removed by md-neighborlist-skin-02-prune; standalone edge sourcing is molix.md.NeighborList's job) ## Testing diff --git a/.claude/specs/md-component-engine-01-compilable-redesign.acceptance.md b/.claude/specs/md-component-engine-01-compilable-redesign.acceptance.md deleted file mode 100644 index 79eeceb..0000000 --- a/.claude/specs/md-component-engine-01-compilable-redesign.acceptance.md +++ /dev/null @@ -1,99 +0,0 @@ ---- -slug: md-component-engine-01-compilable-redesign -criteria: - - id: ac-001 - summary: ForceField/Potential 组件分层(镜像 molpy),无闭包 - type: code - pass_when: | - src/molix/md/forcefield.py 定义 ForceField(nn.Module) 抽象 - (forward(pos: Tensor) -> ForceOutput, 便捷 calc_energy/calc_forces), - PotentialForceField(ForceField) 绑定一个 molpot Potential + 体系模板, - HarmonicForceField(ForceField) 提供解析参考。force_seam.py 已删除, - build_force_fn 闭包不再存在。grep 全 src/molix/md 无 `Callable` - 力别名(ForceFn 删除)、无返回闭包的工厂函数。Integrator 持有 - self.force: ForceField 组件而非 Callable。 - status: verified # last_checked: 2026-06-26 - - id: ac-002 - summary: 热路径静态化 —— 去 getattr/isinstance/None 分支/运行期 gamma 判断 - type: code - pass_when: | - LangevinVerletIntegrator.step 无 `if self.gamma > 0` 分支(O 步恒等式 - 始终执行,γ=0 时 c2=0 数值不变),noise 永远是真实张量(无 None 分支); - mass 在 __init__ 归一化为正张量 buffer(无 isinstance(mass,…)); - MDRunner 的 hook 分发为类型化直调(hook.on_train_batch_end(...)), - 无 getattr(hook, name)。`grep -nE "getattr\(|isinstance\(|is None" - src/molix/md/integrators.py src/molix/md/forcefield.py` 在 step/forward - 热路径上零命中(构造期允许)。 - status: verified # last_checked: 2026-06-26 - - id: ac-003 - summary: typed MDState/ForceOutput 作为组件契约(tensordict↔typed 平衡) - type: code - pass_when: | - src/molix/md/types.py 定义 ForceOutput(NamedTuple: energy, forces) 与 - MDState(NamedTuple: pos, vel, force);二者是 pytree(tree_flatten/ - tree_unflatten round-trip 测试通过)。TensorDict 只出现在 ForceField - 内部(模型 I/O),不出现在 Integrator/MDRunner/Hook 的跨组件签名里; - Integrator.step 的签名是 (MDState, Tensor) -> MDState。 - status: verified # last_checked: 2026-06-26 - - id: ac-004 - summary: Integrator.rollout 全图编译含 PiNet 力,0 graph-break,compile==eager - type: code - pass_when: | - tests/test_molix/test_md_compile.py 用真实 tiny-PiNet 构造 - PotentialForceField,torch.compile(ig.rollout, fullgraph=True) 跑 N>=5 步 - 不抛 graph-break(fullgraph=True 下 break 即抛错→测试失败),且与 eager - rollout 数值一致(fp64 allclose atol<=1e-10);torch.compile(ig.step) 同样 - 0 break 且 ==eager。NVE(γ=0)与 Langevin(γ>0)两路都覆盖。 - status: verified # last_checked: 2026-06-26 - - id: ac-005 - summary: 编译路径 fp32 与 fp64 均正确 - type: scientific - pass_when: | - 同一编译验收在 torch.float32 与 torch.float64 下都通过:fp64 与 eager - 位级/1e-10 一致;fp32 与 eager 在归约精度噪声内(atol 适配,记录实测 - gap)。证明可编译 MD 步对两种精度都数值可信。 - status: verified # last_checked: 2026-06-26 (test_md_compile fp32+fp64 parametrized PASS; fp64==eager, fp32 within noise) - - id: ac-006 - summary: 正确性回归全部保留(防重构倒退) - type: code - pass_when: | - 迁移后这些断言仍绿:活体几何(test_force_seam_tracks_live_geometry 等价: - 能量/力随非刚性位移变化 >1e-6,冻结-PES 不回归)、energy-varies、 - NVE 能量守恒(drift<1e-3)、Langevin 等分(measured kbt 偏差<5%,dof - 约定 γ>0→3N)、force-caching 位级一致、step==step_cached、 - TrajectoryHook shard-flush==单缓冲且清理 shard、ΔF 在 fp64 相减、 - mass<=0 在构造期 raise ValueError。 - status: verified # last_checked: 2026-06-26 - - id: ac-007 - summary: MD 套件 + lint/format 全绿,无外部破坏 - type: runtime - pass_when: | - PYTHONPATH=src:. -m pytest tests/test_molix/test_md_integrators.py - tests/test_molix/test_md_dynamics.py tests/test_molix/test_md_runner.py - tests/test_molix/test_md_compile.py 全绿;ruff check src/molix/md 与 - ruff format --check src/molix/md 干净;无 src 内其它模块导入已删除的 - force_seam/ForceFn(grep 确认)。__init__.py 导出新组件词表。 - status: verified # last_checked: 2026-06-26 - - id: ac-008 - summary: LJ 粒子体系 NVE 5 ns 能量不漂移(长程积分稳定性) - type: scientific - pass_when: | - 新增 LennardJonesForceField(ForceField)(全对 LJ E=4ε[(σ/r)^12−(σ/r)^6], - 无 cutoff/无邻居表→全对天然规避冻结邻居表问题,解析可微)与验证脚本 - benchmarks/verify_md_lj_nve.py:对一个 LJ 团簇(脚本记录 N(~13–55)、 - σ/ε/m、dt、初始温度 T0、平衡构型来源)以 γ=0(NVE)经 torch.compile 的 - rollout 跑满 5 ns(步数 = 5ns/dt,数百万步,只有编译/CUDA-graph 才可行)。 - 判据:总能量无系统漂移 —— 对 E_tot(t) 线性拟合, - |slope·(5 ns)| / |E_tot(0)| < 1e-3,且能量 RMS 涨落有界(不发散、无单调 - 爬升),温度保持有限不爆炸。脚本打印 drift、拟合斜率与 PASS/FAIL; - 同时记录 steps/s 作为编译性能旁证。 - status: verified # last_checked: 2026-06-26 (verify_md_lj_nve --ps 5000: 1e6 steps, rel drift 4.31e-7 < 1e-3, RMS 3.05e-6, PASS) ---- - -# Acceptance criteria - -- **ac-001 / ac-002 / ac-003 是"组件化 + 静态化"三条硬指令的可证伪化**:分别钉死(1)ForceField≠Potential 的 molpy 式分层且无闭包、(2)热路径无动态判断、(3)typed state 作为组件契约。 -- **ac-004 / ac-005 是"`torch.compile(md.run)` 全流程编译"的核心验收**:`rollout`(含 PiNet functorch 力)fullgraph 0 break 且 compile==eager,fp32/fp64 双精度。这是本次重构区别于上一版的根本能力。 -- **ac-006 是防倒退闸**:上一轮 review 修掉的 🚨 冻结-PES 与全部 🔴(显存/ΔF fp64/dof/单位)必须在组件化后逐条仍然成立。 -- **ac-008 是 NVE 积分器金标准**:LJ 团簇(刚硬非谐真实 PES)NVE 5 ns 总能量无系统漂移 —— 比谐振子玩具强得多的长程稳定性检验,且 5ns≈数百万步只有编译路径可行,顺带压测性能。 -- 全部 8 条 verified 后方可 code-complete。 diff --git a/.claude/specs/md-component-engine-01-compilable-redesign.md b/.claude/specs/md-component-engine-01-compilable-redesign.md deleted file mode 100644 index e0ed09f..0000000 --- a/.claude/specs/md-component-engine-01-compilable-redesign.md +++ /dev/null @@ -1,117 +0,0 @@ ---- -title: MD component engine — compilable, component-based redesign of molix.md -status: code-complete -created: 2026-06-26 -chain: md-component-engine ---- - -# MD component engine — compilable, component-based redesign of `molix.md` - -## Summary - -把 `molix.md` 从"闭包 + 动态分发"的研究脚本，重构为一套**基于组件、静态类型化、整步可 `torch.compile` 全图编译**的 MD 引擎。三条硬约束(来自 review 后的用户指令): - -1. **基于组件,删除函数引用** —— 注入的 `force_fn: Callable` 闭包、`build_force_fn` 闭包工厂、测试里的 `_harmonic(k)` 闭包、`MDRunner._call_hooks` 的 `getattr` 动态分发,全部换成 `nn.Module` 组件 / 类型化直调。 -2. **静态、类型化、可编译** —— 去掉热路径里的所有运行期判断(`isinstance(mass,…)`、`if gamma>0`、`noise=None` 分支、`getattr`);`Integrator.step` 与 `Integrator.rollout`(含 PiNet 力)必须 `torch.compile(fullgraph=True)` **零 graph-break**,fp32/fp64 与 eager 数值一致。 -3. **参考 molpy 区分 `Potential` 与 `ForceField`**(两者都是 `nn.Module`,是不同概念,见 Domain basis)。 - -并在 `TensorDict` 与 typed state 之间取平衡:**拓扑/模型 I/O 用 `TensorDict`(模型母语),积分器边界跨越的动力学状态用 typed `MDState` NamedTuple(纯张量,pytree,编译友好)**。 - -本次重构必须**保留上一轮 review 修掉的全部正确性**:活体几何(冻结 PES bug)、ΔF fp64、Langevin dof=3N、有界显存(CPU 流式 + shard-flush)、单位桥接、正质量守卫。 - -## Domain basis - -**molpy 的 `ForceField` vs `Potential`**(`/home/jicli594/work/molcrafts/molpy`,调研确认): - -- `ForceField` 定义 styles/types/parameters(符号、可变、**不含求值核**),通过 `ff.to_potentials(frame)` **绑定到一个具体体系**,产出可求值的 `Potentials`。 -- `Potentials`(=分子力学里的求值集合)是 frame-bound、不可变、可求值:`calc_energy(coords)` / `calc_forces(coords)`;求值接口抽象为 `PotentialLike` Protocol(`optimize/base.py`)。 - -映射到 molnex MD(ML 势场语境): - -| molpy | molnex MD(本 spec) | 职责 | -|---|---|---| -| `Potential`(功能形式/模型) | **`Potential(nn.Module)`** = `molpot.BasePotential` / `PiNetPotential`(已存在,**复用不重写**) | 从 batch `TensorDict` 算能量;力经 `ForceDerivation(method="functorch")` 求 | -| `Potentials`(frame-bound 可求值) | **`ForceField(nn.Module)`**(新增) | 把一个 `Potential` **绑定到体系模板**,对 `pos` 暴露 `forward(pos)->ForceOutput` + `calc_energy/calc_forces` | -| `PotentialLike` Protocol | `ForceField` 的 `calc_energy/calc_forces` | 积分器消费的统一求值接口 | -| `Frame`/`Block`(列式状态) | `MDState` NamedTuple(typed) + `ForceField` 内部持有的 `TensorDict` 模板 | tensordict↔typed 平衡点 | - -**结论**:`Potential` 是"功能形式(能量模型)",`ForceField` 是"绑定体系后的能量+力求值组件"。积分器**只认 `ForceField`**,不认 `Potential`、不认闭包。 - -**可编译性事实(调研确认,决定本设计可行)**:`PiNetPotential` 用 `ForceDerivation(method="functorch")` → `torch.func.grad` **traced into the forward graph**,`energy→force` 单次 backward,`torch.compile(fullgraph=True)` **0 graph breaks**(证据:`tests/test_molzoo/test_pinet_functorch_force.py`、`docs/molix/explanation/throughput-and-compilation.md:21-31,94-99`,GH200 ~130 steps/s);fp32/fp64 均正确(`benchmarks/verify_pinet_cudagraph_ef.py`);力路径**不改输入 batch**(全在内部 clone,`pinet.py:443-469`)→ 跨步复用持久模板安全。编译入口 `molix.compile.maybe_compile(cuda_graphs=True)` = `CUDA_GRAPH_PRESET{fullgraph,dynamic=False,reduce-overhead}`,要求静态 shape(冻结邻居表天然满足)。 - -**积分器物理(不变)**:BAOAB(Leimkuhler & Matthews 2013, DOI 10.1093/amrx/abs010);O 步 `v ← c1·v + c2·σ·ξ`,`c1=e^{-γΔt}`,`c2=√(1-c1²)`,`σ=√(k_BT/m)`;γ=0 ⇒ c1=1,c2=0 ⇒ O 步恒等 ⇒ 退化为 velocity-Verlet(**这正是"always-on O 步去分支"成立的依据**:0·σ·ξ=0,数值不变)。 - -## Design - -### 组件清单(全部 `nn.Module` 或 typed pytree) - -**`molix/md/types.py`(新)** —— 共享 typed 容器,避免循环依赖: -- `class ForceOutput(NamedTuple): energy: Tensor; forces: Tensor` —— 取代 `(E,F)` 裸元组与 `ForceFn` Callable 别名。 -- `class MDState(NamedTuple): pos: Tensor; vel: Tensor; force: Tensor` —— 积分器跨步的动力学状态;NamedTuple 是 pytree,`torch.compile` 原生支持。 - -**`molix/md/forcefield.py`(新,取代 `force_seam.py`)** —— `ForceField` 组件层: -- `class ForceField(nn.Module)`[抽象]:`forward(self, pos: Tensor) -> ForceOutput`;便捷 `calc_energy(pos)->Tensor` / `calc_forces(pos)->Tensor`(镜像 molpy `PotentialLike`)。 -- `class PotentialForceField(ForceField)`:绑定一个 `Potential`(molpot 模型)+ 体系模板。`__init__` 里 `clone()` 模板一次、`del` 掉 `edges.bond_diff/bond_dist`(活体几何,**保留冻结-PES 修复**)、把 `pos` 之外的拓扑作为成员持有;`energy_scale: float` buffer 做单位桥接;`forward` 写入 `pos`(`.to(device,dtype)`)→ `model(batch, compute_forces=True)` → `ForceOutput(energy*scale, forces*scale)`。**无闭包、无每步 clone(模型内部自 clone)**。 -- `class AnalyticForceField(ForceField)` 及具体 `HarmonicForceField(ForceField)`(`E=½k‖x‖²`, `F=-kx`)—— 取代测试与参考用的 `_harmonic(k)` 闭包,使积分器单测不依赖 PiNet 且走同一组件接口。 -- `class LennardJonesForceField(ForceField)`(全对 LJ,`E=4ε[(σ/r)¹²−(σ/r)⁶]`,无 cutoff/无邻居表,**全对天然规避冻结邻居表问题**)—— 用于 NVE 长程能量守恒验收(ac-008):小 LJ 团簇(N~13–55,平衡构型附近)是各向异性、刚硬、非谐的真实 PES,比谐振子玩具强得多的积分器稳定性检验。力解析可微 → `rollout` 全图编译。 - -**`molix/md/integrators.py`(重构)** —— `Integrator` 组件层: -- `class Integrator(nn.Module)`[抽象]:持有 `self.force: ForceField`(组件,非 Callable);`initial(self, pos, vel) -> MDState`(seed 力);`step(self, state: MDState, noise: Tensor) -> MDState`;`rollout(self, state: MDState, n_steps: int) -> MDState`。 -- `class LangevinVerletIntegrator(Integrator)`:`register_buffer` 存 `dt/gamma/kbt/c1/c2/mass_col/inv_mass/sigma`(`mass` 在 `__init__` 一律归一化为正张量 buffer,**去 isinstance**,**正质量守卫**);`step` 走 BAOAB,**O 步恒等式始终执行**(去 `if gamma>0`),`noise` 永远是真实张量(去 `None` 分支);`dof` 约定 γ>0→3N、γ=0→3N-3(**保留 dof 修复**,作为温度估计辅助,放在 runner)。 -- **编译**:`step`/`rollout` 对 traceable `ForceField`(含 PiNet functorch)`torch.compile(fullgraph=True)` 零 break。`rollout` 内部逐步 `torch.randn`(dynamo functionalize RNG;`manual_seed` 复现)→ 满足"`torch.compile(md.run)` 全流程编译"的无 hook 快路径。生产路径见 runner。 - -**`molix/md/runner.py`(重构)** —— 驱动 + 观测: -- `MDRunner`:持有 `Integrator` 组件;`run(state0|pos,vel, n_steps)` 用 **eager Python 循环 + 已编译 `step`**(每步 eager 抽 `noise` 传入已编译 step,hook 在此触发——hook 有副作用不能进编译图);可选 `compile=True` 经 `maybe_compile` 包 `step`。**hook 分发改类型化直调**(`BaseHook` 已有 no-op 默认实现,`hook.on_train_batch_end(...)` 直接调,**删 `getattr`**)。物理量仍只走 `outputs` 通道(保 state 命名空间契约)。 -- `TrajectoryHook`:消费 `MDState`/`ForceOutput`;**保留 shard-flush 有界显存 + 诚实同步拷贝**;`weights_only=True` 重载 shard。 -- 共享质量列工具 `as_mass_col` 仍由 integrators 提供(去 runner 重复)。 - -**`molix/md/dynamics.py`(重构)** —— study 层迁到组件:`run_trajectory`/`evaluate_delta_along_trajectory`/`build_paired_trajectory` 改用 `PotentialForceField` 对(ref/quant),不再调 `build_force_fn` 闭包;**保留 ΔF fp64 与 dof 元数据**。`TrajectoryArtifact` 不变。 -**`molix/md/ase_shim.py`(重构)** —— 基于 `ForceField` 组件而非闭包。 -**`molix/md/__init__.py`** —— 导出组件词表:`ForceField/PotentialForceField/HarmonicForceField`、`ForceOutput/MDState`、`Integrator/LangevinVerletIntegrator`、`MDRunner/TrajectoryHook`、study 层符号。 - -### tensordict ↔ typed 平衡(指令 5) - -- **`TensorDict`** 留在 `ForceField` 内部(模型 I/O 的异构拓扑:`atoms/edges/graphs`),不外泄到积分器。 -- **`MDState`/`ForceOutput`**(typed NamedTuple,纯张量)是积分器/runner/hook 之间的**唯一**数据契约。NamedTuple 是 pytree → `torch.compile`/`vmap` 友好,且静态可标注。 -- 边界函数:`ForceField.forward(pos: Tensor) -> ForceOutput`。一边是 TensorDict,一边是 typed tensor,转换点唯一且类型化。 - -## Files - -- `src/molix/md/types.py` —— 新增 `ForceOutput`、`MDState`。 -- `src/molix/md/forcefield.py` —— 新增,取代 `force_seam.py`;`ForceField`/`PotentialForceField`/`AnalyticForceField`/`HarmonicForceField`。 -- `src/molix/md/force_seam.py` —— 删除(`build_force_fn` 迁为 `PotentialForceField`)。 -- `src/molix/md/integrators.py` —— `Integrator` 基类 + `LangevinVerletIntegrator(nn.Module)`,`MDState` 步进,去分支/去 isinstance,`rollout` 可编译。 -- `src/molix/md/runner.py` —— `MDRunner` 类型化 hook 直调 + 编译 step 循环;`TrajectoryHook` 适配 `MDState`。 -- `src/molix/md/dynamics.py` —— study 层迁到 `PotentialForceField`。 -- `src/molix/md/ase_shim.py` —— 基于 `ForceField`。 -- `src/molix/md/__init__.py` —— 组件词表导出。 -- `tests/test_molix/test_md_*.py` —— 迁到组件 + 新增编译/类型测试。 -- `benchmarks/verify_md_lj_nve.py` —— 新增,LJ 团簇 NVE 5ns 能量守恒长程验证(ac-008)。 -- `CLAUDE.md` —— 更新 `molix.md` 依赖图注记(组件化 + 可编译)。 - -## Tasks - -1. `types.py`:`ForceOutput`、`MDState`(+ pytree 注册校验 round-trip 测试)。 -2. `forcefield.py`:`ForceField` 抽象 + `PotentialForceField`(含活体几何 `del bond_diff`、`energy_scale`、持久模板)+ `HarmonicForceField`。RED:迁移 `test_force_seam_tracks_live_geometry`(活体几何回归)。 -3. `integrators.py`:`Integrator`/`LangevinVerletIntegrator(nn.Module)`,`MDState` 步进,always-on O 步,mass buffer + 正质量守卫,`as_mass_col`。RED:NVE 守恒、Langevin 等分、force-caching、step==step_cached、compile==eager 全部迁移并通过。 -4. `runner.py`:`MDRunner` 类型化 hook 直调 + 编译 step;`TrajectoryHook` 适配 `MDState`,保留 shard-flush/dof/单位。RED:迁移 runner 生命周期 + shard-flush 测试。 -5. `dynamics.py`/`ase_shim.py`/`__init__.py`:study 层与 ASE 壳迁到 `ForceField`;保留 ΔF fp64、dof 元数据;迁移 paired/energy-varies 测试。 -6. **编译验收**:新增 `test_md_compile.py`:`torch.compile(ig.rollout, fullgraph=True)` 对真实 tiny-PiNet 跑 N 步 0 graph-break 且 ≈ eager(fp32 与 fp64);`torch.compile(ig.step)` 同理。 -7. **LJ NVE 5ns 能量守恒验收**:`LennardJonesForceField` + `benchmarks/verify_md_lj_nve.py`,LJ 团簇 γ=0 经编译 rollout 跑满 5 ns,验证总能量无系统漂移(ac-008)。5ns≈数百万步,**只有靠编译/CUDA-graph 才可行**,本验收同时压测性能。 -8. **静态化验收**:热路径无 `getattr`/`isinstance`/`None`-分支(代码审查 + 可选 `ty`/`mypy` 门);`ruff check`+`ruff format` 干净。 -9. 更新 `CLAUDE.md` + specs `INDEX.md`。 - -## Testing - -- **正确性回归(必须全绿,防止重构倒退)**:活体几何(energy/force 随构型变)、能量守恒(NVE)、等分(Langevin,T 无偏 dof)、force-caching 位级一致、shard-flush==单缓冲、ΔF fp64、正质量守卫 raise。 -- **编译**:`fullgraph=True` 的 `rollout`/`step` 对 PiNet 0 break(用 `torch._dynamo` 计数或 `fullgraph` 抛错即失败);compile==eager 数值 `allclose`(fp32 atol 适配、fp64 1e-10);NVE+compile 与 Langevin+compile 均覆盖。 -- **类型/静态**:`ForceFn` Callable 别名已删;`grep` 断言热路径无 `getattr(`/`isinstance(`/`is None`(测试或 CI 检查);NamedTuple pytree round-trip。 -- **LJ NVE 5ns 长程守恒**(ac-008):`benchmarks/verify_md_lj_nve.py` —— LJ 团簇(记录 N/σ/ε/m/dt/T0)γ=0 编译 rollout 跑满 5 ns(步数=5ns/dt);判据:E_tot(t) 线性拟合斜率 `|slope·5ns|/|E_tot(0)| < 1e-3`,RMS 涨落有界、无单调爬升,温度有限不爆。打印 drift/斜率 PASS/FAIL。属长程验证脚本(非快速单测)。 -- 环境:无 editable 安装,`PYTHONPATH=src:. /nobackup/proj/disk/teoroo/personal/jicli594/work/.x86_64/bin/python -m pytest tests/test_molix/test_md_*.py`。 - -## Out of scope - -- **PBC / 最小镜像 / 邻居表重建**:仍是开放体系、冻结邻居表、小位移研究引擎(在 `ForceField` 文档与 `__init__` 明示);活体 minimum-image 与 Verlet skin 另立 spec。 -- **把 study 层(`TrajectoryArtifact`/`build_paired_trajectory`)迁出到 `pinet-quant/csmd`**:API review 提的跨包搬迁,独立处理,本 spec 只做组件化与类型化。 -- **MACE**:走 `ForceDerivation(method="autograd")`(`allow_in_graph` 路径),其全图编译与 cuet 约束(pytorch#170834)是另一条线。 -- **CUDA-graph/`reduce-overhead` 下的 RNG 确定性**:`rollout` 默认 in-graph `torch.randn`;CUDA-graph 路径的预抽噪声缓冲作为后续优化。 diff --git a/.devcontainer/cpu/devcontainer.json b/.devcontainer/cpu/devcontainer.json deleted file mode 100644 index 37970f5..0000000 --- a/.devcontainer/cpu/devcontainer.json +++ /dev/null @@ -1,11 +0,0 @@ -// For format details, see https://aka.ms/devcontainer.json. For config options, see the -// README at: https://github.com/devcontainers/templates/tree/main/src/anaconda -{ - "name": "molnex[cpu]", - "image": "mcr.microsoft.com/devcontainers/base:ubuntu-22.04", - "features": { - "ghcr.io/molcrafts/features/molvis:latest": {} - }, - "mounts": [], - "remoteUser": "root" -} diff --git a/.devcontainer/cuda/devcontainer.json b/.devcontainer/cuda/devcontainer.json deleted file mode 100644 index b4e76b3..0000000 --- a/.devcontainer/cuda/devcontainer.json +++ /dev/null @@ -1,19 +0,0 @@ -// For format details, see https://aka.ms/devcontainer.json. For config options, see the -// README at: https://github.com/devcontainers/templates/tree/main/src/anaconda -{ - "name": "molnex[cuda]", - "image": "mcr.microsoft.com/devcontainers/base:ubuntu-22.04", - "features": { - "ghcr.io/molcrafts/features/molvis:latest": { - "computeBackend": "cuda", - "cudaVersion": "12.1" - } - }, - "runArgs": [ - "--runtime", - "nvidia", - "--gpus", - "all" - ], - "remoteUser": "root" -} diff --git a/.github/assets/moko.svg b/.github/assets/moko.svg index 2d0d688..e37dbcb 100644 --- a/.github/assets/moko.svg +++ b/.github/assets/moko.svg @@ -1 +1 @@ - \ No newline at end of file + diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index a62be6e..35160e6 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -50,5 +50,5 @@ jobs: pip install -e ".[dev]" --no-build-isolation -v timeout-minutes: 20 - name: Run tests - run: pytest tests/ -q -p no:cacheprovider + run: pytest tests/ -q -p no:cacheprovider -n auto --dist worksteal continue-on-error: true diff --git a/.gitignore b/.gitignore index 9f18139..6181e6f 100644 --- a/.gitignore +++ b/.gitignore @@ -132,3 +132,9 @@ molnex-bench-*.out # Benchmark scratch outputs (paper artifacts live in pinet-quant) benchmarks/results/ + +# machine-local benchmark/probe run logs and arch-tagged op builds +benchmarks/*.err +benchmarks/*.out +src/molix/op/build-*/ +.claude/scheduled_tasks.lock diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index ef26766..f13f4fd 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -31,7 +31,7 @@ repos: - id: ci-test name: "CI test (same as ci.yml)" - entry: bash -c 'pytest tests/ -q -p no:cacheprovider || true' + entry: bash -c 'pytest tests/ -q -p no:cacheprovider -n auto --dist worksteal || true' language: system pass_filenames: false always_run: true diff --git a/CHANGELOG.md b/CHANGELOG.md index 769602f..bf5465f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -17,7 +17,7 @@ models. **molix — training & execution** - `torch.compile` / CUDA-graph capture and AOT-Inductor model export (`molix.compile`, `molix.export`). -- In-process Langevin velocity-Verlet MD driver (`molix.md`) with an ASE shim. +- In-process Langevin velocity-Verlet MD driver (`molix.md`). - Trajectory diagnostics and a thermal-noise verdict (`molix.analysis`), plus weight-quantization tooling with paired force-Δ and effective-temperature scalars (`molix.quant`). @@ -46,6 +46,66 @@ models. - GitHub Actions CI + pre-commit (ruff / ty), channel-based logging via `mollog`. ### Changed + +- **Breaking (`molix.md`): the neighbour-rebuild cadence has exactly one owner, the list.** `MD(rebuild_every=)` and `molix.md.NeighborListHook` are **removed** (hard removal, no shim — `stage: experimental`), together with `Integrator.rebuild_every` / `_force_eval_count`. The cadence is now configured where it belongs, on `NeighborList(skin=, every=, delay=, check=)`, and `Integrator.eval_force` asks that policy once per force evaluation, at the positions being evaluated — the step-start hook refreshed the list one displacement behind the positions entering `F = -∇E`, a systematic NVE energy leak. Migration: `MD(ff, ..., rebuild_every=1)` → `NeighborList(..., skin=0.0, every=1, delay=0, check=True)` (the accurate no-skin limit; `skin > 0` is the production setting), and a frozen list for a `fullgraph=True` rollout is `MD(ff, ..., integrator=LangevinVerletIntegrator(ff, ..., rebuild=False))`. `ForceField` gains the read-only `rebuilds_neighbors` capability property (`True` on `LennardJonesCutForceField` / `PeriodicPotentialForceField`, per-instance on `CallableForceField`) that the integrator derives its static `rebuild` bool from; `ForceField.rebuild_neighbors(pos)` now *runs the list's policy* (`neighbors.update(pos)`, which may decline) instead of forcing a build — `neighbors.rebuild(pos)` remains the forced-build escape hatch. +- `molix.md.NeighborList` gains a TensorDict bind surface: `build(batch)` writes the live `edges.edge_index` / `edges.shifts` buffers into the batch by reference (single owner — `PeriodicPotentialForceField._bind_neighbors` is gone) and `update()` now accepts a batch or a raw positions tensor. +- MACE family performance + convention fixes (review-driven). + `MACEMatpes(use_fallback=)` now defaults to `False` (fused cuEq kernels — + measured ~36x faster per MD step; forces are always autograd there, so the + functorch reason for the fallback never applies; pass `True` on CPU), and + `MACE` / `ProductHead` expose `use_fallback` instead of hardwiring the + pure-torch path. The element-table check moved off the per-step path + (`MACEMatpes.validate_elements`, run once per instance). The MACE-side + `(2, E)` edge layout was removed everywhere — `DensityInteraction` / + `ResidualInteraction` / `ZBLRepulsion` and both model cores now take the + repo-wide `(E, 2)` `[:, 0]`=source convention, eliminating the per-step + `.t().contiguous()` copy and restoring the `bond_index` anti-alias guard. + `scatter_sum_compile_safe` defaults to `index_add_`; the bit-exact one-hot + GEMM (2.4–3.3x slower at every profiled shape) is opt-in via + `MOLNEX_SCATTER_ONEHOT=1`. `MACEOMol` gained the `num_interactions` guard + and the strict state-dict loader already written for MatPES (raises on + unfilled learnables/shape mismatches instead of silently dropping weights), + plus an `edge_channels` parameter replacing hardcoded `128x{l}` tables. + New `cueq-cu12` / `cueq-cu13` extras declare the fused-kernel wheel; + `run_nve.py` now reports actual fused-kernel availability, not the request. +- `molix.md` public contract reworked. `MD` is the single documented entry + point; `MD(dtype=)` now governs the **MD side only** (state, integrator + constants, mass) with the potential's precision set independently via + `MD.set_potential_dtype` — the integrator casts force output back to the + state dtype at the boundary. `MD(integrator=)` accepts any constructed + `Integrator` (the ABC now declares `advance` / `advance_n` / `rollout` / + `removed_dof`). `MDRunner` speaks its own `MDHook` protocol + (`on_run_start` / `on_step_start` / `on_step_end` / `on_run_end` with typed + `MDObservables`) instead of impersonating `Trainer` hooks, and returns a + typed `MDState`. Renames: `MDState.force` → `forces`, `molix.md`'s + `CheckpointHook` → `MDCheckpointHook`; velocity sampling moved off `MD` to + `MaxwellBoltzmann`; `LangevinVerletIntegrator.run` and the migrated-out + `dynamics` study layer (`run_trajectory`, `build_paired_trajectory`, + `TrajectoryArtifact`) were removed (they live in the `pinet-quant`/`csmd` + project). +- `PotentialForceField`'s potential contract is monomorphic: `forward(td)` + writes `graphs.energy` / `atoms.forces` per `molix.schema`; force derivation + is fixed at potential construction (`compute_forces=True`), no longer a + per-call flag, and `calc_energy` is no longer cheaper than a full + evaluation. New adapters: `PeriodicPotentialForceField` (rebuilding + fixed-capacity neighbour list bound by reference) and `CallableForceField` + (any `pos -> (energy, forces)` callable, e.g. an AOTI `.pt2`). +- `PeriodicNeighborList.edge_index` is now `(capacity, 2)` per the repo-wide + edge convention (was `(2, capacity)`); `to()` accepts positional + device/dtype. The batch-schema keys (`ENERGY_KEY`, `FORCES_KEY`, …) moved to + the new `molix.schema` (re-exported by `molpot.derivation.protocol`), and + physical constants (`KB_EV_PER_K`, `EV_PER_AMU_A2_FS2`, `KB_AMU_A_FS`) to + the new `molix.units` — both breaking the latent `molix → molpot` import + cycle. +- The MD-side neighbour list is renamed `PeriodicNeighborList` → `NeighborList` + (`molix.md.NeighborList`). Every call site moved in the same commit and there + is **no back-compat alias** (`stage: experimental`) — update imports. The bare + name is now shared with the data-pipeline `SampleTask` + `molix.data.tasks.neighbor.NeighborList`, deliberately: one is a stateful + per-run fixed-capacity buffer owner, the other a stateless pipeline transform, + and each is the natural name in its own layer. `molix.md.neighbors` imports the + task as `NeighborListTask` so the class cannot shadow its own dependency. + Behaviour, signature and defaults are unchanged. - On-disk cache is now the single-file `PackedCache` (`.pt` with packed per-atom/edge/graph buckets, `mmap` loads) — replaces per-sample memmap dirs to stay within HPC inode budgets. diff --git a/CLAUDE.md b/CLAUDE.md index 838b45c..75b3b67 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -2,10 +2,11 @@ mol_project: name: molnex language: python + stage: experimental build: install: "pip install -e '.[dev]'" - check: "ruff check src/ && ruff format --check src/" - test: "python -m pytest tests/ -v" + check: "ruff check src/ tests/ scripts/ regressions/ && ruff format --check src/ tests/ scripts/ regressions/ && ty check src/ --exit-zero-on-warning" + test: "python -m pytest tests/ -q -n 12 --dist worksteal" test_single: "python -m pytest {path} -v" coverage: "python -m pytest tests/ --cov=src --cov-report=term-missing" arch: @@ -23,6 +24,116 @@ mol_project: This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository. + + + +## What this repo is + +MolNex is a dict-first molecular ML framework for unified modeling of +molecular potentials and properties with physics-aware ML. Four packages: +molix (training), molrep (representation), molpot (potentials), molzoo +(encoder recipes). Python + PyTorch, with optional C++/CUDA engine +interfaces under `interface/`. + +## Where things live + +- Source code: `src/` +- Tests: `tests/` (mirrors source: `src/foo/boo.py` → `tests/test_foo/test_boo.py`) +- Public documentation: `docs/` +- Passive project knowledge (notes, decisions, debt, blueprint): `.claude/notes/` +- Active runtime specs (alive, deleted on completion): `.claude/specs/` +- Claude Code runtime config (agents, skills, hooks, settings): + `.claude/agents/`, `.claude/skills/`, `.claude/hooks/`, `.claude/settings.json` + +## Design preferences (default) + +**Default for all MolCrafts projects.** Apply unless the operator +**explicitly** requires a functional (or other) style for a named +subsystem — then capture the exception with `/mol:note` and scope it. +Do **not** invent a functional style on your own. + +### Iron law — no silent debt (all projects) + +Discover anti-pattern / failing test / broken invariant / clear bug +in the surface you touch or depend on → **prioritize or hard-stop**: + +1. **Do not ignore** ("pre-existing, leave it"), skip-mark, weaken + asserts, or land features on known rot. +2. **Fix now** if local + stage-allowed; else **stop**, report + path:line, route `/mol:fix` / `/mol:refactor` / supersede. +3. **Name it** in the summary (found / fixed / blocking). Silence = process failure. + +Outranks "stay in scope" / "minimal diff" when those mean knowingly +leaving rot you already saw. + +### Prefer + +- **OOP by default.** Domain concepts are types with methods + (`NeighborList.build`, `ForceField.energy`), not free-floating + helpers. Module-level functions only for true free operations (pure + math with no natural owner) or thin package re-exports. +- **Primitive, single-responsibility public APIs.** Callers compose: + construct → configure → one concern → read result. Each public + method does one named thing. +- **Inline until the second real use.** A helper used in exactly one + place stays inline (or a private method on the owning type). Extract + only at a second call site, or when a unit test must target that unit. + +### Forbid + +- **Factory functions as the primary constructor story.** No + `make_foo` / `build_bar` / `create_*` wrappers around construction. + Prefer `Foo(...)`. Explicit alternate constructors only when they + have distinct semantics (`Foo.from_file`, `Foo.empty`) — not + `make_foo` aliases of `__init__`. +- **God data structures.** No mega-dict / mega-struct / ambient + "context" blob every layer reaches into. Pass the few fields a call + needs, or a narrow typed view. Split types that accumulate more + than one coherent responsibility. +- **All-in-one façade APIs.** No public `run_everything` / + `compute_all` / `pipeline` that hides multi-step work. Composition + is the **caller's** job (scripts, docs examples, `regressions/`). + The library exposes primitives only. + +### Shape check (before adding a public symbol) + +1. Natural owning type? → method on that type, not a free function. +2. More than one user-visible step? → split into primitives. +3. Only one in-tree call site? → do not extract. +4. Tempted to hang another field on a "context" bag? → new parameter + or smaller type instead. + +### Tests (default) + +- Unit tests **only** under `tests/`, path mirrors source + (`src/foo/boo.py` → `tests/test_foo/test_boo.py`), types mirror + (`FooClass` → `TestFooClass`). Single-function tests — no e2e under + `tests/`. Public-API scenarios → `regressions/` with **hard-coded** + goldens (no live third-party oracles). Details: `tester` agent. + +## Default workflow + +For non-trivial work, prefer: +1. plan (`/mol:spec` or free-form → discuss / grill) +2. implement (`/mol:impl` or `/mol:fix`) +3. review (`/mol:review`) +4. capture decisions (`/mol:note` — harness sync, not append-only) + +## What must never change casually + +- Post-collate batch schema (`atoms` / `edges` / `graphs` / `bonds`) +- Edge convention: `edge_index[:,0]=source`, `edge_diff = pos[target]-pos[source]` +- One-way package dependency: `molix` ← `molrep` ← `molzoo` / `molpot` +- molpy-only imports under `src/` and `tests/` (never bare `molrs`) +- `TrainState` nested-namespace write contract (no slash-key writes) +- `PackedCache` single-file layout (do not migrate to `TensorDict.memmap_`) + + + + + ## Project Overview **MolNex** (v0.1.0) is a dict-first molecular ML framework for unified modeling of molecular potentials and properties with physics-aware ML. It is composed of four packages: @@ -45,6 +156,14 @@ UnitsError / UnitSystem / MolRec / etc. Do **not** `import molrs` or `from molrs import …` — molrs is the Rust core behind molpy, not a molnex dependency surface. +**Allowed third-party surface (hard).** Under `src/`, `tests/`, and in-repo +`scripts/`: MolCrafts packages (`molpy`, `mollog`, `molcfg`, …), PyTorch / +TensorDict / torch ecosystem already in `pyproject.toml`, **numpy**, and +**cuEquivariance** (`cuequivariance` / `cuequivariance_torch`). **No ASE, no +e3nn, no mace-torch** (or other chemistry-ML wrappers). Extxyz I/O uses +`molix.datasets._extxyz`; equivariance uses cue only. Upstream oracles live +out-of-tree if ever needed — do not re-introduce those imports here. + Source ↔ unit-test path mirror: ``` @@ -127,12 +246,20 @@ TensorDict (batch_size=[]) ├── "edges": TensorDict (batch_size=[E]) │ ├── edge_index: source-target pairs (E, 2) # [:,0]=source, [:,1]=target │ ├── edge_diff: edge vectors (E, 3) # pos[target] - pos[source] -│ └── edge_dist: edge distances (E,) +│ ├── edge_dist: edge distances (E,) +│ └── shifts: periodic remainders (E, 3) [optional] # pos[t]-pos[s]+shift = min-image vector, Å └── "graphs": TensorDict (batch_size=[B]) [optional] ├── num_atoms: (B,) + ├── cell: cell vectors (B, 3, 3) [optional] # Å; read by the sonata stress path └── ``` +`("edges", "shifts")` is written by the MD bind path (`molix.md.NeighborList.build`) +and read by `molzoo.mace` potentials; `("graphs", "cell")` is read by +`molpot.composition.sonata` (stress). On the MD bind path the `edges` +namespace has `batch_size=[capacity]` with live edges in `[0, num_edges)` +and dead padding beyond — the list owns `edges` once bound. + Access: `batch["atoms", "Z"]`, `batch["edges", "edge_dist"]`. Encoders mutate the batch in place, writing `node_features` under `atoms` and `edge_features` under `edges` — no subclass swap. @@ -166,6 +293,12 @@ Pass `symmetry=False` to get only the upper-triangle half-pairs (`E = n_pairs`) want to exploit Newton's-3rd-law symmetry. The two modes produce different `task_id`s so pipeline caches are kept separate. +`edge_index` is `(E, 2)` **everywhere** — including every MACE-family block +(`DensityInteraction`, `ResidualInteraction`, `ZBLRepulsion`, both model cores). +Upstream MACE's `(2, E)` layout is transposed away at the checkpoint-port +boundary and never crosses a module seam; unpack as `source, target` +(never `sender, receiver`). `(2, N)` is reserved for `bond_index`. + **Why edge_diff = pos[target] − pos[source]?** This makes the displacement vector point in the same direction as the edge (source → target), which is the convention expected by `SphericalHarmonics` and all `cuEquivariance`-based tensor products in this repo. The C++ `getNeighborPairs` kernel @@ -217,7 +350,7 @@ molrep.embedding → molrep.interaction → molrep.readout molzoo (MACE, Allegro, PiNet encoders) ──→ molpot.heads ↓ ↓ molpot.composition (PotentialComposer, molpot.potentials - Sonata, build_sonata) ↓ + Sonata, Sonata.from_encoder) ↓ ↓ molpot.derivation (EnergyAggregation, molix.core (Trainer, TrainState, Step, Hook) ForceDerivation) ↓ @@ -230,7 +363,9 @@ molix.engine (EngineAdapter/EngineForward/StaticForward, export_for_lammps) ─ Notes on cross-package edges (verified against imports): - `molzoo` consumes `molrep.readout`/`molrep.interaction` and `molpot.derivation` - (PiNet's energy/force head is co-located with the encoder by design). + (PiNet's energy/force head is co-located with the encoder by design; + `MACEMatpes` / `MACEOMol` are full energy/force checkpoint ports under the + same scoped exception — pending the mace-subpackage restructure spec). - `molpot.heads` imports `molrep.embedding` (e.g. `heads/edge.py`) — the arrow runs heads→embedding, **not** readout→heads. - `molrep.heads` (`ScalarHead`, …) is a distinct sub-tree from `molpot.heads`. @@ -482,3 +617,12 @@ One skill + one agent. Both repo-local under `.claude/skills/` and `.claude/agen 5. `molzoo-auditor` MUST print ≥ 1 verdict per invocation (even `✅ confirmed, no drift`), citing the triggering `run_id` or question + paper section + spec row + code file:line. ⚠️ (code-drift) is print-only; 📝/🆚 produce file diffs in `.md`. The auditor never edits code — code changes are always the user's call. The 10-section structure of `.md` is **immutable** — adding / removing / renaming a section is a §10.2 breaking change. Invariants are enforced inside the skill, not via `settings.json` hooks, for now. + +## Learnable classical FF — placement (binding) + +See `.claude/notes/learnable-classical-ff.md`. Short form: + +1. **Reuse first** — molpy `ForceField` / `potential` / Topology enum / SMARTS before inventing molnex twins. +2. **Generalize without `Foo(method=…)`** — single-responsibility types only; no multi-backend method switches on new APIs. +3. **Non-diff sinks to molpy/molrs** — *enumeration*, SMARTS, classical non-torch E/F, FF tables under molpy (≥0.13). Imports: `molpy` only (never bare `molrs`). +4. **Batch topology lives in molix TensorDict**, not molpy Frame. Column form under namespaces, e.g. `batch["angles"]["atomi"]` / `batch["angles", "atomj"]` / `atomk` (propers/impropers add `atoml`; improper **center = atomi**). Stack to COO only at kernel call sites if a potential still wants `[arity, N]`. diff --git a/README.md b/README.md index 069cb04..6b451db 100644 --- a/README.md +++ b/README.md @@ -16,7 +16,7 @@

- Documentation  ·  + Documentation  ·  Quick start  ·  Ecosystem

@@ -62,9 +62,9 @@ becoming the framework itself. pip install molnex ``` -Requires Python >= 3.10 and PyTorch >= 2.10. The package builds native C++ ops +Requires Python >= 3.12 and PyTorch >= 2.10. The package builds native C++ ops via scikit-build-core and CMake >= 4.0; an editable install is -`pip install -e ".[dev]"`. See [Installation](https://molcrafts.github.io/molnex/installation/) +`pip install -e ".[dev]"`. See [Installation](https://docs.molcrafts.org/molnex/installation/) for the full build setup. ## Quick start @@ -85,20 +85,20 @@ state = trainer.train(datamodule, max_epochs=5) print(state["train/loss"]) ``` -See the [Molix Quick Start](https://molcrafts.github.io/molnex/molix/tutorials/quick-start/) +See the [Molix Quick Start](https://docs.molcrafts.org/molnex/molix/tutorials/quick-start/) for the runnable end-to-end version, and -[Train a Graph Model](https://molcrafts.github.io/molnex/molix/tutorials/train-a-graph-model/) +[Train a Graph Model](https://docs.molcrafts.org/molnex/molix/tutorials/train-a-graph-model/) for molecular graph batches. ## Documentation -- [Documentation home](https://molcrafts.github.io/molnex/) -- [Installation](https://molcrafts.github.io/molnex/installation/) -- [Molix](https://molcrafts.github.io/molnex/molix/) — training, hooks, data, and execution -- [MolRep](https://molcrafts.github.io/molnex/molrep/) — representation learning modules -- [MolPot](https://molcrafts.github.io/molnex/molpot/) — potential composition and physical outputs -- [MolZoo](https://molcrafts.github.io/molnex/molzoo/) — reference encoder families -- [API Reference](https://molcrafts.github.io/molnex/api/) +- [Documentation home](https://docs.molcrafts.org/molnex/) +- [Installation](https://docs.molcrafts.org/molnex/installation/) +- [Molix](https://docs.molcrafts.org/molnex/molix/) — training, hooks, data, and execution +- [MolRep](https://docs.molcrafts.org/molnex/molrep/) — representation learning modules +- [MolPot](https://docs.molcrafts.org/molnex/molpot/) — potential composition and physical outputs +- [MolZoo](https://docs.molcrafts.org/molnex/molzoo/) — reference encoder families +- [API Reference](https://docs.molcrafts.org/molnex/api/) ## MolCrafts ecosystem @@ -119,7 +119,7 @@ for molecular graph batches. ## Contributing -Contributions are welcome — see the [documentation](https://molcrafts.github.io/molnex/) +Contributions are welcome — see the [documentation](https://docs.molcrafts.org/molnex/) to get started. ## License diff --git a/benchmarks/bench_mace_matpes.py b/benchmarks/bench_mace_matpes.py new file mode 100644 index 0000000..f05e8b3 --- /dev/null +++ b/benchmarks/bench_mace_matpes.py @@ -0,0 +1,190 @@ +"""MACE-MatPES perf guard: fused-vs-fallback cuEq kernels + compiled energy core. + +The MACE family carries the repo's headline GH200 compile numbers, yet had no +benchmark of its own — which is exactly how a `use_fallback=True` default (a +measured **35.7x** per-step regression with zero correctness upside) shipped +unnoticed. This script is the guard: + + * eager ``energy_forces`` with ``use_fallback=False`` (fused) vs ``True`` + (pure-torch) — asserts the fused arm is at least ``--min-speedup`` faster + when the ``cuequivariance-ops-torch`` wheel is importable; + * the ``Compiler(cuda_graphs=True)``-compiled energy core (the + ``run_nve.py`` production path) vs eager, on the fused arm. + +Random weights (perf only, no physics); a periodic random-dense system builds +real ``(E, 2)`` edges + shifts through :class:`molix.md.NeighborList`. +Run on a GPU node: + + python benchmarks/bench_mace_matpes.py # defaults: N=192, fp32 + python benchmarks/bench_mace_matpes.py --fp64 --steps 30 +""" + +from __future__ import annotations + +import argparse +import math +import time + +import torch + +from molix import config +from molix.compile import Compiler +from molix.md import NeighborList +from molpot.derivation.force import autograd_forces_from_energy + +_Z_TABLE = [1, 6, 7, 8, 14, 26] # small table; dims below are the MatPES-class ones + +#: ~0.045 atoms/A^3 — condensed-phase-ish density, for a realistic edge count. +DENSITY = 0.045 +#: Model ``r_max`` *and* neighbour-list cutoff; the two must stay equal. +CUTOFF = 6.0 + + +def _box_length(n_atoms: int) -> float: + """Cubic box edge in Angstrom holding ``n_atoms`` at :data:`DENSITY`.""" + return float((n_atoms / DENSITY) ** (1.0 / 3.0)) + + +def _min_n_atoms() -> int: + """Smallest ``n_atoms`` whose box half-width strictly exceeds :data:`CUTOFF`. + + Minimum image requires ``cutoff <= box / 2``; below that + :class:`molix.md.NeighborList` raises. + """ + return math.floor(DENSITY * (2.0 * CUTOFF) ** 3) + 1 + + +def _build_model(use_fallback: bool): + from molzoo import MACEMatpes + + torch.manual_seed(0) + return MACEMatpes( + atomic_numbers=_Z_TABLE, + atomic_energies=torch.zeros(len(_Z_TABLE), dtype=config.ftype), + r_max=CUTOFF, + num_bessel=10, + num_polynomial_cutoff=5, + l_max=3, + num_features=128, + max_hidden_l=1, + num_interactions=2, + correlation=3, + mlp_dim=16, + use_fallback=use_fallback, + ).eval() + + +def _system(n_atoms: int, device: torch.device): + torch.manual_seed(1) + box = _box_length(n_atoms) + pos = torch.rand(n_atoms, 3, dtype=config.ftype, device=device) * box + cell = torch.eye(3, dtype=config.ftype, device=device) * box + Z = _Z_TABLE[0] + torch.zeros(n_atoms, dtype=torch.long, device=device) + Z[::3] = _Z_TABLE[2] + Z[::5] = _Z_TABLE[3] + neighbors = NeighborList(cell=cell, cutoff=CUTOFF, positions=pos) + batch = torch.zeros(n_atoms, dtype=torch.long, device=device) + return pos, Z, batch, neighbors + + +def _time(fn, steps: int, warmup: int = 5) -> float: + for _ in range(warmup): + fn() + if torch.cuda.is_available(): + torch.cuda.synchronize() + t0 = time.perf_counter() + for _ in range(steps): + fn() + if torch.cuda.is_available(): + torch.cuda.synchronize() + return (time.perf_counter() - t0) / steps * 1e3 + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument( + "--n-atoms", + type=int, + default=192, + help=f"atoms in the cubic box; must exceed {_min_n_atoms() - 1} so the box " + f"half-width stays above the {CUTOFF} A cutoff", + ) + ap.add_argument("--steps", type=int, default=20) + ap.add_argument("--fp64", action="store_true") + ap.add_argument( + "--min-speedup", + type=float, + default=3.0, + help="required fused/fallback per-step ratio when the ops wheel is present", + ) + args = ap.parse_args() + + # Fail here, not 40 frames deep inside NeighborList: minimum image + # needs cutoff <= box/2, and the box is derived from --n-atoms at DENSITY. + box = _box_length(args.n_atoms) + if box / 2.0 <= CUTOFF: + ap.error( + f"--n-atoms {args.n_atoms} gives a {box:.2f} A box at {DENSITY} atoms/A^3, " + f"whose half-width {box / 2.0:.2f} A does not exceed the {CUTOFF} A cutoff; " + f"the minimum-image neighbour list would miss periodic images. " + f"Use --n-atoms {_min_n_atoms()} or more." + ) + + config.set_precision("fp64" if args.fp64 else "fp32") + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + try: + import cuequivariance_ops_torch # noqa: F401 + + fused_available = True + except ImportError: + fused_available = False + print(f"device={device} dtype={config.ftype} fused_ops_available={fused_available}") + + pos, Z, batch, neighbors = _system(args.n_atoms, device) + print(f"system: N={args.n_atoms} E={neighbors.num_edges} (capacity {neighbors.capacity})") + + times: dict[str, float] = {} + for label, use_fallback in (("fused", False), ("fallback", True)): + model = _build_model(use_fallback).to(device) + + def step(model=model): + out = model.energy_forces( + pos, Z, neighbors.edge_index, batch, num_graphs=1, shifts=neighbors.shifts + ) + return out["forces"] + + times[label] = _time(step, args.steps) + print(f"eager energy+forces [{label:8s}]: {times[label]:9.3f} ms/step") + + ratio = times["fallback"] / times["fused"] + print(f"fallback/fused ratio: {ratio:.2f}x (guard: >= {args.min_speedup} when fused available)") + + # Compiled energy core on the fused arm — the run_nve.py production path. + if device.type == "cuda": + model = _build_model(False).to(device) + energy_fn = Compiler(cuda_graphs=True)(model.energy_core) + + def compiled_step(): + leaf = pos.detach().requires_grad_(True) + with torch.enable_grad(): + energy = energy_fn(leaf, Z, neighbors.edge_index, batch, 1, neighbors.shifts) + return autograd_forces_from_energy(energy, leaf) + + t = _time(compiled_step, args.steps) + print(f"compiled energy core [fused ]: {t:9.3f} ms/step ({times['fused'] / t:.2f}x eager)") + + ok = (not fused_available) or ratio >= args.min_speedup + if not ok: + print( + "RESULT: FAIL — fused kernels available but the speedup collapsed; " + "check use_fallback plumbing / the ops wheel / cuEq versions" + ) + else: + print("RESULT: PASS") + return 0 if ok else 1 + + +if __name__ == "__main__": + import sys + + sys.exit(main()) diff --git a/benchmarks/bench_pinet.py b/benchmarks/bench_pinet.py index 798e9b4..95218cc 100644 --- a/benchmarks/bench_pinet.py +++ b/benchmarks/bench_pinet.py @@ -148,9 +148,7 @@ def gpu_matrix(n_steps: int) -> None: { "compiler": cname, "precision": precision, - "wall_ms": f"{res.wall_ms_per_step:.3f}" - if res.wall_ms_per_step - else "n/a", + "wall_ms": f"{res.wall_ms_per_step:.3f}" if res.wall_ms_per_step else "n/a", "atoms/s": f"{res.throughput_atoms_per_sec:,.0f}", "lb%": f"{res.launch_bound_pct:.0f}" if res.launch_bound_pct else "n/a", } diff --git a/benchmarks/bench_pinet_force_compile.py b/benchmarks/bench_pinet_force_compile.py new file mode 100644 index 0000000..3a6ec92 --- /dev/null +++ b/benchmarks/bench_pinet_force_compile.py @@ -0,0 +1,398 @@ +"""PiNet force path: eager vs torch.compile + peak GPU memory. + +Regimes: + +* energy-only (baseline compile surface) +* force inf / force train — method=func and method=grad +* eager vs compile(fullgraph=False) vs compile(fullgraph=True) +* optional reduce-overhead (CUDA graphs; needs static shapes — same MockBatch size) + +Reports ms/step and peak allocated / reserved MiB (torch.cuda memory stats). + +Usage (GPU node):: + + PYTHONPATH=src python benchmarks/bench_pinet_force_compile.py +""" + +from __future__ import annotations + +import argparse +import gc +import time +from collections.abc import Callable +from dataclasses import dataclass + +import torch + +_orig_load_library = torch.ops.load_library + + +def _soft_load_library(path): # noqa: ANN001 + try: + return _orig_load_library(path) + except OSError as exc: + print(f"[bench] skip native op library ({path}): {exc}") + + +torch.ops.load_library = _soft_load_library # type: ignore[method-assign] + +from molix.profiler import MockBatch # noqa: E402 +from molzoo.pinet import PiNetPotential # noqa: E402 + +ATOM_TYPES = list(range(1, 8)) + + +@dataclass +class Row: + name: str + ms: float + peak_alloc_mib: float + peak_reserved_mib: float + status: str = "ok" + notes: str = "" + + +def _sync() -> None: + if torch.cuda.is_available(): + torch.cuda.synchronize() + + +def _reset_mem() -> None: + if not torch.cuda.is_available(): + return + gc.collect() + torch.cuda.empty_cache() + torch.cuda.reset_peak_memory_stats() + torch.cuda.synchronize() + + +def _peak_mib() -> tuple[float, float]: + if not torch.cuda.is_available(): + return float("nan"), float("nan") + alloc = torch.cuda.max_memory_allocated() / (1024**2) + reserved = torch.cuda.max_memory_reserved() / (1024**2) + return alloc, reserved + + +def build( + *, + compute_forces: bool, + method: str = "func", + rank: int = 3, + depth: int = 3, + hidden: int = 64, +) -> PiNetPotential: + return PiNetPotential( + atom_types=ATOM_TYPES, + r_max=5.0, + n_basis=5, + pp_nodes=[hidden, hidden], + pi_nodes=[hidden, hidden], + ii_nodes=[hidden, hidden], + depth=depth, + rank=rank, + hidden_dim=hidden, + compute_forces=compute_forces, + method=method, # type: ignore[arg-type] + emit_property_features=False, + ) + + +def factory(device: str, n_atoms: int, n_edges: int, n_graphs: int) -> Callable[[], object]: + return MockBatch( + n_atoms=n_atoms, + n_edges=n_edges, + n_graphs=n_graphs, + atomic_numbers=7, + device=device, + seed=0, + ) + + +def _ef_loss(batch) -> torch.Tensor: + return ( + batch["graphs", "energy"].float().pow(2).mean() + + batch["atoms", "forces"].float().pow(2).mean() + ) + + +def _e_loss(batch) -> torch.Tensor: + return batch["graphs", "energy"].float().pow(2).mean() + + +def time_and_mem( + step: Callable[[], None], + *, + n_warmup: int, + n_steps: int, +) -> tuple[float, float, float]: + for _ in range(n_warmup): + step() + _sync() + _reset_mem() + # one step after reset so peak includes steady-state activations + step() + _sync() + t0 = time.perf_counter() + for _ in range(n_steps): + step() + _sync() + ms = (time.perf_counter() - t0) * 1000.0 / n_steps + peak_a, peak_r = _peak_mib() + return ms, peak_a, peak_r + + +def run_cell( + name: str, + *, + compute_forces: bool, + method: str, + train: bool, + compile_mode: str | None, + fac: Callable[[], object], + device: torch.device, + n_warmup: int, + n_steps: int, +) -> Row: + """compile_mode: None | 'default' | 'fullgraph' | 'reduce-overhead'.""" + torch.compiler.reset() + try: + pot = build(compute_forces=compute_forces, method=method).to(device) + pot.train(train) + + notes = f"method={method}" + if compile_mode is not None: + kw: dict = {"backend": "inductor"} + if compile_mode == "fullgraph": + kw.update(fullgraph=True, dynamic=False) + notes += " fullgraph=True" + elif compile_mode == "reduce-overhead": + kw.update(fullgraph=True, dynamic=False, mode="reduce-overhead") + notes += " reduce-overhead" + elif compile_mode == "default": + kw.update(fullgraph=False) + notes += " fullgraph=False" + else: + raise ValueError(compile_mode) + pot = torch.compile(pot, **kw) # type: ignore[assignment] + # warm compile + for _ in range(3): + if train: + out = pot(fac()) + if compute_forces: + _ef_loss(out).backward() + else: + _e_loss(out).backward() + pot.zero_grad(set_to_none=True) + else: + with torch.enable_grad() if compute_forces else torch.no_grad(): + pot(fac()) + _sync() + + opt = torch.optim.Adam(pot.parameters(), lr=1e-3) if train else None + + def step(): + if train: + assert opt is not None + opt.zero_grad(set_to_none=True) + out = pot(fac()) + loss = _ef_loss(out) if compute_forces else _e_loss(out) + loss.backward() + opt.step() + else: + if compute_forces: + with torch.enable_grad(): + pot(fac()) + else: + with torch.no_grad(): + pot(fac()) + + ms, pa, pr = time_and_mem(step, n_warmup=n_warmup, n_steps=n_steps) + return Row(name=name, ms=ms, peak_alloc_mib=pa, peak_reserved_mib=pr, notes=notes) + except Exception as e: # noqa: BLE001 + msg = f"{type(e).__name__}: {str(e).splitlines()[0][:80]}" + return Row( + name=name, + ms=float("nan"), + peak_alloc_mib=float("nan"), + peak_reserved_mib=float("nan"), + status="FAIL", + notes=msg, + ) + finally: + torch.compiler.reset() + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + +def print_table(rows: list[Row]) -> None: + print(f"\n{'name':<48} {'ms/step':>9} {'peak_alloc':>12} {'peak_rsrv':>10} status notes") + print("-" * 110) + for r in rows: + ms = f"{r.ms:9.3f}" if r.ms == r.ms else f"{'nan':>9}" + pa = ( + f"{r.peak_alloc_mib:10.1f} MiB" + if r.peak_alloc_mib == r.peak_alloc_mib + else f"{'nan':>12}" + ) + pr = ( + f"{r.peak_reserved_mib:8.1f} MiB" + if r.peak_reserved_mib == r.peak_reserved_mib + else f"{'nan':>10}" + ) + print(f"{r.name:<48} {ms} {pa:>12} {pr:>10} {r.status:<5} {r.notes}") + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu") + ap.add_argument("--steps", type=int, default=40) + ap.add_argument("--warmup", type=int, default=12) + ap.add_argument("--atoms", type=int, default=512) + ap.add_argument("--edges", type=int, default=4096) + ap.add_argument("--graphs", type=int, default=16) + ap.add_argument( + "--skip-reduce-overhead", + action="store_true", + help="Skip CUDA-graph reduce-overhead cells (slow first compile)", + ) + args = ap.parse_args() + + if args.device == "cuda" and not torch.cuda.is_available(): + raise SystemExit("CUDA requested but not available") + + device = torch.device(args.device) + print(f"torch {torch.__version__} device={device}") + if device.type == "cuda": + print(f"GPU: {torch.cuda.get_device_name(0)}") + props = torch.cuda.get_device_properties(0) + print(f"total VRAM: {props.total_memory / (1024**3):.1f} GiB") + + fac = factory(str(device), args.atoms, args.edges, args.graphs) + print( + f"batch atoms={args.atoms} edges={args.edges} graphs={args.graphs} " + f"steps={args.steps} warmup={args.warmup}" + ) + + rows: list[Row] = [] + + # --- energy only -------------------------------------------------------- + for cmode, tag in ( + (None, "energy-inf eager"), + ("default", "energy-inf compile"), + ("fullgraph", "energy-inf compile fullgraph"), + ): + rows.append( + run_cell( + tag, + compute_forces=False, + method="func", + train=False, + compile_mode=cmode, + fac=fac, + device=device, + n_warmup=args.warmup, + n_steps=args.steps, + ) + ) + + # --- force inference ---------------------------------------------------- + for method in ("func", "grad"): + for cmode, tag in ( + (None, f"force-inf {method} eager"), + ("default", f"force-inf {method} compile"), + ("fullgraph", f"force-inf {method} fullgraph"), + ): + rows.append( + run_cell( + tag, + compute_forces=True, + method=method, + train=False, + compile_mode=cmode, + fac=fac, + device=device, + n_warmup=args.warmup, + n_steps=args.steps, + ) + ) + + # --- force train -------------------------------------------------------- + for method in ("func", "grad"): + for cmode, tag in ( + (None, f"force-train {method} eager"), + ("default", f"force-train {method} compile"), + ("fullgraph", f"force-train {method} fullgraph"), + ): + rows.append( + run_cell( + tag, + compute_forces=True, + method=method, + train=True, + compile_mode=cmode, + fac=fac, + device=device, + n_warmup=args.warmup, + n_steps=args.steps, + ) + ) + + if not args.skip_reduce_overhead: + for method, train, tag in ( + ("func", False, "force-inf func reduce-overhead"), + ("func", True, "force-train func reduce-overhead"), + ("grad", True, "force-train grad reduce-overhead"), + ): + rows.append( + run_cell( + tag, + compute_forces=True, + method=method, + train=train, + compile_mode="reduce-overhead", + fac=fac, + device=device, + n_warmup=max(args.warmup, 15), + n_steps=args.steps, + ) + ) + + print_table(rows) + + # quick ratios vs eager force-inf func + by = {r.name: r for r in rows if r.status == "ok" and r.ms == r.ms} + base = by.get("force-inf func eager") + if base is not None: + print("\n# force-inf func vs eager") + for k in ( + "force-inf func compile", + "force-inf func fullgraph", + "force-inf func reduce-overhead", + ): + if k in by: + print( + f" {k}: {by[k].ms / base.ms:.2f}x time, " + f"alloc {by[k].peak_alloc_mib / base.peak_alloc_mib:.2f}x" + ) + + base_tr = by.get("force-train func eager") + if base_tr is not None: + print("\n# force-train func vs eager") + for k in ( + "force-train func compile", + "force-train func fullgraph", + "force-train func reduce-overhead", + ): + if k in by: + print( + f" {k}: {by[k].ms / base_tr.ms:.2f}x time, " + f"alloc {by[k].peak_alloc_mib / base_tr.peak_alloc_mib:.2f}x" + ) + + print("\n=== DONE ===") + + +if __name__ == "__main__": + main() diff --git a/benchmarks/bench_pinet_rank3_speed.py b/benchmarks/bench_pinet_rank3_speed.py new file mode 100644 index 0000000..a4a089f --- /dev/null +++ b/benchmarks/bench_pinet_rank3_speed.py @@ -0,0 +1,478 @@ +"""PiNet rank=3 train / inference speed microbench. + +Measures wall time for the real operating regimes with the current API: + +* energy-only inference (``compute_forces=False`` at init) +* force inference (``compute_forces=True`` at init, eval + enable_grad) +* force training (train, energy + force loss + ``loss.backward``) +* encode-only forward + +``forces`` / method knobs are **init-only** — ``forward(batch)`` has no +flags and no runtime branching. + +Usage:: + + PYTHONPATH=src python benchmarks/bench_pinet_rank3_speed.py --device cuda --steps 50 + PYTHONPATH=src python benchmarks/bench_pinet_rank3_speed.py --device cpu --steps 20 +""" + +from __future__ import annotations + +import argparse +import time +from collections.abc import Callable +from dataclasses import dataclass + +import torch + +# PiNet is pure-PyTorch; the arch-tagged C++ op .so may be linked against a +# different torch ABI (e.g. cpu build vs this CUDA wheel). Soft-skip so the +# microbench still runs — scatter/PME native ops are not on the PiNet path. +_orig_load_library = torch.ops.load_library + + +def _soft_load_library(path): # noqa: ANN001 + try: + return _orig_load_library(path) + except OSError as exc: + print(f"[bench] skip native op library ({path}): {exc}") + + +torch.ops.load_library = _soft_load_library # type: ignore[method-assign] + +from molix.profiler import MockBatch # noqa: E402 +from molpot.derivation import EnergyReadout, ForceReadout # noqa: E402 +from molzoo.pinet import PiNet, PiNetPotential # noqa: E402 + +ATOM_TYPES = list(range(1, 8)) + + +@dataclass +class Row: + name: str + ms: float + atoms_per_s: float + notes: str = "" + + +def _sync(device: torch.device) -> None: + if device.type == "cuda": + torch.cuda.synchronize() + + +def _time_it( + fn: Callable[[], None], + *, + device: torch.device, + n_warmup: int, + n_steps: int, +) -> float: + for _ in range(n_warmup): + fn() + _sync(device) + t0 = time.perf_counter() + for _ in range(n_steps): + fn() + _sync(device) + return (time.perf_counter() - t0) * 1000.0 / n_steps + + +def _factory( + device: str, + *, + n_atoms: int = 512, + n_edges: int = 4096, + n_graphs: int = 16, +) -> Callable[[], object]: + return MockBatch( + n_atoms=n_atoms, + n_edges=n_edges, + n_graphs=n_graphs, + atomic_numbers=7, + device=device, + seed=0, + ) + + +def _ef_loss(batch) -> torch.Tensor: + return ( + batch["graphs", "energy"].float().pow(2).mean() + + batch["atoms", "forces"].float().pow(2).mean() + ) + + +def _energy_loss(batch) -> torch.Tensor: + return batch["graphs", "energy"].float().pow(2).mean() + + +def build_potential( + *, + rank: int = 3, + depth: int = 3, + hidden: int = 64, + emit_property_features: bool | None = None, + compute_forces: bool = True, + method: str = "func", +) -> PiNetPotential: + kwargs: dict = dict( + atom_types=ATOM_TYPES, + r_max=5.0, + n_basis=5, + pp_nodes=[hidden, hidden], + pi_nodes=[hidden, hidden], + ii_nodes=[hidden, hidden], + depth=depth, + rank=rank, + hidden_dim=hidden, + compute_forces=compute_forces, + ) + if emit_property_features is not None: + kwargs["emit_property_features"] = emit_property_features + pot = PiNetPotential(**kwargs) + # Rebuild pipeline when method differs from the hard-coded default ("func"). + if method != "func": + pot.deriv_method = method + if compute_forces: + energy_ro = EnergyReadout(pot, method=method, backward=True) # type: ignore[arg-type] + force_ro = ForceReadout(pot, method=method) # type: ignore[arg-type] + + def _pipeline(batch, e=energy_ro, f=force_ro): + return f(e(batch)) + + pot._pipeline = _pipeline + else: + pot._pipeline = EnergyReadout(pot, method=method, backward=False) # type: ignore[arg-type] + return pot + + +def build_encoder(*, rank: int = 3, hidden: int = 64, emit: bool = False) -> PiNet: + return PiNet( + atom_types=ATOM_TYPES, + r_max=5.0, + n_basis=5, + pp_nodes=[hidden, hidden], + pi_nodes=[hidden, hidden], + ii_nodes=[hidden, hidden], + depth=3, + rank=rank, + emit_property_features=emit, + ) + + +def _row(name: str, ms: float, n_atoms: int, notes: str = "") -> Row: + return Row(name=name, ms=ms, atoms_per_s=n_atoms / (ms / 1000.0), notes=notes) + + +def print_table(rows: list[Row], baseline: str | None = None) -> None: + base_ms = next((r.ms for r in rows if r.name == baseline), None) if baseline else None + print(f"\n{'name':<42} {'ms/step':>10} {'atoms/s':>12} {'rel':>8} notes") + print("-" * 90) + for r in rows: + rel = f"{r.ms / base_ms:5.2f}x" if base_ms and base_ms > 0 and r.ms == r.ms else " -" + print(f"{r.name:<42} {r.ms:10.3f} {r.atoms_per_s:12,.0f} {rel:>8} {r.notes}") + + +def run_matrix( + *, + device: torch.device, + n_steps: int, + n_warmup: int, + n_atoms: int, + n_edges: int, + n_graphs: int, + skip_compile: bool = False, +) -> list[Row]: + dev = str(device) + factory = _factory(dev, n_atoms=n_atoms, n_edges=n_edges, n_graphs=n_graphs) + rows: list[Row] = [] + + # ---- encoder-only ------------------------------------------------------- + for rank, emit, tag in ( + (3, False, "enc rank3 emit=0"), + (3, True, "enc rank3 emit=1"), + (1, False, "enc rank1 emit=0"), + ): + enc = build_encoder(rank=rank, emit=emit).to(device) + enc.eval() + + def step(e=enc): + with torch.no_grad(): + e(factory()) + + ms = _time_it(step, device=device, n_warmup=n_warmup, n_steps=n_steps) + rows.append(_row(tag, ms, n_atoms)) + + # ---- energy-only inference --------------------------------------------- + pot_e = build_potential(rank=3, emit_property_features=False, compute_forces=False).to(device) + pot_e.eval() + + def energy_inf(): + with torch.no_grad(): + pot_e(factory()) + + ms = _time_it(energy_inf, device=device, n_warmup=n_warmup, n_steps=n_steps) + rows.append(_row("energy-inf rank3 (no forces)", ms, n_atoms, "baseline")) + + # compiled energy (inference only) + if skip_compile: + rows.append(_row("energy-inf + torch.compile", float("nan"), n_atoms, "skipped")) + else: + pot_ec = build_potential(rank=3, emit_property_features=False, compute_forces=False).to( + device + ) + pot_ec.eval() + try: + pot_ec = torch.compile(pot_ec, backend="inductor", fullgraph=False) + for _ in range(3): + with torch.no_grad(): + pot_ec(factory()) + _sync(device) + + def energy_compiled(): + with torch.no_grad(): + pot_ec(factory()) + + ms = _time_it( + energy_compiled, device=device, n_warmup=max(5, n_warmup), n_steps=n_steps + ) + rows.append(_row("energy-inf + torch.compile", ms, n_atoms)) + except Exception as e: # noqa: BLE001 + rows.append( + _row( + "energy-inf + torch.compile", + float("nan"), + n_atoms, + f"FAIL {type(e).__name__}", + ) + ) + finally: + torch.compiler.reset() + + # ---- force inference (eval fused path, method=func) -------------------- + for emit, tag in ((False, "force-inf eval emit=0"), (True, "force-inf eval emit=1")): + pot = build_potential( + rank=3, emit_property_features=emit, compute_forces=True, method="func" + ).to(device) + pot.eval() + assert pot.encoder.emit_property_features is emit + + def step_eval(m=pot): + m.eval() + with torch.enable_grad(): + m(factory()) + + ms = _time_it(step_eval, device=device, n_warmup=n_warmup, n_steps=n_steps) + rows.append(_row(tag, ms, n_atoms, "func has_aux 1-pass")) + + # rank-1 force inf for ratio + pot_r1 = build_potential( + rank=1, emit_property_features=False, compute_forces=True, method="func" + ).to(device) + pot_r1.eval() + + def force_r1(): + with torch.enable_grad(): + pot_r1(factory()) + + ms = _time_it(force_r1, device=device, n_warmup=n_warmup, n_steps=n_steps) + rows.append(_row("force-inf eval rank1", ms, n_atoms)) + + # force-inf grad method (autograd 1-pass) + pot_ig = build_potential( + rank=3, emit_property_features=False, compute_forces=True, method="grad" + ).to(device) + pot_ig.eval() + + def force_inf_grad(): + with torch.enable_grad(): + pot_ig(factory()) + + ms = _time_it(force_inf_grad, device=device, n_warmup=n_warmup, n_steps=n_steps) + rows.append(_row("force-inf eval grad", ms, n_atoms, "autograd 1-pass")) + + # ---- force training ---------------------------------------------------- + for method, tag, note in ( + ("func", "force-train func emit=0", "func has_aux + loss.backward"), + ("grad", "force-train grad emit=0", "autograd create_graph + loss.backward"), + ): + pot = build_potential( + rank=3, + emit_property_features=False, + compute_forces=True, + method=method, + ).to(device) + pot.train() + opt = torch.optim.Adam(pot.parameters(), lr=1e-3) + + def step(m=pot, o=opt): + o.zero_grad(set_to_none=True) + out = m(factory()) + _ef_loss(out).backward() + o.step() + + ms = _time_it(step, device=device, n_warmup=n_warmup, n_steps=n_steps) + rows.append(_row(tag, ms, n_atoms, note)) + + # emit=1 force train (func default path) + pot = build_potential( + rank=3, emit_property_features=True, compute_forces=True, method="func" + ).to(device) + pot.train() + opt = torch.optim.Adam(pot.parameters(), lr=1e-3) + + def step_emit1(m=pot, o=opt): + o.zero_grad(set_to_none=True) + _ef_loss(m(factory())).backward() + o.step() + + ms = _time_it(step_emit1, device=device, n_warmup=n_warmup, n_steps=n_steps) + rows.append(_row("force-train func emit=1", ms, n_atoms, "func has_aux + loss.backward")) + + # energy-only train (no forces) + pot_et = build_potential(rank=3, emit_property_features=False, compute_forces=False).to(device) + pot_et.train() + opt_et = torch.optim.Adam(pot_et.parameters(), lr=1e-3) + + def energy_train(): + opt_et.zero_grad(set_to_none=True) + out = pot_et(factory()) + _energy_loss(out).backward() + opt_et.step() + + ms = _time_it(energy_train, device=device, n_warmup=n_warmup, n_steps=n_steps) + rows.append(_row("energy-train (no forces)", ms, n_atoms)) + + # ---- forward-only (no loss.backward) train vs eval --------------------- + pot = build_potential( + rank=3, emit_property_features=False, compute_forces=True, method="func" + ).to(device) + b = factory() + + def just_forward_train(): + pot.train() + with torch.enable_grad(): + pot(b.clone()) + + ms_fwd = _time_it(just_forward_train, device=device, n_warmup=n_warmup, n_steps=n_steps) + rows.append(_row("force-train forward only", ms_fwd, n_atoms, "no loss.backward")) + + def just_forward_eval(): + pot.eval() + with torch.enable_grad(): + pot(b.clone()) + + ms_ev = _time_it(just_forward_eval, device=device, n_warmup=n_warmup, n_steps=n_steps) + rows.append( + _row( + "force-eval forward only", + ms_ev, + n_atoms, + f"train/eval fwd ratio={ms_fwd / ms_ev:.2f}x" if ms_ev > 0 else "", + ) + ) + + print(f"\n# batch: atoms={n_atoms} edges={n_edges} graphs={n_graphs} device={device}") + print(f"# torch {torch.__version__} cuda={torch.cuda.is_available()}") + if device.type == "cuda": + print(f"# GPU: {torch.cuda.get_device_name(0)}") + + return rows + + +def profile_cpu_ops(n_steps: int, n_atoms: int, n_edges: int, n_graphs: int) -> None: + from torch.profiler import ProfilerActivity, profile + + factory = _factory("cpu", n_atoms=n_atoms, n_edges=n_edges, n_graphs=n_graphs) + pot = build_potential(rank=3, emit_property_features=False, compute_forces=True) + pot.train() + for _ in range(3): + _ef_loss(pot(factory())).backward() + pot.zero_grad(set_to_none=True) + + with profile(activities=[ProfilerActivity.CPU], record_shapes=False) as p: + for _ in range(n_steps): + _ef_loss(pot(factory())).backward() + pot.zero_grad(set_to_none=True) + print("\n# Top CPU ops — force-train step (self time):") + print(p.key_averages().table(sort_by="self_cpu_time_total", row_limit=20)) + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu") + ap.add_argument("--steps", type=int, default=40) + ap.add_argument("--warmup", type=int, default=10) + ap.add_argument("--atoms", type=int, default=512) + ap.add_argument("--edges", type=int, default=4096) + ap.add_argument("--graphs", type=int, default=16) + ap.add_argument("--profile-ops", action="store_true") + ap.add_argument("--skip-compile", action="store_true", help="Skip torch.compile cells") + args = ap.parse_args() + + if args.device == "cuda" and not torch.cuda.is_available(): + raise SystemExit( + "CUDA requested but torch.cuda.is_available() is False " + f"(torch={torch.__version__}). Install a CUDA build or use --device cpu." + ) + + device = torch.device(args.device) + rows = run_matrix( + device=device, + n_steps=args.steps, + n_warmup=args.warmup, + n_atoms=args.atoms, + n_edges=args.edges, + n_graphs=args.graphs, + skip_compile=args.skip_compile, + ) + print_table(rows, baseline="energy-inf rank3 (no forces)") + + by_name = {r.name: r.ms for r in rows if r.ms == r.ms} + + def ratio(a: str, b: str) -> str: + if a in by_name and b in by_name and by_name[b] > 0: + return f"{by_name[a] / by_name[b]:.2f}x" + return "n/a" + + print("\n# Key ratios") + print( + " force-train-func / energy-train = " + f"{ratio('force-train func emit=0', 'energy-train (no forces)')}" + ) + print( + " force-train-grad / energy-train = " + f"{ratio('force-train grad emit=0', 'energy-train (no forces)')}" + ) + print( + " force-train-func / force-inf = " + f"{ratio('force-train func emit=0', 'force-inf eval emit=0')}" + ) + print( + " force-train-grad / force-inf = " + f"{ratio('force-train grad emit=0', 'force-inf eval emit=0')}" + ) + print( + " force-train-grad / func = " + f"{ratio('force-train grad emit=0', 'force-train func emit=0')}" + ) + print( + " force-inf-grad / func = " + f"{ratio('force-inf eval grad', 'force-inf eval emit=0')}" + ) + print( + " force-inf train-fwd / eval-fwd = " + f"{ratio('force-train forward only', 'force-eval forward only')}" + ) + print(f" enc emit=1 / emit=0 = {ratio('enc rank3 emit=1', 'enc rank3 emit=0')}") + print(f" enc rank3 / rank1 = {ratio('enc rank3 emit=0', 'enc rank1 emit=0')}") + print( + " force-inf rank3 / rank1 = " + f"{ratio('force-inf eval emit=0', 'force-inf eval rank1')}" + ) + + if args.profile_ops and device.type == "cpu": + profile_cpu_ops(min(args.steps, 15), args.atoms, args.edges, args.graphs) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/pinet_compile_probe.py b/benchmarks/pinet_compile_probe.py new file mode 100644 index 0000000..08112ae --- /dev/null +++ b/benchmarks/pinet_compile_probe.py @@ -0,0 +1,304 @@ +"""Probe how torch.compile can wrap current PiNetPotential. + +Run on a GPU node (GH200):: + + PYTHONPATH=src python benchmarks/pinet_compile_probe.py +""" + +from __future__ import annotations + +import traceback + +import torch + +_orig = torch.ops.load_library + + +def soft(p): # noqa: ANN001 + try: + return _orig(p) + except OSError as e: + print(f"[probe] skip op: {e}") + + +torch.ops.load_library = soft # type: ignore[method-assign] + +from molix.compile import Compiler # noqa: E402 +from molix.profiler import MockBatch # noqa: E402 +from molpot.derivation import EnergyReadout, ForceReadout # noqa: E402 +from molzoo.pinet import PiNet, PiNetPotential # noqa: E402 + +ATOM = list(range(1, 8)) +device = "cuda" if torch.cuda.is_available() else "cpu" + + +def factory(n=64, e=256, g=4): + return MockBatch(n_atoms=n, n_edges=e, n_graphs=g, atomic_numbers=7, device=device, seed=0) + + +def make_pot(*, forces=False, method="func", rank=3, depth=2, hidden=32): + pot = PiNetPotential( + atom_types=ATOM, + r_max=5.0, + n_basis=5, + pp_nodes=[hidden, hidden], + pi_nodes=[hidden, hidden], + ii_nodes=[hidden, hidden], + depth=depth, + rank=rank, + hidden_dim=hidden, + compute_forces=forces, + emit_property_features=False, + ) + if method != "func": + pot.deriv_method = method + if forces: + er = EnergyReadout(pot, method=method, backward=True) + fr = ForceReadout(pot, method=method) + pot._pipeline = lambda b, e=er, f=fr: f(e(b)) + else: + pot._pipeline = EnergyReadout(pot, method=method, backward=False) + if device == "cuda": + pot = pot.cuda() + return pot + + +def try_call(label, fn): + print(f"\n=== {label} ===", flush=True) + try: + out = fn() + print("OK", out, flush=True) + return True + except Exception as e: + print(f"FAIL {type(e).__name__}: {str(e).splitlines()[0][:220]}", flush=True) + lines = traceback.format_exc().strip().splitlines() + for line in lines[-25:]: + print(line, flush=True) + return False + finally: + torch.compiler.reset() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + +def main() -> None: + print( + "torch", + torch.__version__, + "device", + device, + "gpu", + torch.cuda.get_device_name(0) if torch.cuda.is_available() else "n/a", + flush=True, + ) + + def a(): + pot = make_pot(forces=False).eval() + c = torch.compile(pot, backend="inductor", fullgraph=False) + with torch.no_grad(): + b = c(factory()()) + return ("energy", float(b["graphs", "energy"].sum())) + + try_call("A energy-only torch.compile(fullgraph=False)", a) + + def b(): + pot = make_pot(forces=False).eval() + c = torch.compile(pot, backend="inductor", fullgraph=True) + with torch.no_grad(): + b = c(factory()()) + return ("energy", float(b["graphs", "energy"].sum())) + + try_call("B energy-only torch.compile(fullgraph=True)", b) + + def c(): + enc = PiNet( + atom_types=ATOM, + r_max=5.0, + n_basis=5, + pp_nodes=[32, 32], + pi_nodes=[32, 32], + ii_nodes=[32, 32], + depth=2, + rank=3, + emit_property_features=False, + ) + if device == "cuda": + enc = enc.cuda() + enc = enc.eval() + cenc = torch.compile(enc, backend="inductor", fullgraph=False) + with torch.no_grad(): + b = cenc(factory()()) + return ("nf", tuple(b["atoms", "node_features"].shape)) + + try_call("C encoder torch.compile(fullgraph=False)", c) + + def d(): + pot = make_pot(forces=False).eval() + + class Core(torch.nn.Module): + def __init__(self, p): + super().__init__() + self.p = p + + def forward(self, batch): + return self.p._write_energy(batch) + + core = torch.compile(Core(pot), backend="inductor", fullgraph=False) + with torch.no_grad(): + b = core(factory()()) + return ("energy", float(b["graphs", "energy"].sum())) + + try_call("D Core(_write_energy) compile fullgraph=False", d) + + def e(): + pot = make_pot(forces=False).eval() + batch = factory()() + expl = torch._dynamo.explain(pot)(batch) + print("graph_count", expl.graph_count, flush=True) + print("graph_break_count", expl.graph_break_count, flush=True) + if expl.break_reasons: + for i, r in enumerate(expl.break_reasons[:8]): + print(f" break[{i}] {str(r)[:240]}", flush=True) + return ("breaks", expl.graph_break_count) + + try_call("E dynamo.explain energy-only", e) + + def f(): + pot = make_pot(forces=True, method="grad").eval() + c = torch.compile(pot, backend="inductor", fullgraph=False) + with torch.enable_grad(): + b = c(factory()()) + return ( + "E", + float(b["graphs", "energy"].sum()), + "F", + float(b["atoms", "forces"].abs().mean()), + ) + + try_call("F force grad torch.compile(fullgraph=False)", f) + + def g(): + pot = make_pot(forces=True, method="func").eval() + c = torch.compile(pot, backend="inductor", fullgraph=False) + with torch.enable_grad(): + b = c(factory()()) + return ( + "E", + float(b["graphs", "energy"].sum()), + "F", + float(b["atoms", "forces"].abs().mean()), + ) + + try_call("G force func torch.compile(fullgraph=False)", g) + + def h(): + pot = make_pot(forces=False).eval() + pot = Compiler(cuda_graphs=True)(pot) + fac = factory(n=64, e=256, g=4) + with torch.no_grad(): + for _ in range(3): + pot(fac()) + if device == "cuda": + torch.cuda.synchronize() + b = pot(fac()) + return ("energy", float(b["graphs", "energy"].sum())) + + try_call("H Compiler(cuda_graphs=True) energy-only", h) + + def i(): + pot = make_pot(forces=False).eval() + pot.encoder = torch.compile(pot.encoder, backend="inductor", fullgraph=False) + pot._pipeline = EnergyReadout(pot, method="func", backward=False) + with torch.no_grad(): + b = pot(factory()()) + return ("energy", float(b["graphs", "energy"].sum())) + + try_call("I compile encoder submodule, energy readout eager", i) + + def j(): + pot = make_pot(forces=True, method="grad").train() + pot = torch.compile(pot, backend="inductor", fullgraph=False) + opt = torch.optim.Adam(pot.parameters(), lr=1e-3) + opt.zero_grad(set_to_none=True) + b = pot(factory()()) + loss = b["graphs", "energy"].pow(2).mean() + b["atoms", "forces"].pow(2).mean() + loss.backward() + opt.step() + return ("loss", float(loss.detach())) + + try_call("J force-train grad compiled fullgraph=False one step", j) + + # K: fullgraph force with grad (old production claim) + def k(): + pot = make_pot(forces=True, method="grad").eval() + pot = torch.compile(pot, backend="inductor", fullgraph=True, dynamic=False) + with torch.enable_grad(): + b = pot(factory()()) + return ( + "E", + float(b["graphs", "energy"].sum()), + "F", + float(b["atoms", "forces"].abs().mean()), + ) + + try_call("K force grad torch.compile(fullgraph=True)", k) + + # L: plain tensors path — extract energy as function of (Z, pos, edge_index, batch) + # (future compile surface if TensorDict is the problem) + def l_flat(): + pot = make_pot(forces=False).eval() + batch = factory()() + + class FlatEnergy(torch.nn.Module): + def __init__(self, p): + super().__init__() + self.encoder = p.encoder + self.out_layers = p.out_layers + self.energy_aggregation = p.energy_aggregation + + def forward(self, Z, pos, edge_index, atom_batch, num_graphs: int): + from tensordict import TensorDict + + td = TensorDict( + { + "atoms": TensorDict( + {"Z": Z, "pos": pos, "batch": atom_batch}, + batch_size=[Z.shape[0]], + ), + "edges": TensorDict( + {"edge_index": edge_index}, + batch_size=[edge_index.shape[0]], + ), + "graphs": TensorDict({}, batch_size=[num_graphs]), + }, + batch_size=[], + ) + td = self.encoder(td) + block_outputs = td["atoms", "p1_block_outputs"] + output = block_outputs.new_zeros(block_outputs.shape[0], 1) + for i, out_layer in enumerate(self.out_layers): + output = out_layer(block_outputs[:, i, :], output) + atom_energy = output.squeeze(-1) + return self.energy_aggregation(atom_energy, atom_batch, num_graphs=num_graphs) + + flat = FlatEnergy(pot) + if device == "cuda": + flat = flat.cuda() + flat = flat.eval() + cflat = torch.compile(flat, backend="inductor", fullgraph=True) + Z = batch["atoms", "Z"] + pos = batch["atoms", "pos"] + ei = batch["edges", "edge_index"] + ab = batch["atoms", "batch"] + ng = int(batch["graphs"].batch_size[0]) + with torch.no_grad(): + e = cflat(Z, pos, ei, ab, ng) + return ("energy", float(e.sum()), "shape", tuple(e.shape)) + + try_call("L flat (Z,pos,edge_index) energy fullgraph=True", l_flat) + + print("\n=== DONE ===", flush=True) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/run_gh200_ljcut_nve.sbatch b/benchmarks/run_gh200_ljcut_nve.sbatch new file mode 100755 index 0000000..087ee9b --- /dev/null +++ b/benchmarks/run_gh200_ljcut_nve.sbatch @@ -0,0 +1,80 @@ +#!/bin/bash +#SBATCH --job-name=ljcut-nve-prec +#SBATCH --account=naiss2026-4-715-gpu +#SBATCH --partition=gpu +#SBATCH --gres=gpu:1 +#SBATCH --nodes=1 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=16 +#SBATCH --time=04:00:00 +#SBATCH --output=%x-%j.out + +# One precision-matrix arm per job (submit 3 in parallel). dtype is the only +# free variable; all MD knobs are pinned. Submit from the repo root: +# +# sbatch --export=ALL,ARM=md64_pot64 benchmarks/run_gh200_ljcut_nve.sbatch +# sbatch --export=ALL,ARM=md64_pot32 benchmarks/run_gh200_ljcut_nve.sbatch +# sbatch --export=ALL,ARM=md32_pot32 benchmarks/run_gh200_ljcut_nve.sbatch +# +# Or omit ARM to run all three arms sequentially in one job. +# Extra CLI args after the script name are forwarded to the verify script. + +set -euo pipefail +echo "=== node $(hostname) arch $(uname -m) $(date) ===" +echo "=== ARM=${ARM:-ALL} ===" + +source "${LMOD_PKG:-/software/sse2/el9_epyc9005/manual/lmod/lmod/lmod}/init/bash" +module purge +module load GPU/buildenv-gcccuda/2026.03-cu13.0 +module load GPU/Python/3.13.5-bare-gcc-2025b-eb + +REPO="${SLURM_SUBMIT_DIR:?submit from the repo root}" +PY=/nobackup/proj/disk/teoroo/personal/jicli594/work/.aarch64/bin/python +cd "$REPO" + +GCC_LIBDIR=$(dirname "$(g++ -print-file-name=libstdc++.so)") +export LD_LIBRARY_PATH="$GCC_LIBDIR:${LD_LIBRARY_PATH:-}" + +"$PY" -c "import torch, platform; print('arch', platform.machine(), 'torch', torch.__version__, 'cuda', torch.cuda.is_available())" + +cmake -S src/molix/op -B src/molix/op/build-gh200 -DMOLNEX_OP_ENABLE_CUDA=ON \ + -DCMAKE_PREFIX_PATH="$("$PY" -c 'import torch; print(torch.utils.cmake_prefix_path)')" +cmake --build src/molix/op/build-gh200 -j + +EXTRA=("$@") +COMMON=( + --ps 5000 + --n 5 + --dt 4.0 + --skin 0 + --every 1 + --delay 0 + --sample-every 1000 + --seed 1 + --device cuda +) + +run_arm () { + local name="$1"; shift + echo "=== arm: $name ===" + PYTHONPATH=src:. "$PY" benchmarks/verify_md_ljcut_nve.py \ + "${COMMON[@]}" ${EXTRA[@]+"${EXTRA[@]}"} "$@" \ + 2>&1 | grep -vE "UserWarning|warnings.warn|W[0-9]+ " +} + +case "${ARM:-ALL}" in + md64_pot64) run_arm md64_pot64 --precision fp64 ;; + md64_pot32) run_arm md64_pot32 --precision fp64 --potential-precision fp32 ;; + md32_pot32) run_arm md32_pot32 --precision fp32 ;; + ALL) + run_arm md64_pot64 --precision fp64 + run_arm md64_pot32 --precision fp64 --potential-precision fp32 + run_arm md32_pot32 --precision fp32 + ;; + *) + echo "unknown ARM=${ARM} (want md64_pot64|md64_pot32|md32_pot32|ALL)" >&2 + exit 2 + ;; +esac + +echo "=== DONE $(date) ===" diff --git a/benchmarks/run_gh200_mace_nve_precision.sbatch b/benchmarks/run_gh200_mace_nve_precision.sbatch new file mode 100755 index 0000000..c892816 --- /dev/null +++ b/benchmarks/run_gh200_mace_nve_precision.sbatch @@ -0,0 +1,77 @@ +#!/bin/bash +#SBATCH --job-name=mace-nve-prec +#SBATCH --account=naiss2026-4-715-gpu +#SBATCH --partition=gpu +#SBATCH --gres=gpu:1 +#SBATCH --nodes=1 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=16 +#SBATCH --time=1-12:00:00 +#SBATCH --output=%x-%j.out + +# One precision-matrix arm per job (submit 3 in parallel). dtype is the only +# free variable; all MD knobs are pinned. 5 ns @ dt=0.5 fs = 10_000_000 steps. +# Submit from the repo root: +# +# sbatch --export=ALL,ARM=md64_pot64 benchmarks/run_gh200_mace_nve_precision.sbatch +# sbatch --export=ALL,ARM=md64_pot32 benchmarks/run_gh200_mace_nve_precision.sbatch +# sbatch --export=ALL,ARM=md32_pot32 benchmarks/run_gh200_mace_nve_precision.sbatch +# +# Or omit ARM to run all three arms sequentially in one job. +# Extra CLI args after the script name are forwarded to every arm. + +set -euo pipefail +echo "=== node $(hostname) arch $(uname -m) $(date) ===" +echo "=== ARM=${ARM:-ALL} ===" + +source "${LMOD_PKG:-/software/sse2/el9_epyc9005/manual/lmod/lmod/lmod}/init/bash" +module purge +module load GPU/buildenv-gcccuda/2026.03-cu13.0 +module load GPU/Python/3.13.5-bare-gcc-2025b-eb + +REPO="${SLURM_SUBMIT_DIR:?submit from the repo root}" +PY=/nobackup/proj/disk/teoroo/personal/jicli594/work/.aarch64/bin/python +WD=/nobackup/proj/disk/teoroo/personal/jicli594/work/mace-nve/precision-matrix +WEIGHTS=/home/jicli594/work/mace_models +cd "$REPO" + +GCC_LIBDIR=$(dirname "$(g++ -print-file-name=libstdc++.so)") +export LD_LIBRARY_PATH="$GCC_LIBDIR:${LD_LIBRARY_PATH:-}" + +"$PY" -c "import torch, platform; print('arch', platform.machine(), 'torch', torch.__version__, 'cuda', torch.cuda.is_available())" + +cmake -S src/molix/op -B src/molix/op/build-gh200 -DMOLNEX_OP_ENABLE_CUDA=ON \ + -DCMAKE_PREFIX_PATH="$("$PY" -c 'import torch; print(torch.utils.cmake_prefix_path)')" +cmake --build src/molix/op/build-gh200 -j + +EXTRA=("$@") +COMMON=(--system "$WD/system.pt" --weights-dir "$WEIGHTS" + --steps 10000000 --dt 0.5 --temperature 300 --stride 10000 + --skin 0 --every 1 --delay 0 --seed 0 + --checkpoint-every 1000000 --flush-every 200 + --compile --device cuda) + +run_arm () { + local name="$1"; shift + echo "=== arm: $name ===" + PYTHONPATH=src:. "$PY" scripts/matpes_port/run_nve.py \ + "${COMMON[@]}" --out "$WD/nve5ns_${name}.pt" ${EXTRA[@]+"${EXTRA[@]}"} "$@" \ + 2>&1 | grep -vE "UserWarning|warnings.warn|cuequivariance_ops|W[0-9]+ " +} + +case "${ARM:-ALL}" in + md64_pot64) run_arm md64_pot64 --precision fp64 ;; + md64_pot32) run_arm md64_pot32 --precision fp64 --potential-precision fp32 ;; + md32_pot32) run_arm md32_pot32 --precision fp32 ;; + ALL) + run_arm md64_pot64 --precision fp64 + run_arm md64_pot32 --precision fp64 --potential-precision fp32 + run_arm md32_pot32 --precision fp32 + ;; + *) + echo "unknown ARM=${ARM} (want md64_pot64|md64_pot32|md32_pot32|ALL)" >&2 + exit 2 + ;; +esac + +echo "=== DONE $(date) ===" diff --git a/benchmarks/run_gh200_pinet_rank3_speed.sbatch b/benchmarks/run_gh200_pinet_rank3_speed.sbatch new file mode 100644 index 0000000..92e43fd --- /dev/null +++ b/benchmarks/run_gh200_pinet_rank3_speed.sbatch @@ -0,0 +1,59 @@ +#!/bin/bash +#SBATCH --job-name=pinet-r3-speed +#SBATCH --account=naiss2026-4-715-gpu +#SBATCH --partition=gpu +#SBATCH --gres=gpu:1 +#SBATCH --nodes=1 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=16 +#SBATCH --time=01:00:00 +#SBATCH --output=benchmarks/pinet-r3-speed-%j.out +#SBATCH --error=benchmarks/pinet-r3-speed-%j.err + +set -euo pipefail +echo "=== node $(hostname) arch $(uname -m) $(date) ===" +nvidia-smi -L || true + +source "${LMOD_PKG:-/software/sse2/el9_epyc9005/manual/lmod/lmod/lmod}/init/bash" +module purge +module load GPU/buildenv-gcccuda/2026.03-cu13.0 +module load GPU/Python/3.13.5-bare-gcc-2025b-eb + +ROOT=/nobackup/proj/disk/teoroo/personal/jicli594/work/molcrafts/molnex +cd "$ROOT" +VENV="$ROOT/.venv-gh200" +export PYTHONPATH="$ROOT/src" +export OMP_NUM_THREADS="${SLURM_CPUS_PER_TASK:-16}" + +if [ ! -x "$VENV/bin/python" ]; then + echo "=== creating aarch64 venv ===" + python -m venv "$VENV" + "$VENV/bin/pip" install -q --upgrade pip + "$VENV/bin/pip" install -q "torch>=2.10" "scikit-build-core>=0.10" "cmake>=4.0.0" ninja \ + numpy tensordict pydantic "molcrafts-molcfg>=1.0" "molcrafts-mollog>=1.0" \ + "molcrafts-molpy>=0.3" cuequivariance cuequivariance-torch opt_einsum_fx tqdm typer zarr numcodecs + # Native op optional for PiNet pure-torch path; soft-skip in the bench if ABI mismatch. + MOLNEX_OP_ENABLE_CUDA=OFF "$VENV/bin/pip" install -q --no-build-isolation -e . --no-deps || true +else + echo "=== reusing venv $VENV ===" +fi + +"$VENV/bin/python" - <<'PY' +import torch, platform +print("arch", platform.machine()) +print("torch", torch.__version__, "cuda", torch.cuda.is_available()) +if torch.cuda.is_available(): + print("gpu", torch.cuda.get_device_name(0)) + x = torch.randn(4, device="cuda") + print("cuda smoke", float(x.sum())) +PY + +export PYTHONUNBUFFERED=1 + +"$VENV/bin/python" benchmarks/bench_pinet_rank3_speed.py --device cuda --steps 50 --warmup 15 \ + --atoms 512 --edges 4096 --graphs 16 + +"$VENV/bin/python" benchmarks/bench_pinet_rank3_speed.py --device cuda --steps 30 --warmup 10 \ + --atoms 2048 --edges 16384 --graphs 64 + +echo "=== DONE $(date) ===" diff --git a/benchmarks/run_l40_pinet_rank3.sbatch b/benchmarks/run_l40_pinet_rank3.sbatch new file mode 100644 index 0000000..ff11d07 --- /dev/null +++ b/benchmarks/run_l40_pinet_rank3.sbatch @@ -0,0 +1,63 @@ +#!/bin/bash +#SBATCH --job-name=pinet-r3-speed +#SBATCH --account=naiss2026-4-715-gpu +#SBATCH --partition=gpu +#SBATCH --gres=gpu:1 +#SBATCH --nodes=1 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=8 +#SBATCH --time=00:45:00 +#SBATCH --output=benchmarks/pinet-r3-speed-%j.out +#SBATCH --error=benchmarks/pinet-r3-speed-%j.err +# +# Note: Arrhenius ``gpu`` partition is GH200 (aarch64). Login-node L40 is +# not a compute target. Prefer ``run_gh200_pinet_rank3_speed.sbatch``. +# This script auto-picks the arch-matching venv if both exist. + +set -euo pipefail +echo "=== node $(hostname) arch $(uname -m) $(date) ===" +nvidia-smi -L || true + +ROOT=/nobackup/proj/disk/teoroo/personal/jicli594/work/molcrafts/molnex +cd "$ROOT" + +ARCH="$(uname -m)" +if [ "$ARCH" = "aarch64" ]; then + # GH200 stack + source "${LMOD_PKG:-/software/sse2/el9_epyc9005/manual/lmod/lmod/lmod}/init/bash" + module purge + module load GPU/buildenv-gcccuda/2026.03-cu13.0 2>/dev/null || true + module load GPU/Python/3.13.5-bare-gcc-2025b-eb 2>/dev/null || true + PY="$ROOT/.venv-gh200/bin/python" +else + PY="$ROOT/.venv-l40/bin/python" +fi + +if [ ! -x "$PY" ]; then + echo "ERROR: missing python at $PY (arch=$ARCH)" >&2 + exit 1 +fi + +export PYTHONPATH="$ROOT/src" +export OMP_NUM_THREADS="${SLURM_CPUS_PER_TASK:-8}" +export PYTHONUNBUFFERED=1 + +"$PY" - <<'PY' +import torch, platform +print("arch", platform.machine()) +print("torch", torch.__version__, "cuda", torch.cuda.is_available()) +if torch.cuda.is_available(): + print("gpu", torch.cuda.get_device_name(0)) + x = torch.randn(4, device="cuda") + print("cuda smoke", float(x.sum())) +PY + +echo "=== small batch atoms=512 graphs=16 ===" +"$PY" benchmarks/bench_pinet_rank3_speed.py --device cuda --steps 50 --warmup 15 \ + --atoms 512 --edges 4096 --graphs 16 + +echo "=== large batch atoms=2048 graphs=64 ===" +"$PY" benchmarks/bench_pinet_rank3_speed.py --device cuda --steps 30 --warmup 10 \ + --atoms 2048 --edges 16384 --graphs 64 + +echo "=== DONE $(date) ===" diff --git a/benchmarks/verify_md_lj_nve.py b/benchmarks/verify_md_lj_nve.py index 014fb2b..0b76bcb 100644 --- a/benchmarks/verify_md_lj_nve.py +++ b/benchmarks/verify_md_lj_nve.py @@ -26,8 +26,7 @@ import torch -from molix.md import LangevinVerletIntegrator, LennardJonesForceField -from molix.md.runner import KB_AMU_A_FS +from molix.md import KB_AMU_A_FS, LangevinVerletIntegrator, LennardJonesForceField # Persistent artifact dir (versioned with the repo for the paper). _OUT_DEFAULT = Path(__file__).resolve().parent / "results" / "md_lj_nve" @@ -165,8 +164,7 @@ def _save_artifacts(out, t, e, pes, kes, traj, e0, rel_drift, rms_rel, ns, args) ax0.axhline(0.0, color="k", lw=0.5, ls=":") ax0.set_ylabel(r"$(E_{\rm tot}-E_0)/|E_0|$ [ppm]") ax0.set_title( - f"LJ$_{{13}}$ NVE, {ns:.1f} ns, dt={_DT:g} fs — " - f"drift {rel_drift:.1e}, RMS {rms_rel:.1e}" + f"LJ$_{{13}}$ NVE, {ns:.1f} ns, dt={_DT:g} fs — drift {rel_drift:.1e}, RMS {rms_rel:.1e}" ) ax1.plot(t.numpy(), pe.numpy(), lw=0.7, color="C0", label="potential") ax1.plot(t.numpy(), ke.numpy(), lw=0.7, color="C1", label="kinetic") diff --git a/benchmarks/verify_md_ljcut_nve.py b/benchmarks/verify_md_ljcut_nve.py new file mode 100644 index 0000000..a0dcad6 --- /dev/null +++ b/benchmarks/verify_md_ljcut_nve.py @@ -0,0 +1,359 @@ +"""Bulk lj/cut NVE on GPU with a compiled force field (LAMMPS-melt state point). + +Drives an FCC argon lattice at the classic ``melt`` state point (ρ* = 0.8442, +T0* = 1.44, r_c = 2.5σ) under NVE with :class:`molix.md.LennardJonesCutForceField` +over a rebuilding :class:`molix.md.NeighborList`. The force evaluation +is ``torch.compile``d — fullgraph inductor by default, ``--cuda-graphs`` for the +``reduce-overhead`` preset — while the *list* decides when to rebuild, under its +own Verlet skin and LAMMPS ``every`` / ``delay`` / ``check`` gate +(``--skin`` / ``--every`` / ``--delay`` / ``--no-check``). ``Integrator.eval_force`` +asks it once per force evaluation, at the positions being evaluated, and runs the +answer eagerly *between* compiled force calls; the fixed-capacity buffers keep +every tensor shape static across rebuilds, which is what lets the compiled force +path survive the whole run. + +The default ``--skin 1.02`` Å is 0.3 σ — the ``neighbor 0.3 bin`` setting shipped +with the LAMMPS ``melt`` example — giving a half-skin of 0.51 Å against a +ballistic per-step displacement of order 0.013 Å at T0, i.e. a rebuild every few +tens of steps rather than every one. ``--skin 0`` is the no-skin limit (rebuild +whenever anything moved at all), useful as the reference arm. + +Pure-GPU requires the molix op built with ``MOLNEX_OP_ENABLE_CUDA=ON`` so the +neighbour rebuild runs on-device (a CPU-only op build fails at list +construction with a dispatch error). + +Units: (amu, Å, fs) with energy in amu·Å²/fs² (see molix.md.integrators); +``skin`` / ``r_build`` in Å, ``every`` / ``delay`` in MD steps. + +Run:: + + PYTHONPATH=src:. python benchmarks/verify_md_ljcut_nve.py # 100 ps + PYTHONPATH=src:. python benchmarks/verify_md_ljcut_nve.py --ps 5 # smoke + PYTHONPATH=src:. python benchmarks/verify_md_ljcut_nve.py --cuda-graphs + PYTHONPATH=src:. python benchmarks/verify_md_ljcut_nve.py --skin 0 # no-skin arm + PYTHONPATH=src:. python benchmarks/verify_md_ljcut_nve.py \\ + --precision fp64 --potential-precision fp32 # split-precision arm + +Pass when ``|slope·duration| / |E_tot(0)| < 1e-3`` with bounded RMS fluctuation +and — at ``skin > 0``, where the counter carries information — no dangerous +builds (``ndanger == 0``). +""" + +from __future__ import annotations + +import argparse +import time +from pathlib import Path + +import torch + +from molix.compile import Compiler +from molix.md import ( + KB_EV_PER_K, + MD, + LennardJonesCutForceField, + MaxwellBoltzmann, + MDHook, + NeighborList, +) + +# Persistent artifact dir (versioned with the repo for the paper). +_OUT_DEFAULT = Path(__file__).resolve().parent / "results" / "md_ljcut_nve" + +# Argon in (amu, Å, fs): ε = 0.0103 eV, σ = 3.4 Å, m = 39.95 amu. +_EPS_EV = 0.0103 +_EPS = _EPS_EV / 103.6426965638 # eV -> amu·Å²/fs² +_SIGMA = 3.4 +_MASS = 39.95 +_RHO_STAR = 0.8442 # LAMMPS melt reduced density +_T0_STAR = 1.44 # LAMMPS melt initial reduced temperature +_CUTOFF = 2.5 * _SIGMA + +_DTYPE_MAP = {"fp64": torch.float64, "fp32": torch.float32} + + +def _fcc(n_cells: int, a: float, *, dtype: torch.dtype) -> tuple[torch.Tensor, torch.Tensor]: + """FCC lattice: ``4 n³`` atoms in a cubic box of side ``n·a``.""" + basis = torch.tensor( + [[0.0, 0.0, 0.0], [0.0, 0.5, 0.5], [0.5, 0.0, 0.5], [0.5, 0.5, 0.0]], dtype=dtype + ) + grid = torch.arange(n_cells, dtype=dtype) + offsets = torch.stack(torch.meshgrid(grid, grid, grid, indexing="ij"), dim=-1).reshape(-1, 3) + pos = (offsets.unsqueeze(1) + basis.unsqueeze(0)).reshape(-1, 3) * a + cell = torch.eye(3, dtype=dtype) * (n_cells * a) + return pos, cell + + +class _EnergySampler(MDHook): + """Record (t, PE, KE, E_tot, T) at every hook-visible chunk boundary.""" + + def __init__(self, dt_fs: float) -> None: + self._dt = float(dt_fs) + self.t_ps: list[float] = [] + self.pe: list[float] = [] + self.ke: list[float] = [] + self.etot: list[float] = [] + self.temp: list[float] = [] + + def clear(self) -> None: + for series in (self.t_ps, self.pe, self.ke, self.etot, self.temp): + series.clear() + + def on_step_end(self, runner, step, obs) -> None: + self.t_ps.append(step * self._dt / 1000.0) + self.pe.append(float(obs.potential)) + self.ke.append(float(obs.kinetic)) + self.etot.append(float(obs.total)) + self.temp.append(float(obs.temperature)) + + +def main() -> int: + ap = argparse.ArgumentParser() + ap.add_argument("--ps", type=float, default=100.0, help="trajectory length in ps") + ap.add_argument("--n", type=int, default=5, help="FCC cells per side (N = 4n^3 atoms)") + ap.add_argument("--dt", type=float, default=4.0, help="timestep (fs)") + ap.add_argument( + "--t0", + type=float, + default=_T0_STAR * _EPS_EV / KB_EV_PER_K, + help="initial temperature (K); default T0* = 1.44", + ) + ap.add_argument( + "--precision", + choices=("fp64", "fp32"), + default="fp64", + help="MD-side precision (trajectory state, integrator, mass)", + ) + ap.add_argument( + "--potential-precision", + choices=("fp64", "fp32"), + default=None, + help="force-field / inference precision; defaults to --precision. " + "Setting it below --precision is the split-precision arm " + "(e.g. MD fp64 over pot fp32)", + ) + ap.add_argument( + "--skin", + type=float, + default=1.02, + help="Verlet skin in A (default 1.02 = 0.3 sigma, the LAMMPS melt setting); " + "the list is built at cutoff + skin and stays complete to cutoff while no " + "atom has moved more than skin/2. 0 = the no-skin limit", + ) + ap.add_argument( + "--every", type=int, default=1, help="attempt a rebuild only every N steps (LAMMPS every)" + ) + ap.add_argument( + "--delay", + type=int, + default=0, + help="attempt no rebuild until N steps after the last one (LAMMPS delay; " + "must be a multiple of --every)", + ) + ap.add_argument( + "--no-check", + action="store_true", + help="rebuild on cadence alone, without the half-skin displacement test " + "(cheaper, and never a free optimisation: it accepts missed pairs)", + ) + ap.add_argument("--capacity-factor", type=float, default=1.5) + ap.add_argument("--sample-every", type=int, default=100, help="steps between energy samples") + ap.add_argument("--seed", type=int, default=1) + ap.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu") + ap.add_argument("--no-compile", action="store_true") + ap.add_argument( + "--cuda-graphs", + action="store_true", + help="compile with the reduce-overhead CUDA-graph preset instead of default inductor", + ) + ap.add_argument("--out", type=Path, default=_OUT_DEFAULT, help="artifact dir (npz + png)") + ap.add_argument("--no-save", action="store_true", help="skip writing artifacts") + args = ap.parse_args() + + md_dtype = _DTYPE_MAP[args.precision] + pot_name = args.potential_precision or args.precision + pot_dtype = _DTYPE_MAP[pot_name] + device = torch.device(args.device) + n_steps = int(args.ps * 1000.0 / args.dt) + a = _SIGMA * (4.0 / _RHO_STAR) ** (1.0 / 3.0) + pos, cell = _fcc(args.n, a, dtype=md_dtype) + n_atoms = pos.shape[0] + pos, cell = pos.to(device), cell.to(device) + + try: + neighbors = NeighborList( + cell=cell, + cutoff=_CUTOFF, + positions=pos, + skin=args.skin, + every=args.every, + delay=args.delay, + check=not args.no_check, + capacity_factor=args.capacity_factor, + ) + except (RuntimeError, NotImplementedError) as err: + if device.type == "cuda": + raise SystemExit( + f"on-device neighbour rebuild failed ({err}).\n" + "Build the molix op with CUDA kernels: " + "cmake -S src/molix/op -B src/molix/op/build -DMOLNEX_OP_ENABLE_CUDA=ON" + ) from err + raise + + # Geometry / nblist stays on the MD-state dtype (like MACE's run_nve path); + # only the LJ parameter buffers (and the pair arithmetic that reads them) + # sit on pot_dtype. Casting the whole module to pot_dtype would drag the + # list with it via _apply and smuggle a second free variable into the + # "precision matrix". + ff = LennardJonesCutForceField(epsilon=_EPS, sigma=_SIGMA, neighbors=neighbors) + ff = ff.to(device=device, dtype=md_dtype) + if pot_dtype != md_dtype: + for name in ("epsilon", "sigma", "cutoff_sq", "energy_shift"): + ff.register_buffer(name, getattr(ff, name).detach().to(dtype=pot_dtype)) + compiled = not args.no_compile + if compiled: + ff = Compiler(cuda_graphs=args.cuda_graphs, fullgraph=True)(ff) + + sampler = _EnergySampler(args.dt) + # No cadence kwarg: the list owns the policy and the integrator derives its + # switch from LennardJonesCutForceField.rebuilds_neighbors. + md = MD( + ff, + mass=_MASS, + dt=args.dt, + gamma=0.0, # NVE + dtype=md_dtype, + device=device, + hooks=[sampler], + ) + vel = MaxwellBoltzmann(_MASS, n_atoms=n_atoms).sample(args.t0, seed=args.seed) + vel = vel.to(device=device, dtype=md_dtype) + + md.run(pos, vel, min(3 * args.sample_every, n_steps), chunk=args.sample_every) # warmup/compile + sampler.clear() + # The warmup left the list built at the *warmup's* final configuration, and + # its counters carrying the warmup's rebuilds. Re-phase it onto the timed + # run's initial configuration with the forced-build escape hatch (so the + # first force evaluation is not served a list held elsewhere under a coarse + # every/delay gate), then read the counters, and report deltas: what is + # printed is exactly what the timed trajectory paid. + neighbors.rebuild(pos) + rebuilds_before, ndanger_before = neighbors.rebuild_count, neighbors.ndanger + if device.type == "cuda": + torch.cuda.synchronize() + t_wall = time.perf_counter() + md.run(pos, vel, n_steps, chunk=args.sample_every) + if device.type == "cuda": + torch.cuda.synchronize() + wall = time.perf_counter() - t_wall + rebuilds = neighbors.rebuild_count - rebuilds_before + ndanger = neighbors.ndanger - ndanger_before + + t = torch.tensor(sampler.t_ps, dtype=torch.float64) + e = torch.tensor(sampler.etot, dtype=torch.float64) + e0 = float(e[0]) + tc = t - t.mean() + slope = (tc * (e - e.mean())).sum() / (tc * tc).sum() + duration = n_steps * args.dt / 1000.0 # ps + rel_drift = float(abs(slope * duration) / abs(e0)) + rms_rel = float((e - e.mean()).pow(2).mean().sqrt() / abs(e0)) + rate = n_steps / wall if wall > 0 else float("nan") + t_mean = sum(sampler.temp[len(sampler.temp) // 2 :]) / max(1, len(sampler.temp) // 2) + + print( + f"lj/cut NVE melt: N={n_atoms} (fcc {args.n}^3) rho*={_RHO_STAR} rc={_CUTOFF:.2f} A " + f"dt={args.dt} fs steps={n_steps} duration={duration / 1000:.3f} ns" + ) + print( + f" device={device} compiled={compiled} cuda_graphs={args.cuda_graphs} " + f"precision: MD {args.precision}, potential {pot_name}" + ) + print( + f" policy: skin={neighbors.skin:g} A every={neighbors.every} delay={neighbors.delay} " + f"check={neighbors.check} r_build={neighbors.r_build:.2f} A" + ) + print( + f" rebuilds={rebuilds} ({rebuilds / max(1, n_steps):.3f} of {n_steps} steps) " + f"ndanger={ndanger}" + ) + print(f" T0={args.t0:.1f} K (2nd half)={t_mean:.1f} K E_tot(0)={e0:.6e}") + print(f" rel energy drift (|slope*dur|/|E0|) = {rel_drift:.3e} (bound 1e-3)") + print(f" rel RMS energy fluctuation = {rms_rel:.3e}") + print(f" steps/s = {rate:.0f} ({rate * args.dt / 1e6 * 86400:.1f} ns/day)") + # At skin=0 every rebuild lands on the first permitted opportunity by + # construction, so ndanger just counts rebuilds and carries no information; + # at skin>0 a nonzero count means a rebuild came too late and pairs were + # missed, which is a wrong PES however small the drift happens to look. + ok = rel_drift < 1e-3 and bool(torch.isfinite(e).all()) + if neighbors.skin > 0.0: + ok = ok and ndanger == 0 + print("RESULT:", "PASS" if ok else "FAIL") + + if not args.no_save: + _save_artifacts( + args.out, sampler, e0, rel_drift, rms_rel, duration, args, n_atoms, pot_name + ) + + return 0 if ok else 1 + + +def _save_artifacts(out, sampler, e0, rel_drift, rms_rel, duration, args, n_atoms, pot_name): + """Write the energy/temperature series (.npz) and a conservation figure (.png).""" + import matplotlib + + matplotlib.use("Agg") + import matplotlib.pyplot as plt + import numpy as np + + out.mkdir(parents=True, exist_ok=True) + t = np.asarray(sampler.t_ps) + e = np.asarray(sampler.etot) + pe = np.asarray(sampler.pe) + ke = np.asarray(sampler.ke) + temp = np.asarray(sampler.temp) + + tag = f"md{args.precision[-2:]}_pot{pot_name[-2:]}" + npz = out / f"ljcut_nve_{tag}.npz" + np.savez_compressed( + npz, + time_ps=t, + e_total=e, + e_pot=pe, + e_kin=ke, + temperature=temp, + meta=np.array( + f"lj/cut argon NVE melt; N={n_atoms}; rho*={_RHO_STAR}; rc={_CUTOFF}A; " + f"dt={args.dt}fs; T0={args.t0:.1f}K; duration={duration:.1f}ps; seed={args.seed}; " + f"skin={args.skin}A; every={args.every}; delay={args.delay}; " + f"check={not args.no_check}; MD={args.precision}; pot={pot_name}; " + f"rel_drift={rel_drift:.3e}; rel_rms={rms_rel:.3e}; units=(amu,A,fs)" + ), + ) + + fig, (ax0, ax1, ax2) = plt.subplots(3, 1, figsize=(6.0, 7.0), sharex=True) + ax0.plot(t, (e - e0) / abs(e0) * 1e6, lw=0.8, color="C3") + ax0.axhline(0.0, color="k", lw=0.5, ls=":") + ax0.set_ylabel(r"$(E_{\rm tot}-E_0)/|E_0|$ [ppm]") + ax0.set_title( + f"lj/cut NVE melt, N={n_atoms}, {duration:.0f} ps, MD {args.precision}/pot {pot_name} — " + f"drift {rel_drift:.1e}, RMS {rms_rel:.1e}" + ) + ax1.plot(t, pe, lw=0.7, color="C0", label="potential") + ax1.plot(t, ke, lw=0.7, color="C1", label="kinetic") + ax1.plot(t, e, lw=0.9, color="k", label="total") + ax1.set_ylabel(r"energy [amu$\cdot$Å$^2$/fs$^2$]") + ax1.legend(loc="best", fontsize=8, ncol=3) + ax2.plot(t, temp, lw=0.7, color="C2") + ax2.set_xlabel("time [ps]") + ax2.set_ylabel("T [K]") + fig.tight_layout() + png = out / f"ljcut_nve_{tag}_energy.png" + fig.savefig(png, dpi=200) + plt.close(fig) + + print(f" saved: {npz}") + print(f" saved: {png}") + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/docs/api/molix.md b/docs/api/molix.md index 7c1919b..403d669 100644 --- a/docs/api/molix.md +++ b/docs/api/molix.md @@ -6,7 +6,9 @@ ::: molix.core.state -::: molix.core.hooks +::: molix.core.hook + +::: molix.hooks ::: molix.core.losses @@ -14,14 +16,16 @@ ::: molix.data -::: molix.data.types - ::: molix.data.collate ::: molix.data.pipeline ::: molix.data.datamodule +::: molix.data.task + +::: molix.data.tasks + ## Datasets ::: molix.datasets diff --git a/docs/api/molzoo.md b/docs/api/molzoo.md index cf29f8b..6e1ab53 100644 --- a/docs/api/molzoo.md +++ b/docs/api/molzoo.md @@ -4,6 +4,20 @@ ::: molzoo.mace +::: molzoo.mace.spec + +::: molzoo.mace.geometry + +::: molzoo.mace.encoder + +::: molzoo.mace.potential + +::: molzoo.mace.checkpoint + +::: molzoo.mace.variants + +::: molzoo.mace.research + ## Allegro ::: molzoo.allegro diff --git a/docs/molix/explanation/batch-schema.md b/docs/molix/explanation/batch-schema.md index ad0ff96..f52342e 100644 --- a/docs/molix/explanation/batch-schema.md +++ b/docs/molix/explanation/batch-schema.md @@ -6,16 +6,19 @@ shapes are intentionally different — you must know which side you're on: | Stage | Container | Example access | |---|---|---| | Pre-collate (source, pipeline task I/O, `MmapDataset[i]`) | **flat `dict`** | `sample["Z"]`, `sample["edge_index"]` | -| Post-collate (`GraphBatch` from `collate_molecules`) | **nested `TensorDict`** | `batch["atoms", "Z"]`, `batch["edges", "edge_index"]` | +| Post-collate (output of `collate_molecules`) | **nested `TensorDict`** | `batch["atoms", "Z"]`, `batch["edges", "edge_index"]` | The single conversion point is `collate_molecules` (invoked by `DataModule._CollateFn`). Tuple-key access like `batch["atoms", "Z"]` is a `TensorDict`-only feature and **does not work** on a raw sample dict — that's why `sample["edges", "edge_index"]` raises `KeyError`. -Nested `TensorDict` subclasses (defined in `molix.data.types`) are the -batch-side containers. Each level carries its own batch size, enabling -natural per-atom, per-edge, and per-graph operations. +The batch side is a **plain `tensordict.TensorDict`** — there is no subclass and +no `@tensorclass` wrapper. What makes it a "molecular batch" is the namespace +layout (`atoms` / `edges` / `graphs`, plus `bonds` when the samples carry +covalent topology), not a Python type. Each namespace is itself a `TensorDict` +carrying its own `batch_size`, which is what lets per-atom, per-edge and +per-graph tensors of different lengths live in one container. ## Sample Schema (pre-collate, single molecule, plain flat dict) @@ -32,39 +35,50 @@ top-level keys** (no `"atoms"` / `"edges"` nesting): Access with flat keys: `sample["Z"]`, `sample["edge_index"]`, `sample["targets"]["U0"]`. The nested tuple-key syntax below is for the -post-collate `GraphBatch` only. +post-collate batch only. ## Batch Schema (nested TensorDict) -`collate_molecules` converts a list of sample dicts into a `GraphBatch`: +`collate_molecules` converts a list of sample dicts into one nested +`TensorDict`: ``` -GraphBatch (batch_size=[]) -├── "atoms": AtomData (batch_size=[N_total]) +TensorDict (batch_size=[]) +├── "atoms": TensorDict (batch_size=[N_total]) │ ├── Z: LongTensor[N_total] │ ├── pos: FloatTensor[N_total, 3] │ ├── batch: LongTensor[N_total] # graph membership │ └── -├── "edges": EdgeData (batch_size=[E_total]) +├── "edges": TensorDict (batch_size=[E_total]) │ ├── edge_index: LongTensor[E_total, 2] │ ├── edge_diff: FloatTensor[E_total, 3] │ └── edge_dist: FloatTensor[E_total] -└── "graphs": GraphData (batch_size=[B]) - ├── num_atoms: LongTensor[B] - └── +├── "graphs": TensorDict (batch_size=[B]) +│ ├── num_atoms: LongTensor[B] +│ └── +└── "bonds": TensorDict (batch_size=[]) # only when samples carry bonds + ├── bond_index: LongTensor[2, N_bonds] + └── bond_types: Tensor[N_bonds] ``` -## Type Hierarchy +## Namespaces -| Type | Extends | batch_size | Purpose | -|------|---------|------------|---------| -| `AtomData` | `TensorDict` | `[N]` | Per-atom tensors (encoder adds `node_features` in place) | -| `EdgeData` | `TensorDict` | `[E]` | Per-edge tensors (encoder adds `edge_features` in place) | -| `GraphData` | `TensorDict` | `[B]` | Per-graph tensors + targets | -| `GraphBatch` | `TensorDict` | `[]` | Top-level container | +| Namespace | batch_size | Purpose | +|------|------------|---------| +| `atoms` | `[N_total]` | Per-atom tensors (encoder adds `node_features` in place) | +| `edges` | `[E_total]` | Per-edge tensors (encoder adds `edge_features` in place) | +| `graphs` | `[B]` | Per-graph tensors + graph-level targets | +| `bonds` | `[]` | Covalent topology, present only when samples supply it | +| *(top level)* | `[]` | Container holding the namespaces above | -Encoder outputs are written into the existing `AtomData` / `EdgeData` -sub-dicts by key addition — no subclass swap. +Every level is a plain `TensorDict`; the batch sizes differ because the +number of atoms, edges, graphs and bonds in a batch are unrelated counts. +`bonds` is the exception with `batch_size=[]`: `bond_index` is COO-shaped +`[2, N_bonds]`, so its leading dimension is 2, not the bond count, and it +cannot share a batch axis with `bond_types` `[N_bonds]`. + +Encoder outputs are written into the existing `atoms` / `edges` sub-dicts by +key addition — the batch object is mutated in place, never replaced. ## Access Patterns @@ -77,10 +91,15 @@ batch["graphs", "energy"] # graph-level target (B,) ## Conventions -- Graph-level targets (energy, U0, etc.) are stored in `GraphData`, shape `[B]`. -- Atom-level targets (forces) are stored in `AtomData`, shape `[N_total, ...]`. -- `edge_index` is always `[E, 2]` with `[:, 0] = source`, `[:, 1] = destination`. -- Models receive the `GraphBatch` directly and access nested keys as needed. +- Graph-level targets (energy, U0, etc.) go under `graphs`, shape `[B]`. +- Atom-level targets (forces) go under `atoms`, shape `[N_total, ...]`. + Which target name is routed where is declared by the `TargetSchema` passed to + `collate_molecules`; names in `atom_level` go to `atoms`, everything else is + flattened to `[B]` under `graphs`. +- `edge_index` is always `[E, 2]` with `[:, 0] = source`, `[:, 1] = target`, and + `edge_diff = pos[target] - pos[source]`. Per-molecule edge indices are rebased + onto the concatenated atom numbering during collation. +- Models receive the whole batch and access nested keys as needed. - Loss functions receive `(predictions, batch)` and read targets from the batch. ## Related Pages diff --git a/docs/molix/explanation/throughput-and-compilation.md b/docs/molix/explanation/throughput-and-compilation.md index 1a20c05..47bfc2c 100644 --- a/docs/molix/explanation/throughput-and-compilation.md +++ b/docs/molix/explanation/throughput-and-compilation.md @@ -185,3 +185,42 @@ At production journal cadence the full hook stack costs only ~2 % (129.8 vs 132.2) — the earlier ~91 steps/s figure was an artifact of a deliberately frequent `journal_every=20`. So the real production throughput with the winning combo is **~130 steps/s** (≈10x eager's ~12, ≈4–5x plain compile's ~25–34). + +## Does this transfer to cuEquivariance models? (MACE-MatPES, 2026-08-07) + +The sweep above is PiNet: **pure torch**, so dynamo has nothing exotic to trace. +MACE runs cuEquivariance fused kernels, which are custom `autograd.Function`s — +a plausible reason for dynamo to break the graph and for `fullgraph=True` to +fail outright. Measured on `molzoo.MACEMatpes` with official +`mace-matpes-r2scan-0` weights, 193-atom periodic water box (17344 edges), one +GH200, compiling `_compute_energy` with `autograd.grad` taken *outside* the +compiled region: + +| | fp64 ms (step/s) | fp32 ms (step/s) | +|---|---|---| +| eager | 35.7 (28.0) | 35.0 (28.6) | +| inductor default | 19.0 (52.5) | 19.6 (50.9) | +| **inductor + `reduce-overhead`** | **4.18 (239)** | **2.40 (417)** | +| + `fullgraph=True` | 4.19 (239) | 2.40 (417) | + +**It transfers unchanged.** `CUDA_GRAPH_PRESET` is the winner here too (8.5x +eager at fp64, 14.6x at fp32). + +Three things worth knowing: + +1. **cuEq does not break the graph**: `graph_breaks=0, graphs=1, ops=381`. + `fullgraph=True` is therefore free rather than beneficial — it closes no + breaks, it only asserts there are none. Keep it as the assertion. +2. **Compiling fp64 is numerically free**: ΔE = 0, ΔF = 2.6e-14 against eager. + fp32 shifts results ~2e-3 eV / ~1.6e-3 eV/Å (inductor fusion reassociates + float adds) — expected, not a defect. +3. **Precision is invisible until the launches are gone.** Eager fp32 ≈ eager + fp64 (35.0 vs 35.7 ms) because the step is latency-bound. Under CUDA graphs + fp32 is 1.75x fp64. Anyone benchmarking precision *in eager mode* on a small + system will wrongly conclude precision does not matter. + +Static shapes come free for MD here. Open-system runs freeze the neighbour +list for the trajectory; periodic runs rebuild it on a cadence into +fixed-capacity buffers (`molix.md.NeighborList` — contents change in +place, shapes never do). Either way no padding registry is needed, unlike the +training path. diff --git a/docs/molix/index.md b/docs/molix/index.md index 32681a3..6adecc9 100644 --- a/docs/molix/index.md +++ b/docs/molix/index.md @@ -19,9 +19,14 @@ plain-`TensorDict` batch contract used by models and losses (namespaces lifecycle behavior. - [Data Pipeline](user-guide/data.md): understand sources, preprocessing, caching, collation, and data modules. -- [Data Loading](user-guide/data-loading.md): convert flat samples into - `GraphBatch` objects. +- [Data Loading](user-guide/data-loading.md): convert flat sample dicts into + collated nested `TensorDict` batches. - [Data Modules](user-guide/data-modules.md): wire datasets into `Trainer`. +- [Profiling](user-guide/profiling.md): measure task, module, DataLoader, + Trainer and dataset cost with the `molix.profiler` suite. +- [Molecular Dynamics](user-guide/md.md): run trajectories over a trained + potential with `MD` — force fields, MD hooks, and the split MD/inference + precision model. ## Explanation diff --git a/docs/molix/tutorials/train-a-graph-model.md b/docs/molix/tutorials/train-a-graph-model.md index 10078f3..d070682 100644 --- a/docs/molix/tutorials/train-a-graph-model.md +++ b/docs/molix/tutorials/train-a-graph-model.md @@ -6,7 +6,8 @@ from flat molecule samples. Here, the batch is built directly. ## 1. Define the Model -Models that receive a `GraphBatch` access data with nested tuple keys: +A collated batch is a plain nested `TensorDict`, so models read it with tuple +keys — the first element names the namespace, the second the field: ```python import torch @@ -36,32 +37,40 @@ class SimpleGraphModel(nn.Module): ## 2. Build a Batch +Every level is a `tensordict.TensorDict`; there is no molecule-specific +subclass. The `batch_size` you give each namespace is what makes the three +different lengths (5 atoms, 0 edges, 1 graph) coexist in one container. + ```python -from molix.data.types import AtomData, EdgeData, GraphBatch, GraphData +from tensordict import TensorDict -atoms = AtomData( +atoms = TensorDict( Z=torch.tensor([6, 1, 1, 1, 1]), pos=torch.randn(5, 3), batch=torch.zeros(5, dtype=torch.long), batch_size=[5], ) -edges = EdgeData( +edges = TensorDict( edge_index=torch.zeros(0, 2, dtype=torch.long), edge_diff=torch.zeros(0, 3), edge_dist=torch.zeros(0), batch_size=[0], ) -graphs = GraphData( +graphs = TensorDict( num_atoms=torch.tensor([5]), energy=torch.tensor([-40.5]), batch_size=[1], ) -batch = GraphBatch(atoms=atoms, edges=edges, graphs=graphs, batch_size=[]) +batch = TensorDict(atoms=atoms, edges=edges, graphs=graphs, batch_size=[]) ``` +This is one methane molecule with no edges — enough to exercise the training +loop, since the model above only embeds `Z` and pools per graph. A real run +would let `NeighborList` fill the `edges` namespace. + ## 3. Train ```python diff --git a/docs/molix/user-guide/data-loading.md b/docs/molix/user-guide/data-loading.md index 448dd0b..57d1477 100644 --- a/docs/molix/user-guide/data-loading.md +++ b/docs/molix/user-guide/data-loading.md @@ -3,8 +3,8 @@ MolNex data flows from plain dict samples to nested TensorDict batches: 1. `DataSource.__getitem__` returns a single-sample dict (`Z`, `pos`, `targets`, ...) -2. `DataLoader(collate_fn=collate_molecules)` merges samples into a `GraphBatch` (nested TensorDict) -3. `Trainer` passes the `GraphBatch` to the model and loss function +2. `DataLoader(collate_fn=collate_molecules)` merges samples into one nested `TensorDict` +3. `Trainer` passes that batch to the model and loss function ## Collation @@ -15,15 +15,33 @@ from molix.data.collate import collate_molecules loader = DataLoader(dataset, batch_size=32, shuffle=True, collate_fn=collate_molecules) ``` -`collate_molecules` produces a nested `GraphBatch`: +`collate_molecules` produces a plain nested `TensorDict` whose namespaces each +carry their own batch size: -- Atom-level fields (`Z`, `pos`) → `AtomData` (batch_size=[N_total]) -- Edge fields (`edge_index`, `edge_diff`, `edge_dist`) → `EdgeData` (batch_size=[E_total]) -- Graph-level metadata + targets → `GraphData` (batch_size=[B]) +- Atom-level fields (`Z`, `pos`, `batch`) → `atoms` (batch_size=[N_total]) +- Edge fields (`edge_index`, `edge_diff`, `edge_dist`) → `edges` (batch_size=[E_total]) +- Graph-level metadata (`num_atoms`) + graph targets → `graphs` (batch_size=[B]) +- Covalent topology (`bond_index`, `bond_types`), when present → `bonds` (batch_size=[]) + +Read it with tuple keys: `batch["atoms", "Z"]`, `batch["edges", "edge_index"]`, +`batch["graphs", "energy"]`. + +Two collate paths exist and are required to agree leaf for leaf. +`collate_molecules` walks a list of per-sample dicts; `collate_packed` builds +the same batch directly out of the packed `PackedCache` tensors by gathering +rows, skipping the unpack-then-repack round trip. `collate_molecules` is the +equivalence oracle for `collate_packed`. ## Preprocessing Tasks -- **NeighborList**: Compute neighbor edges in the `prepare()` stage -- **AtomicDress**: Remove atomic baselines from graph-level scalar targets +Tasks run *before* collation, on flat sample dicts, and are composed into a +`Pipeline`. They come in two flavours: + +- **`NeighborList`** is a `SampleTask`: it sees one molecule at a time and adds + `edge_index`, `edge_diff` and `edge_dist` for every pair within `cutoff`. +- **`AtomicDress`** is a `DatasetTask`: it needs the whole training set, so it + runs in two phases — `fit` solves a least-squares problem for a per-element + baseline energy, then `execute` subtracts that baseline from each sample's + scalar target. For the full batch structure, see [Batch Schema](../explanation/batch-schema.md). diff --git a/docs/molix/user-guide/data-modules.md b/docs/molix/user-guide/data-modules.md index 6bbae09..bb4b61f 100644 --- a/docs/molix/user-guide/data-modules.md +++ b/docs/molix/user-guide/data-modules.md @@ -1,42 +1,76 @@ # Data Modules -`DataModule` integrates the full data pipeline: source → pipeline → collation → DataLoader. +`DataModule` is the last stage of the data pipeline: it wraps two pre-built +datasets in DataLoaders that collate, shard across DDP ranks, and prefetch. +Everything upstream — downloading, per-sample transforms, caching, splitting — +happens before it. The minimal protocol requires: -- `setup(stage)` - Prepare datasets -- `train_dataloader()` - Returns iterable of `GraphBatch` -- `val_dataloader()` - Returns iterable of `GraphBatch` +- `setup(stage)` - Prepare datasets/samplers for a stage (e.g. `"fit"`) +- `train_dataloader()` - Returns an iterable of collated training batches +- `val_dataloader()` - Returns an iterable of collated validation batches -Each batch element is a `GraphBatch` (nested TensorDict, see -[Batch Schema](../explanation/batch-schema.md)). +`DataModuleProtocol` also declares `on_epoch_start(epoch)`, which the `Trainer` +calls when the module defines it (e.g. to reseed a sampler). + +Each item yielded by those loaders is a nested `TensorDict` with `atoms` / +`edges` / `graphs` namespaces, see +[Batch Schema](../explanation/batch-schema.md). ## Using the Built-in DataModule +The wiring is source → `PipelineSpec.cache` → dataset → split → `DataModule`: + ```python -from molix.data import DataModule, Pipeline, NeighborList, AtomicDress +from molix.data import ( + AtomicDress, DataModule, NeighborList, Pipeline, SubsetDataset, SubsetSource, +) from molix.datasets import QM9Source source = QM9Source(root="./data/qm9", total=1000) +train_idx = list(range(800)) +val_idx = list(range(800, 1000)) + +# AtomicDress fits its per-element baseline with a least-squares solve over a +# whole dataset. Fit it on the training indices only, or the val split leaks +# into the baseline. +train_source = SubsetSource(source, train_idx) + pipe = ( Pipeline("qm9") .add(NeighborList(cutoff=5.0)) - .add(AtomicDress(target_key="U0")) + .add(AtomicDress(elements=(1, 6, 7, 8, 9), target_key="U0")) .build() ) -dm = DataModule(source=source, pipeline=pipe, batch_size=32) +dag = pipe.cache(source, base_dir="./cache", fit_source=train_source) +full = dag.dataset(mmap=True) +train_ds = SubsetDataset(full, train_idx) +val_ds = SubsetDataset(full, val_idx) + +dm = DataModule( + train_ds, + val_ds, + target_schema=QM9Source.TARGET_SCHEMA, + batch_nodes=pipe.batch_nodes, + batch_size=32, +) dm.setup("fit") for batch in dm.train_dataloader(): - # batch is a GraphBatch + # batch is a nested TensorDict Z = batch["atoms", "Z"] # (N_total,) pos = batch["atoms", "pos"] # (N_total, 3) energy = batch["graphs", "U0"] # (B,) break ``` +`target_schema` decides which target names land under `graphs` and which under +`atoms`; `batch_nodes` carries any `BatchTask` nodes in the pipeline, which run +after collation and so cannot be cached with the rest. + ## Minimal Custom DataModule ```python diff --git a/docs/molix/user-guide/data.md b/docs/molix/user-guide/data.md index 6ba39e4..7523b48 100644 --- a/docs/molix/user-guide/data.md +++ b/docs/molix/user-guide/data.md @@ -2,12 +2,13 @@ `molix.data` provides the molecular data pipeline: -- **Types** (`types.py`): Nested TensorDict subclasses — `AtomData`, `EdgeData`, `GraphData`, `GraphBatch` -- **Sources**: `DataSource` protocol and `InMemorySource` / `SubsetSource` implementations -- **Pipeline**: Task-based preprocessing pipeline (sample-level, dataset-level, batch-level) -- **Tasks**: Built-in preprocessing tasks (`NeighborList`, `AtomicDress`) -- **Collation**: `collate_molecules` converts sample dicts into nested `GraphBatch` -- **DataModule**: DDP-aware data module integrating pipeline + collation + DataLoader +- **Sources** (`source.py`): `DataSource` protocol and `InMemorySource` / `SubsetSource` implementations +- **Tasks** (`task.py`, `tasks/`): transform primitives (`SampleTask`, `DatasetTask`, `BatchTask`) and the built-ins `NeighborList`, `AtomicDress`, `UnitConvert`, `ConstantLabel`, `PadMolecularBatch` +- **Pipeline** (`pipeline.py`): declarative `Pipeline` / `PipelineSpec` container — which tasks run, in what order, under what cache identity +- **Cache** (`cache.py`): `PackedCache`, the single-file packed store a materialized pipeline writes +- **Datasets** (`dataset.py`): `MmapDataset` / `CachedDataset` / `SubsetDataset` readers over a `PackedCache` +- **Collation** (`collate.py`): `collate_molecules` turns sample dicts into a plain nested `TensorDict` with `atoms` / `edges` / `graphs` (and `bonds`) namespaces; `collate_packed` is the equivalent fast path straight off packed cache tensors +- **DataModule** (`datamodule.py`): DDP-aware data module integrating pipeline + collation + DataLoader Recommended reading order: diff --git a/docs/molix/user-guide/hooks.md b/docs/molix/user-guide/hooks.md index cf55dce..a281465 100644 --- a/docs/molix/user-guide/hooks.md +++ b/docs/molix/user-guide/hooks.md @@ -4,23 +4,47 @@ Hooks add behavior around the training loop without replacing `Trainer`. Use them for logging, metrics, checkpoints, profiling, telemetry, learning-rate events, and custom lifecycle logic. +Hooks live in two modules, and the split matters when you write imports: + +- `molix.core.hook` (singular) — the *contract* layer: the `Hook` protocol, the + `BaseHook` no-op base class, and `ScalarHook`. +- `molix.hooks` (plural) — every *concrete* implementation (`Log`, `TensorBoardHook`, + `MetricsHook`, `StepSpeedHook`, `CheckpointHook`, `JournalHook`, + `GradClipHook`, `ProgressBarHook`, `ProfilerHook`, `GPUMemoryHook`, + `GPUUtilsHook`, `MolRecMetricsHook`, `ActivationCheckpointingHook`, + `EarlyStop`). + ## Registration ```python -from molix.core.hooks import Log, TensorBoardHook from molix.core.trainer import Trainer +from molix.hooks import Log, StepSpeedHook, TensorBoardHook + +speed = StepSpeedHook() trainer = Trainer( model=model, loss_fn=loss_fn, optimizer_factory=opt_factory, hooks=[ - Log(every_n_steps=100), - TensorBoardHook(log_dir="runs/experiment-1"), + speed, + Log(every_n_steps=100, keys=[("train", "loss"), speed]), + TensorBoardHook(every_n_steps=100, log_dir="runs/experiment-1"), ], ) ``` +`Log` has no default column set: `keys` is required and each entry is either a +state path (`("train", "loss")`, or the equivalent slash string `"train/loss"`) +or a `ScalarHook` instance, in which case `Log` expands that hook's +`scalar_keys`. At `on_train_start`, `Log` rejects any key that is neither a +built-in state path nor advertised by a registered `ScalarHook`, so a typo +fails loudly instead of printing a column of dashes. + +`TensorBoardHook` needs no key list — it scans the `train`, `performance` and +`gpu` namespaces on each logged train step and `eval` on each eval completion, +writing every numeric or 0-d tensor value it finds. + Hooks run in registration order by default. To force an order, pass `(hook, priority)` tuples. Lower priorities run earlier. @@ -37,7 +61,7 @@ hooks = [ Subclass `BaseHook` and override only the methods you need: ```python -from molix.core.hooks import BaseHook +from molix.core.hook import BaseHook class NaNStopperHook(BaseHook): @@ -47,7 +71,11 @@ class NaNStopperHook(BaseHook): raise RuntimeError(f"Non-finite loss at step {state.global_step}") ``` -Common hook points include: +The default steps return `{"loss": ..., "predictions": ...}` (the train step +adds `"optimizer_applied"`), which is what `outputs` holds in the batch-end +callbacks. + +The full set of callbacks the `Trainer` dispatches: - `on_train_start` - `on_train_end` @@ -55,11 +83,20 @@ Common hook points include: - `on_epoch_end` - `on_train_batch_start` - `on_train_batch_end` +- `on_after_backward` +- `on_eval_phase_start` - `on_eval_batch_start` - `on_eval_batch_end` -- `on_after_backward` +- `on_eval_step_complete` + +`on_eval_phase_start` and `on_eval_step_complete` bracket **every** eval phase — +both the step-based one triggered by `eval_every_n_steps` and the epoch-end one. +Accumulating hooks should reset their eval buffers in `on_eval_phase_start` and +publish in `on_eval_step_complete`; publishing there (rather than in +`on_epoch_end`) also puts the value in `state` before the LR scheduler reads it. -Hook exceptions propagate. A hook that detects an invalid run should raise +Hook exceptions propagate: `Trainer._call_hooks` logs the failure with a +traceback and then re-raises. A hook that detects an invalid run should raise instead of silently logging and continuing. ## State Writes @@ -68,17 +105,26 @@ instead of silently logging and continuing. access: ```python -state["train"]["loss"] = loss.item() +state["train"]["loss"] = loss.detach() state["eval"]["MAE"] = mae state["performance"]["step_per_second"] = rate state["gpu"]["peak_gib"] = peak ``` +Note the `detach()` rather than `item()` on the per-step value: `.item()` forces +a CPU↔GPU synchronization, and doing that on every training step drains the GPU +queue and serializes an otherwise launch-bound loop. Everything on the per-step +path (`DefaultTrainStep`, `GradClipHook`, `MetricsHook`'s train side) therefore stores +the 0-d device tensor and leave materialization to the consumers, which sample +on their own throttled cadence. Cold-path values — an eval metric published once +per eval phase — are converted to `float` at the write, because the LR scheduler +and `CheckpointHook` compare and serialize them as plain numbers. + Do not write slash or tuple paths: ```python -state["train/loss"] = loss.item() # raises ValueError -state[("train", "loss")] = loss.item() # raises ValueError +state["train/loss"] = loss.detach() # raises ValueError +state[("train", "loss")] = loss.detach() # raises ValueError ``` Reads support all three forms: @@ -89,17 +135,48 @@ state["eval/MAE"] state[("eval", "MAE")] ``` +Which namespace a hook may write depends on the callback it is writing from: + +- `on_train_batch_end` / `on_after_backward` → `train`, `performance`, `gpu` +- `on_eval_batch_end` / `on_eval_step_complete` → `eval` + +A hook that reports the same metric on both sides must hold two independent +accumulators (`MetricsHook` deep-copies its metrics into `train_metrics` and +`val_metrics` for exactly this reason) — sharing one buffer across phases lets +an eval-side `reset()` corrupt the train-side value. + ## ScalarHook Hooks that produce scalar values for other hooks to consume should subclass -`ScalarHook` and declare `scalar_keys`. +`ScalarHook` and declare `scalar_keys` — a tuple of state paths, where each path +is either a top-level string key or a `(namespace, name)` tuple. ```python -from molix.core.hooks import ScalarHook +from molix.core.hook import ScalarHook + +class LearningRateHook(ScalarHook): + scalar_keys = (("train", "lr"),) -class StepSpeedHook(ScalarHook): - scalar_keys = (("performance", "step_per_second"),) + def on_train_batch_end(self, trainer, state, batch, outputs): + state["train"]["lr"] = trainer.optimizer.param_groups[0]["lr"] ``` -The `Log` hook can use these paths to decide which columns to render. +`Log` reads `scalar_keys` to expand a hook passed in its `keys` list into +columns, and uses the union of all registered hooks' `scalar_keys` (plus the +built-in paths `epoch`, `global_step`, `stage`, `steps_since_last_eval`, +`best_metric`, `train/loss`, `eval/loss`) to validate the keys you asked for. + +If the paths depend on runtime configuration — `MetricsHook` derives them from +the metric class names it was given — override `scalar_keys` as a `@property` +instead of setting it as a class attribute. + +Shipped `ScalarHook` subclasses and the paths they publish: + +| Hook | Writes | +|---|---| +| `StepSpeedHook` | `performance/step_per_second` | +| `GradClipHook` | `train/grad_norm` (pre-clip L2 norm) | +| `MetricsHook` | `train/` and `eval/`, one per metric class | +| `GPUMemoryHook` | `gpu/alloc_gib`, `gpu/resv_gib`, `gpu/peak_gib` (selected subset) | +| `GPUUtilsHook` | `gpu/util_pct`, `gpu/mem_util_pct` (selected subset) | diff --git a/docs/molix/user-guide/md.md b/docs/molix/user-guide/md.md new file mode 100644 index 0000000..557e98d --- /dev/null +++ b/docs/molix/user-guide/md.md @@ -0,0 +1,136 @@ +# Molecular Dynamics + +`molix.md` is the in-process MD engine: BAOAB Langevin velocity-Verlet over +any trained potential, compilable end to end. **`MD` is the entry point** — +the lower layers (`ForceField`, `Integrator`, `MDRunner`) are the primitives +it composes, and you reach for them directly only when you need a custom loop. + +## Quick start + +```python +import torch +from molix.md import MD, MaxwellBoltzmann, PotentialForceField, TrajectoryHook + +force = PotentialForceField(potential, template) # bind a potential to a system +md = MD(force, mass=masses, dt=0.5, gamma=0.1, temperature=300.0, + dtype=torch.float64, + hooks=[TrajectoryHook("traj.pt", stride=10, numbers=Z)]) + +vel = MaxwellBoltzmann(masses).sample(300.0, seed=0) +final = md.run(pos, vel, n_steps=100_000) # -> MDState(pos, vel, forces, energy) +``` + +Units are (amu, Å, fs); energy in amu·Å²/fs². Drive an eV/Å potential with +`energy_scale=1 / molix.units.EV_PER_AMU_A2_FS2` on the force field. + +## Two precisions, deliberately independent + +`MD(dtype=)` governs the **MD side only**: trajectory state, integrator step +constants, mass. The potential's precision is a separate axis: + +```python +md = MD(force, mass=m, dt=0.5, dtype=torch.float64) # fp64 trajectory ... +md.set_potential_dtype(torch.float32) # ... over fp32 inference +``` + +This is what lets you study the MD process and the inference process +separately — an fp64 trajectory over a quantized/fp32 model is a supported, +meaningful configuration. At the component boundary the integrator casts the +force field's output back into the state dtype, so the two precisions never +silently promote mid-step. For mixed precision inside the model only, use +`MD(autocast_dtype=torch.bfloat16)`. + +## Choosing a force field + +| You have | Use | +|---|---| +| A molpot potential + collated template (open system) | `PotentialForceField` | +| A TensorDict potential reading `edges.shifts` + a periodic cell | `PeriodicPotentialForceField` (owns a rebuilding `NeighborList`) | +| A bulk Lennard-Jones system (periodic, truncated-shifted, LAMMPS `lj/cut`) | `LennardJonesCutForceField` over a `NeighborList` | +| Any `pos -> (energy, forces)` callable (AOTI `.pt2`, compiled closure, external engine) | `CallableForceField` | +| An analytic test PES | `HarmonicForceField`, `LennardJonesForceField` | + +The **list owns the rebuild cadence**: `NeighborList(skin=, every=, delay=, +check=)` carries the LAMMPS `neigh_modify` policy, and `Integrator.eval_force` +merely asks once per force evaluation, *at the positions being evaluated* — +the list decides, into fixed-capacity buffers whose shapes never change, which +is what lets the force path stay inside a CUDA graph across the whole +trajectory. `skin=0, every=1, delay=0, check=True` is the accurate no-skin +limit (rebuild whenever anything moved); `skin > 0` is the production setting. +A frozen list for a fully compiled rollout goes through the integrator seam: +`MD(ff, ..., integrator=LangevinVerletIntegrator(ff, ..., rebuild=False))`. + +Migration: `MD(rebuild_every=)` and `NeighborListHook` are removed — +the cadence lives on the list now. + +A pure-GPU compiled bulk run composes the primitives directly — compile the +force field, keep the rebuild eager: + +```python +from molix import Compiler +from molix.md import MD, LennardJonesCutForceField, MaxwellBoltzmann, NeighborList + +nl = NeighborList(cell=cell, cutoff=2.5 * sigma, positions=pos, skin=1.02) +ff = LennardJonesCutForceField(epsilon=eps, sigma=sigma, neighbors=nl).to("cuda", torch.float64) +ff = Compiler(cuda_graphs=True)(ff) # or Compiler(fullgraph=True) +md = MD(ff, mass=39.95, dt=4.0, gamma=0.0, # γ=0 → NVE; the list owns the cadence + dtype=torch.float64, device="cuda") +vel = MaxwellBoltzmann(39.95, n_atoms=len(pos)).sample(172.0, seed=1) +state = md.run(pos, vel, n_steps=25_000) +``` + +See `benchmarks/verify_md_ljcut_nve.py` for the full melt-benchmark version +with energy-conservation checks. + +A batch `TensorDict` can carry the live neighbour buffers directly — +`build(batch)` binds `edges.edge_index` / `edges.shifts` **by reference**, so +every in-place rebuild is visible through the batch with no re-binding: + +```python +nl = NeighborList(cell=cell, cutoff=r_cut, positions=pos, skin=1.0) +nl.build(batch) # bind + first build; returns the same batch +out = potential(nl.build(batch)) # composes with forward(td) -> td + +for _ in range(n_steps): + ... + nl.update(batch) # per step: policy-gated rebuild +``` + +Note the two-statement idiom: `nl.build(batch).update(batch)` would call +**`TensorDict.update`** (a silent self-merge), not the neighbour policy — +`build` returns the batch for pipeline composition, so the policy call goes +through the list. + +For large condensed-phase systems, `bin=` selects a pure-torch binned +(cell-list) O(N) build instead of the O(N²) kernel: `bin=0.0` picks the +automatic `r_build / 2` size (LAMMPS `nbin_standard`), a positive float is an +explicit perpendicular bin thickness in Å. The bin size is a cost knob, never +a physics knob — both backends produce the identical edge set. Measured at +N=4096 (CPU, `OMP_NUM_THREADS=4`, 2026-08-09): binned 0.037 s vs kernel +0.450 s per rebuild; at 48 threads the tiny-op OpenMP overhead inverts the +ratio, so pin the thread count when profiling. No timing threshold is +asserted anywhere. + +## Observing a run: MD hooks + +`MDRunner` drives a small, MD-specific hook protocol (`MDHook`) — these are +*not* `Trainer` hooks. Built-ins: + +- `TrajectoryHook` — strided frames to a `.pt` (+ optional extended-XYZ), + sharded to disk so host memory stays bounded. +- `MDCheckpointHook` — restartable `(pos, vel, step)` every N steps, written + atomically; doubles as a heartbeat line in the log. + +A custom hook overrides any of `on_run_start` / `on_step_start` / +`on_step_end(runner, step, obs)` / `on_run_end`; `obs` is a typed +`MDObservables` (pos, vel, forces, potential, kinetic, total, temperature). +A hook that acts every N steps declares `cadence = N` so `run(chunk=...)` +can refuse a chunking that would silently skip it. + +## Custom integrators + +`MD(integrator=...)` accepts any constructed `Integrator` subclass. A +conforming subclass implements `advance` (one eager step) and `rollout` +(compile-friendly loop) and inherits `initial` / `advance_n`; override the +`removed_dof` property if your scheme thermostats all 3N degrees of freedom +(the temperature estimator reads it — Langevin reports 0, NVE 3). diff --git a/docs/molix/user-guide/profiling.md b/docs/molix/user-guide/profiling.md new file mode 100644 index 0000000..1efdfe8 --- /dev/null +++ b/docs/molix/user-guide/profiling.md @@ -0,0 +1,519 @@ +# Profiling + +`molix.profiler` is a suite of five standalone measurement tools. Each one +answers a single "where does the time go?" question about one layer of the +stack, and none of them asks you to wire up a real training run first. + +| Profiler | Question it answers | +| --- | --- | +| `TaskProfiler` | How long does one preprocessing task take per sample? | +| `ModuleProfiler` | How long does one `nn.Module` take per forward and backward pass? | +| `DataLoaderProfiler` | How long does the consumer block waiting for the next batch? | +| `TrainerProfiler` | How much per-step time does `Trainer` add on top of a bare loop? | +| `DatasetProfiler` | What is in my dataset, and what does reading one sample cost? | + +All five have the same shape: **configuration goes into the constructor, data +goes into `run()`, and `run()` returns a result dataclass that knows how to +print itself.** + +```python +from molix.profiler import DatasetProfiler + +profiler = DatasetProfiler(n_samples=200) # configuration +result = profiler.run(dataset) # data +result.print_report() # prints to stdout, returns None +``` + +Every field on a result object is a plain attribute, so anything the report +prints can also be asserted on in a script — `result.avg_num_neighbors` from +`DatasetResult`, `result.timing.p95_ms` from `TaskResult`, +`result.overhead_ms_per_step` from `TrainerResult`, and so on. + +This suite is not the same thing as `ProfilerHook` in `molix.hooks`. That hook +wraps `torch.profiler` around a *live* training run and writes a Chrome trace +file; the profilers here run outside training, on components you hand them +directly. + +## Reading the Reports + +A few conventions are shared by every report. + +**Warmup.** The first few iterations of anything in PyTorch are unrepresentative: +memory allocators grow their pools, kernels are selected and cached, and +memory-mapped file pages are faulted in from disk. Each profiler therefore takes +an `n_warmup` count of iterations that are executed but discarded before timing +starts. + +**Wall-clock time.** Unless a report says otherwise, times are wall-clock +milliseconds measured with `time.perf_counter` — real elapsed time, not CPU +time. `mean` / `std` / `p50` / `p95` are the mean, standard deviation, median +and 95th percentile over the measured iterations. `p95` matters more than the +mean for anything that stalls: a loader that is fast on average but occasionally +blocks for 200 ms will show a mean close to zero and a large `p95`. + +**The `Data:` line.** Reports name their input by calling `describe()` on it if +it has one (`MockSource` and `MockBatch` do), and otherwise fall back to the +class name — a `CachedDataset` simply prints as `CachedDataset`. + +**`[WARN]` lines.** Diagnostics are printed, never raised. A profiler that +notices something suspicious — a skewed size distribution, a missing pointer, a +`NaN` label — appends a `[WARN]` line to the report and keeps going. Exceptions +are reserved for inputs the profiler genuinely cannot measure. + +## Synthetic Inputs + +When you want a measurement before a dataset exists, `molix.profiler` ships +three stand-ins. + +`MockSource` implements the `DataSource` protocol — `__len__` plus +`__getitem__` returning a flat sample `dict` with `Z` `(N,)` (atomic numbers) +and `pos` `(N, 3)` (Cartesian positions). Atom counts are either fixed or drawn +from an inclusive `(lo, hi)` range. + +`MockBatch` is a callable factory returning a post-collate nested `TensorDict` +with the `atoms` / `edges` / `graphs` namespaces the models expect. Atom, edge +and graph counts are fixed or drawn from ranges on every call, which is how you +stress-test a module against variable input shapes. + +```python +from molix.profiler import MockBatch, MockSource + +source = MockSource(n_samples=500, n_atoms=(5, 20), seed=0) +sample = source[0] # {"Z": (N,), "pos": (N, 3)} + +factory = MockBatch(n_atoms=(32, 96), n_edges=(100, 600), n_graphs=4, seed=0) +batch = factory() # nested TensorDict +``` + +Both are **seed-reproducible**: with a seed set, two instances constructed in +the same process produce the same shape sequence, because the size draws come +from the instance's own `random.Random` rather than the process-global `random` +module. Leave `MockBatch(seed=None)` unseeded if you deliberately want a +different shape sequence per instance. + +`MockModel` is the third stand-in: an `nn.Module` that honours the encoder +contract (reads `atoms.pos`, writes `atoms.node_features` `(N, 1, n_features)`) +with one scalar multiply and one scalar `Parameter`. Its compute is negligible, +which is exactly what `TrainerProfiler` needs. Its companion loss, +`mock_node_feature_loss`, sums those features so `backward()` reaches the +parameter. + +## TaskProfiler + +A *task* is one step of the preprocessing pipeline (`molix.data.task`): +`SampleTask` transforms one sample dict, `DatasetTask` first fits global +parameters over the whole dataset and then transforms each sample, and +`BatchTask` transforms an already-collated batch. `TaskProfiler` times +`task.execute(...)` on its own, with no DataLoader, workers or collation in the +picture. + +```python +from molix.data.tasks import NeighborList +from molix.profiler import MockSource, TaskProfiler + +source = MockSource(n_samples=500, n_atoms=(5, 20), seed=0) + +result = TaskProfiler(NeighborList(cutoff=5.0)).run(source, n_samples=100, n_warmup=10) +result.print_report() + +print(result.timing.p95_ms) +``` + +The task is the constructor argument; the sample provider is the `run()` +argument. Any object with `__len__` and `__getitem__` works — a `MockSource`, a +real `DataSource`, or a dataset. Indices wrap modulo the source length, so +`n_samples` may exceed the number of distinct samples available. + +`TaskResult` carries `task_name`, `task_id` (the task's cache-key identity, e.g. +`nlist:cut=5.0:max=512:pbc=False:sym=True`), the `timing` statistics in +milliseconds, `n_samples`, and `data_description`. The report is a single +"Execute time" row with mean, std, p50, p95, min and max. + +Two things to know before pointing it at a `DatasetTask`. First, the profiler +calls `fit()` on the *entire* source before timing begins, mirroring what +`PipelineSpec.run` does — that materialises every sample in memory, so use a +subset for large sources. Second, the fit itself is not timed; only the +per-sample `execute` is. For a `BatchTask`, whatever `source[i]` returns is fed +straight to `execute`, so the "source" must already yield collated batches. + +## ModuleProfiler + +`ModuleProfiler` times any `torch.nn.Module` in isolation: no `Trainer`, no +`DataModule`, no hooks. + +```python +import torch +from molix.profiler import MockBatch, ModuleProfiler + +optimizer = torch.optim.Adam(model.parameters(), lr=1e-4) +profiler = ModuleProfiler(model, loss_fn=my_loss, device="cuda:0", optimizer=optimizer) + +factory = MockBatch(n_atoms=(32, 96), n_edges=(100, 600), n_graphs=4, device="cuda:0") +result = profiler.run(factory, n_steps=200, n_warmup=10) +result.print_report() +``` + +`run()` accepts a `MockBatch` factory (called once per step), a list of +pre-built batches (cycled), or a `DataLoader` or other iterable (cycled). +Batches are moved to `self.device` before each step, so generating them on the +target device — as above — avoids paying for a host-to-device copy inside the +measurement. Passing `loss_fn` adds the backward pass; passing `optimizer` as +well (one that already wraps `module.parameters()`) adds the optimizer step, +making the measured step identical in structure to what `Trainer` runs. The +optimizer step is only taken when a `loss_fn` is present. + +Three separate measurements go into one `ModuleResult`. + +**Component breakdown** — forward, backward and optimizer timed individually. +On CUDA each component is bracketed by a pair of `torch.cuda.Event` markers and +resolved by synchronizing, which serialises CPU and GPU. Read these as +*relative* costs and upper bounds, not as throughput. + +**Full-step block** — `n_steps` steps run back to back with no synchronization +in between, followed by a single one at the end. This is `wall_ms_per_step`, the +number that reflects real throughput. It is measured on CUDA only; on CPU the +field stays `None` and the report omits the block. + +**Op breakdown** — a short `torch.profiler` window, also CUDA only. Summing +per-kernel *device self-time* (time inside the kernel itself, excluding anything +it called) gives `gpu_active_ms_per_step`: on a single stream kernels run +serially, so the sum is the true GPU-busy time. A single event pair around the +whole block would instead report the GPU-timeline span, which already contains +the idle gaps and therefore just reproduces the wall time. + +From those two numbers comes the most useful line in the report: + +$$ +\mathrm{launch\ bound\ \%} = +\frac{t_\mathrm{wall} - t_\mathrm{gpu\ active}}{t_\mathrm{wall}} \times 100 +$$ + +where $t_\mathrm{wall}$ is `wall_ms_per_step` and $t_\mathrm{gpu\ active}$ is +`gpu_active_ms_per_step`, both in milliseconds per step. This is the percentage +of each step during which the GPU sits idle because the CPU has not yet +enqueued ("launched") the next kernel. The report converts it +into a one-line verdict: at or above 60 % it prints `LAUNCH-BOUND — too many +tiny kernels / batch too small`, at or above 30 % `partially launch-bound`, +and below that `compute-bound`. A launch-bound model does not get faster from a +faster GPU; it gets faster from bigger batches, fused kernels, CUDA graphs or +`torch.compile`. + +The rest of `ModuleResult`: `peak_memory_mb` (peak CUDA allocation during the +forward pass, `0` on CPU), `throughput_atoms_per_sec` and +`throughput_graphs_per_sec` (computed from the **forward** mean only, using the +atom and graph counts read out of each batch), `n_params`, +`op_calls_per_step`, and the pre-rendered `op_table` of the top operators by +self-CUDA time. Pass `submodules=True` to add a second, forward-only hooked +pass that attributes time to the module's top-level named children — +`ModuleList` / `Sequential` / `ModuleDict` containers are expanded into their +entries, since the container itself is never called. The gap between the sum of +the children and the reported forward mean is work done inline in `forward` +rather than inside a named child. + +For sub-modules that do not take a batch `TensorDict` — a radial basis function +that takes a distance tensor, say — use `run_fn` with explicit callables: + +```python +distances = torch.rand(512, device="cuda:0") + +result = ModuleProfiler(rbf, device="cuda:0").run_fn( + forward_fn=lambda: rbf(distances), + backward_fn=lambda out: out.sum(), + n_steps=200, + label="edge_dist E=512", +) +``` + +`run_fn` reports the forward/backward breakdown and memory only; throughput is +reported as zero, because raw tensors carry no atom or graph counts. + +## DataLoaderProfiler + +`DataLoaderProfiler` measures **stall time**: how long the consumer of a +`DataLoader` blocks waiting for the next batch to arrive. It uses the +inter-batch gap technique — the wall-clock elapsed from the moment the loop +finishes with batch *i* to the moment batch *i+1* is handed over. That interval +covers worker scheduling, sample reads, collation and pinning, which is exactly +the time a training step cannot overlap with compute. + +```python +from molix.profiler import DataLoaderProfiler, MockSource + +profiler = DataLoaderProfiler(batch_size=32, num_workers=4, pin_memory=True) +result = profiler.run(MockSource(n_samples=2000, n_atoms=(5, 20), seed=0), n_batches=100) +result.print_report() +``` + +Configuration is the DataLoader configuration you want to test: `batch_size`, +`num_workers` (worker subprocesses; `0` means load in the main process), +`pin_memory` (allocate host buffers in page-locked memory so the copy to the GPU +can run asynchronously), `persistent_workers` (keep workers alive across epochs; +silently forced off when `num_workers == 0`), `target_schema` (how labels are +routed during collation, defaulting to `DEFAULT_TARGET_SCHEMA`), and an optional +`PipelineSpec` whose `batch_nodes` are applied inside `collate_fn`. + +`run()` resolves its argument in three ways. An existing `DataLoader` is used +as-is and your constructor settings only affect what the report *says*. A +`torch.utils.data.Dataset` is wrapped in a new DataLoader built from those +settings. Anything else is treated as a `DataSource`: every sample is +materialised, written to a `PackedCache` file in a fresh temporary directory, +and reopened as a `CachedDataset`, so the measurement exercises the same +cache-file path a real workflow uses. Be aware that this last route writes to +disk and holds the whole source in memory while packing. + +With `num_workers > 0` the loader is built with the `spawn` start method and a +top-level picklable collate object, so the profiler works under `spawn` and +`forkserver` where a closure would fail to pickle. + +`DataLoaderResult` reports `load_time` (the stall statistics), +`throughput_graphs_per_sec` and `throughput_atoms_per_sec`, and the per-batch +size distributions `batch_graph_stats` and `batch_atom_stats` — the latter is +the one to watch when batches are size-heterogeneous, because a large `std` +means padded batches waste compute. The report ends with a `[WARN]` when +`p95 > 3 × mean`, which points at worker stalls or collation spikes and suggests +more workers or `persistent_workers=True`. + +If the loader yields no batch beyond the warmup window, `run()` raises +`RuntimeError` rather than reporting statistics over an empty sample. + +## TrainerProfiler + +The other four profilers measure work you asked for. `TrainerProfiler` measures +the work you did not: the per-step cost of the `Trainer` machinery itself — +`Step` protocol dispatch, hook calls, `batch_to` device transfer, `TrainState` +writes, and the optimizer/scheduler/eval-cadence bookkeeping. + +Isolating that cost needs two tricks. First, the model must be nearly free, or +its FLOPs drown everything else out; the default is `MockModel`. Second, a bare +training loop still costs something — forward, backward, `optimizer.step()` — +and that cost is not the Trainer's fault. So the profiler measures a raw loop +with no Trainer at all and subtracts it. + +```python +from molix.profiler import TrainerProfiler + +result = TrainerProfiler(device="cpu").run(n_steps=2000, n_warmup=50, top=15) +result.print_report() + +print(result.overhead_ms_per_step) +``` + +The constructor takes `model` (default `MockModel`), `loss_fn` (default +`mock_node_feature_loss`), `device`, and `hooks`. Hooks are the interesting +knob: run once with none to price the bare loop, then again with the hook list +you actually train with to price hook dispatch. + +```python +from molix.hooks import StepSpeedHook + +TrainerProfiler(hooks=[StepSpeedHook()], device="cpu").run(n_steps=2000).print_report() +``` + +`run(n_steps, n_warmup, batch, top)` drives a real `Trainer` for one epoch over +a fixed batch replayed `n_steps` times. `batch` defaults to a `MockBatch` with +32 atoms, 128 edges and 4 graphs on the configured device; pass your own +`TensorDict` to see how the overhead behaves at your batch size. The internal +optimizer is SGD with `lr=0.0`, so nothing is learned — weights stay put and +only machinery is measured. `top` caps the number of hotspot rows kept. + +Four passes are made after the warmup: the raw loop timed, the raw loop under +`cProfile`, the Trainer loop timed, and the Trainer loop under `cProfile`. +Timing and profiling are deliberately separate runs, because `cProfile` adds +per-call overhead that would distort the wall-clock numbers. + +The report opens with three lines: + +- `Raw loop (baseline)` — `baseline_ms_per_step`, the irreducible per-step work: + `zero_grad` → forward → loss → `backward` → `optimizer.step()`. +- `Trainer loop` — `wall_ms_per_step` for the same work driven through + `Trainer`, with `steps_per_sec = 1000 / wall_ms_per_step` alongside. +- `Trainer overhead` — `overhead_ms_per_step`, the difference (clamped at zero), + and what percentage of the loop it represents. + +Below them is the hotspot table, headed *Trainer self-time minus raw-loop +baseline (per step)*. A function's **self-time** is the time spent in its own +body, excluding the functions it calls, as attributed by `cProfile`. For each +`(file, line, function)` the profiler subtracts the raw-loop self-time from the +Trainer-loop self-time and keeps the positive remainders, sorted descending: + +| column | meaning | +| --- | --- | +| `func` | `file:line(function)`, truncated to 48 characters | +| `added_us` | microseconds per step this function adds *over* the baseline | +| `gross_us` | microseconds per step it costs in the Trainer run, unsubtracted | +| `calls` | calls per step | + +The subtraction is what makes the table readable. Autograd, the model forward +and the optimizer appear in both runs at roughly equal cost and cancel out, +leaving framework machinery — the loop body, the `Step` wrapper, device +transfer, hook dispatch, state writes — at the top. + +The same numbers are available on `TrainerResult` as `n_steps`, `device`, +`n_hooks`, `wall_ms_per_step`, `steps_per_sec`, `baseline_ms_per_step`, +`overhead_ms_per_step`, `model_name` and the `hotspots` list of dicts. + +## DatasetProfiler + +`DatasetProfiler` characterises the data itself rather than the loop around it. +Point it at anything that yields **flat sample dicts** — the raw-sample tier of +the two-tier data contract, so `sample["Z"]`, `sample["pos"]`, +`sample["edge_index"]`, `sample["targets"]["U0"]` — and it answers four +questions at once: how big are the records, what does one access cost, how are +the fields laid out, and what do the labels look like. + +```python +import torch + +from molix.data.cache import PackedCache +from molix.data.dataset import CachedDataset +from molix.profiler import DatasetProfiler + + +def ring(i: int) -> dict: + """One literal sample: a ring molecule of 2–5 carbon atoms.""" + n = 2 + i % 4 + src = torch.arange(n) + dst = (src + 1) % n + return { + "Z": torch.full((n,), 6), + "pos": torch.arange(3 * n, dtype=torch.float32).reshape(n, 3), + "edge_index": torch.cat( + [torch.stack([src, dst], dim=1), torch.stack([dst, src], dim=1)] + ), + "edge_dist": torch.full((2 * n,), 1.5), + "targets": {"U0": torch.tensor([float(i)])}, + } + + +sink = "/tmp/ring-cache.pt" +PackedCache(sink).save([ring(i) for i in range(100)], overwrite=True) +dataset = CachedDataset(sink) + +result = DatasetProfiler(n_samples=20).run(dataset) +result.print_report() + +assert result.counts_exact # size stats cover all 100 records +assert result.avg_num_neighbors == 2.0 # 700 edges / 350 atoms +``` + +You can also hand `run()` a `MockSource(n_samples=100, seed=0)` or a plain +`list[dict]` directly, without packing anything. Those carry no packed pointers, +so the size statistics come from the sampled records instead — see below — and +`MockSource` samples have no edges at all, so the edges row is simply absent +from the size section. + +Configuration is `n_samples` (how many records the sampled path may read), +`stride` (the step between inspected indices, so `stride > 1` spreads the sample +across an ordered dataset instead of reading a prefix), and `n_warmup`. Indices +are `range(0, len(data), stride)` truncated to `n_samples`. + +### The exact fast path + +Cache-backed datasets store their samples packed: every key is concatenated +across all records into one big tensor, with `atom_ptr` and `edge_ptr` cumulative +sum ("cumsum") vectors marking where each record's slice begins. Per-record +counts are then just `ptr[i+1] - ptr[i]`, one vectorised subtraction over the +whole file with no sample unpacked. + +So when the object exposes `atom_counts`, `edge_counts`, `avg_num_neighbors`, +`max_atoms` and `max_edges` — `CachedDataset`, `MmapDataset`, and `SubsetDataset` +which remaps them to its own indices — the profiler reads those properties +directly. Size statistics then describe **every** record in the dataset, not the +sampled subset, and the result reports `counts_exact=True`. The report prints +the flag next to the neighbour count so you always know which regime produced +the numbers. + +`avg_num_neighbors` deserves a definition, since models consume it. It is + +$$ +\langle |N(i)| \rangle = \frac{E_\mathrm{total}}{N_\mathrm{total}} +$$ + +where $E_\mathrm{total}$ is the number of edges and $N_\mathrm{total}$ the +number of atoms, both summed over the whole dataset — or, on a +`SubsetDataset`, over that split only, since a training subset must not peek at +validation data. The ratio is dimensionless. With +`NeighborList(symmetry=True)`, the default, the neighbour graph is fully +bidirectional, so this ratio is the mean number of neighbours per atom, which is +the normalisation constant MACE and Allegro divide their aggregated messages by. + +Without a packed cache — a plain `list[dict]`, a `MockSource` — none of those +properties exist. The profiler then derives size statistics from the records it +sampled, sets `counts_exact=False`, and says so in a `[WARN]` line. The numbers +are still useful; they are just estimates from a subset. + +One legitimate half-way case: a cache built by a pipeline that never ran +`NeighborList` has atoms but no edges, so `edge_counts` raises `ValueError`. The +profiler catches it, sets `edge_stats=None`, drops the edge row from the report, +and appends a `[WARN]` carrying the underlying message. Missing edges are a fact +about your pipeline, not a profiler failure. + +### The sampled path + +Some things cannot be read off a pointer vector and genuinely require touching +records one at a time: + +- `cold_access_ms` — the very first `data[i]`, including memory-map page-in from + disk. On an `MmapDataset` this is usually much larger than the steady state. +- `access_ms` — steady-state `__getitem__` latency after `n_warmup` discarded + accesses, as a full `TimingStat`. +- `sample_bytes` — the leaf-tensor footprint of one record: the sum of + `numel() * element_size()` over its tensor leaves. Non-tensor leaves + contribute nothing. +- `est_total_mb` — `sample_bytes.mean × n_total / 1e6`, i.e. what a full in-RAM + materialisation would cost in megabytes (10⁶ bytes). It is an extrapolation + from the sampled mean, so it is only as good as the sample on a skewed dataset. +- `targets` — for every leaf under `targets.` that is a one-element tensor or a + plain number, a `TargetStat` with mean, std, p50, p95 over the **finite** + values, plus `min`, `max` and `n_nonfinite`. Non-finite entries are counted and + then excluded from the moments, so a single `NaN` row cannot poison the column. + +### Field layout + +The `Fields` section lists one `FieldSpec` per key: the dotted key path, the +packing `axis`, the `dtype`, and the trailing shape after the packing axis +(`(3,)` for `pos` `(N, 3)`, `()` for `Z` `(N,)`). + +| axis | meaning | +| --- | --- | +| `atom` | concatenated along dim 0 — one row per atom | +| `edge` | concatenated along dim 0 — one row per edge | +| `graph` | stacked on a new leading dim — one entry per record; `extra_shape` is the full per-sample shape | +| `scalar` | a non-tensor Python value; `dtype` is the type name and `extra_shape` is `()` | + +When the dataset exposes `packed_view()`, this comes straight from the cache's +`payload["schema"]` — the layout inferred across *all* records at packing time — +and `fields_exact=True`. Otherwise the layout is inferred from the sampled +records (leading dimension matches the atom count, else the edge count, else +per-graph), `fields_exact=False`, and a `[WARN]` notes that axes and trailing +shapes could differ on unsampled records. + +If the dataset has a `stats()` method, its fitted `DatasetTask` state is +collected into `task_states` and the names are listed in the report — that is +how you see at a glance that this cache was baked with, say, `AtomicDress`. + +### Diagnostics and errors + +The dividing line is strict. **Diagnostics never raise.** Non-finite labels, +inexact counts, inexact fields, a missing edge pointer, and an atom-count skew of +`p95 / p50 > 3` (padded batches will waste compute; consider a token-budget +sampler) are all `[WARN]` lines under the report's closing rule. + +`ValueError` is reserved for inputs that cannot be profiled at all, and each +message says what to pass instead: `n_samples <= 0` or `stride <= 0` at +construction; an empty dataset, or an object with no `__len__` / `__getitem__`, +at `run()`. + +No unit conversion is performed and no units are guessed. Positions, distances +and every target are printed exactly as the dataset stores them — if your +positions are in Ångström and your energies in eV, that is what you are reading. + +## Where to Start + +Working outwards from the data usually converges fastest. Run `DatasetProfiler` +first to learn what the records look like and whether sizes are skewed; then +`TaskProfiler` on any preprocessing step that looks expensive; then +`DataLoaderProfiler` to see whether batch production keeps up; then +`ModuleProfiler` to find out whether the model is compute-bound or launch-bound; +and finally `TrainerProfiler` if per-step time still exceeds what the model and +the loader together explain. diff --git a/docs/molix/user-guide/trainer.md b/docs/molix/user-guide/trainer.md index 058fd47..3240d88 100644 --- a/docs/molix/user-guide/trainer.md +++ b/docs/molix/user-guide/trainer.md @@ -23,16 +23,22 @@ state = trainer.train(datamodule, max_epochs=100) ## Loss Function Contract -The default train and eval steps call: +The default train and eval steps pass the batch through unmodified: ```python -predictions = model(...) +predictions = model(batch) loss = loss_fn(predictions, batch) ``` -For plain `dict` batches, keys named `targets` and `extras` are not forwarded to -the model. For non-dict batches such as `GraphBatch`, the whole batch is passed -to `model(batch)`. +The whole batch object goes to `model(batch)` — no key filtering, no unpacking +into keyword arguments — so the model and the loss see exactly the same object. +For the collated nested `TensorDict` that means the model reads its inputs with +tuple keys (`batch["atoms", "Z"]`) and the loss reads its targets from the same +batch (`batch["graphs", "energy"]`). + +`loss_fn` must return a scalar tensor; the step calls `.backward()` on it (after +dividing by `accumulate_grad_batches`) and records the un-scaled value in +`state["train"]["loss"]`. ## Loop Control diff --git a/docs/molpot/user-guide/gradients.md b/docs/molpot/user-guide/gradients.md index fa0ec61..00c7f32 100644 --- a/docs/molpot/user-guide/gradients.md +++ b/docs/molpot/user-guide/gradients.md @@ -42,6 +42,59 @@ print(forces.shape) # torch.Size([10, 3]) `energy_fn` must recompute any position-derived geometry (edge vectors, distances) inside itself so the gradient flows `pos → geometry → energy`. +## Batch-level force-pass kernels (批级力学传递内核) + +`ForceDerivation` / `force.py` is the **tensor-level** layer: you hand it +`energy_fn(pos) -> scalar` and get `forces (N, 3)` back. A potential, though, +works on a post-collate batch — it has to own the position leaf, write +`graphs.energy` / `atoms.forces` in place, and decide whether the returned +energy stays attached. That is the **batch-level** layer, +`molpot.derivation.kernels`: + +| Layer | Module | Signature | Owns | +|-------|--------|-----------|------| +| tensor | `molpot.derivation.force` | `energy_fn(pos) -> scalar` | the only `torch.autograd.grad` / `torch.func.grad` calls | +| batch | `molpot.derivation.kernels` | `energy_core(batch) -> batch` | position-leaf ownership, key writes, `detach_energy` | + +```python +from molpot.derivation import func_force_pass, grad_force_pass + +# one energy forward + torch.autograd.grad on the position leaf +batch = grad_force_pass(self._write_energy, batch, detach_energy=False) + +# single torch.func.grad(..., has_aux=True) pass — fullgraph-compilable +batch = func_force_pass(self._write_energy, batch) +``` + +The kernels **compose** `force.py` — they never re-derive a gradient — so the +backend boundary of the table above still holds: `grad_force_pass` for cuEq +fused kernels (MACE-shaped), `func_force_pass` for pure-PyTorch graphs (PiNet) +that want `torch.compile(fullgraph=True)`. `PiNetPotential`, `GradMode` and +`FuncMode` all bind these; a potential must not hand-roll a third pass body. + +`grad_force_pass(energy_core=None, ...)` skips the forward and differentiates +an energy already materialised on a live position leaf — that is what keeps +`EnergyReadout(backward=True)` + `ForceReadout` at a single model forward. + +### `detach_energy` (three states) + +Who owns the position leaf decides whether the returned `graphs.energy` can +still be part of a loss: + +| `detach_energy` | Behaviour | Caller | +|-----------------|-----------|--------| +| `False` | never detach — energy stays attached for an energy loss | `PiNetPotential`, `GradMode` | +| `True` | always detach — energy is a reported quantity only | inference / logging | +| `None` (default) | detach **iff the kernel created the position leaf**, i.e. the caller had no graph to lose | MACE-style `get_outputs` | + +`func_force_pass` has no such knob: its only callers want the energy attached, +and the fused kernels that need leaf-detaching cannot use functorch anyway +(pytorch#170834). + +Units throughout: positions Å, energies eV, forces eV/Å. The kernels only +differentiate whatever the energy core wrote, so a core on another unit system +produces mismatched forces silently. + ## With PotentialComposer `PotentialComposer` derives forces when positions are present in `data`: diff --git a/docs/molzoo/index.md b/docs/molzoo/index.md index a20ba2c..ad9b6b1 100644 --- a/docs/molzoo/index.md +++ b/docs/molzoo/index.md @@ -12,7 +12,7 @@ physics heads, derivation, and composition stay in `molpot`. | MACE | encoder | `src/molzoo/specs/mace.md` | Writes `atoms.node_features` | | PiNet | encoder + temporary potential | `src/molzoo/specs/pinet2.md` (source tree) | Package `molzoo.pinet/` (`encoder`, `potential`, `properties`); long-term potential home is molpot | | MACE-OMOL | full energy/force | [`specs/mace_omol.md`](specs/mace_omol.md) | Lazy import; not encoder-only | -| Sonata | composition | lives in **`molpot.composition`**, not molzoo | `build_sonata` | +| Sonata | composition | lives in **`molpot.composition`**, not molzoo | `Sonata.from_encoder` | ## Documentation Layout diff --git a/docs/molzoo/specs/mace_omol.md b/docs/molzoo/specs/mace_omol.md index 4e2f9c1..08d7adb 100644 --- a/docs/molzoo/specs/mace_omol.md +++ b/docs/molzoo/specs/mace_omol.md @@ -1,12 +1,13 @@ # MACEOMol Specification -This page is the implementation contract for `molzoo.mace_omol`. It is not a -tutorial; use the MolZoo user guide for theory narrative and worked examples. +This page is the implementation contract for `molzoo.mace.variants.MACEOMol`. +It is not a tutorial; use the MolZoo user guide for theory narrative and worked +examples. | Field | Value | |-------|-------| -| Module | `molzoo.mace_omol` | -| Entry point | `MACEOMol` (plain `nn.Module`; constructor kwargs, no pydantic Spec) | +| Module | `molzoo.mace.variants` (in the `molzoo.mace` package) | +| Entry point | `MACEOMol` — a thin, keyword-compatible alias over `molzoo.mace.potential.MACEPotential`. The constructor still takes the OMOL keywords directly; internally it builds a `molzoo.mace.spec.MACEOMolSpec`. Weights: `MACEPotential.from_checkpoint` with the `molzoo.mace.checkpoint.OMOL_REMAP` preset; `load_omol_state_dict` stays as a back-compat free function in `molzoo.mace.variants`. | | Paper | Batatia et al., "MACE: Higher Order Equivariant Message Passing Neural Networks for Fast and Accurate Force Fields", NeurIPS 2022. OMol25 / MACE-omol-0 foundation model. | | arXiv | https://arxiv.org/abs/2206.07697 | | DOI | not applicable | @@ -20,9 +21,14 @@ tutorial; use the MolZoo user guide for theory narrative and worked examples. `l_max=3`, 3 residual interactions, product `correlation=2`, 83 elements, single `omol` head, with `total_charge` / `total_spin` conditioning, 52.7M parameters). -It **owns**: the full energy/force forward of MACE-OMOL, assembled from generic -`molrep` / `molpot` blocks (no MACE prefix), and the `load_omol_state_dict` -converter that imports official weights (after `mace.cli.convert_e3nn_cueq`). +It **owns**: the OMOL preset of the full energy/force forward — the block graph +in `molzoo.mace.encoder`, the energy/force pipeline in +`molzoo.mace.potential`, and the `OMOL_REMAP` key dialect in +`molzoo.mace.checkpoint` that imports official weights (after +`mace.cli.convert_e3nn_cueq`), reached through `load_omol_state_dict` or +`MACEPotential.from_checkpoint`. The blocks themselves come from `molrep` / +`molpot` — some generic (`ResidualInteraction`, `EquivariantProductBasis`), +some in the MACE-only namespaces (`molrep.readout.mace`). It does **not** own: the building blocks themselves (they live in `molrep` / `molpot` and are reused by other models), neighbor-list construction, dataset / @@ -43,8 +49,8 @@ not an encoder-only feature extractor: its `forward` writes `graphs.energy` and | In | `atoms.pos` | `(N, 3)` | float (`config.ftype`) | Cartesian positions; forces are `-∂E/∂pos` | | In | `atoms.batch` | `(N,)` | long | Graph membership index `0..B-1` | | In | `edges.edge_index` | `(E, 2)` | long | `[:,0]`=source/sender, `[:,1]`=target/receiver (MolNex convention) | -| In | `graphs.total_charge` | `(B,)` | long | Optional; per-graph total charge. Absent → neutral (0) | -| In | `graphs.total_spin` | `(B,)` | long | Optional; per-graph spin. Absent → singlet (0) | +| In | `graphs.total_charge` | `(B,)` | long | Optional; per-graph total charge in units of `e`. Absent → neutral (`0`) | +| In | `graphs.total_spin` | `(B,)` | long | Optional; per-graph spin channel. Absent → `1` (singlet multiplicity `2S+1 = 1`, all electrons paired). **Not `0`** — `MACEPotential._condition_charge_spin` fills `torch.ones(...)`, and spin `0` would index an untrained embedding row (`spin_offset = 0`) and return garbage | `edges.edge_diff` / `edges.edge_dist` are **not** consumed: `forward` recomputes edge vectors from `atoms.pos` so the energy is differentiable w.r.t. positions. @@ -54,12 +60,14 @@ edge vectors from `atoms.pos` so the energy is differentiable w.r.t. positions. | Direction | TensorDict path | Shape | Written when | Contract | |-----------|------------------|-------|--------------|----------| | Out | `graphs.energy` | `(B,)` | always | Per-graph total energy = E0 + scale·shift(readout) | -| Out | `atoms.forces` | `(N, 3)` | always | `F = -∂E/∂pos` via `molpot.derivation.ForceDerivation` | +| Out | `atoms.forces` | `(N, 3)` | always | `F = -∂E/∂pos` via `molpot.derivation.kernels.grad_force_pass` | `forward` mutates `td` in place (creating the `graphs` sub-dict if absent) and returns the same object. The raw-tensor entry point `MACEOMol.energy_forces(...)` returns `{"energy", "forces"}` and is used by the `scripts/omol_port/verify_*.py` -block/E2E checks. +block/E2E checks; the differentiable energy alone is +`MACEPotential.energy_core(positions, Z, edge_index, batch, num_graphs, +shifts=None, total_charge=None, total_spin=None)`. ## 3. Forward Contract @@ -67,11 +75,11 @@ block/E2E checks. | Symbol | Meaning | Code anchor | |--------|---------|-------------| -| $N, E, B$ | atoms, edges, graphs | `MACEOMol._compute_energy` | +| $N, E, B$ | atoms, edges, graphs | `molzoo.mace.potential.MACEPotential.energy_core` | | $Z_i$ | atomic number of atom $i$ | `atoms.Z` | | $\mathbf r_i$ | position of atom $i$ | `atoms.pos` | -| $s,t$ | sender / receiver of an edge | `edge_index[0]`, `edge_index[1]` | -| $\mathbf v_e=\mathbf r_t-\mathbf r_s$ | edge vector | `_compute_energy` | +| $s,t$ | sender / receiver of an edge | `edge_index[:,0]`, `edge_index[:,1]` | +| $\mathbf v_e=\mathbf r_t-\mathbf r_s$ | edge vector | `molzoo.mace.geometry.edge_vectors` | | $h_i$ | node features (irreps, `cue.ir_mul`) | `node_feats` | | $E_0$ | per-element reference + charge/spin readout | `e0` | | $q,\sigma$ | per-graph total charge / spin | `total_charge`, `total_spin` | @@ -127,9 +135,9 @@ $$ | Quantity | Shape | Code anchor | |----------|-------|-------------| -| readout | `(N,)` | `NonLinearBiasReadout` (`molrep.readout.scalar`) | +| readout | `(N,)` | `NonLinearBiasReadout` (`molrep.readout.mace`) | | scale_shift | `(N,)` | `GlobalRescale` (`molpot.heads.rescale`) | -| $\mathbf F$ | `(N,3)` | `ForceDerivation` (`forward`) / `autograd.grad` (`energy_forces`) | +| $\mathbf F$ | `(N,3)` | `molpot.derivation.kernels.grad_force_pass` (`forward`) / `molpot.derivation.force.autograd_forces_from_energy` (`energy_forces`) — both `torch.autograd.grad` | ## 4. Configuration Contract @@ -161,12 +169,12 @@ $$ | Radial MLP | `modules.radial.RadialMLP` | `molrep.interaction.RadialMLP` | matched | | Gated nonlinearity | e3nn `Gate` | `molrep.interaction.GatedNonlinearity` | matched | | Product basis | `EquivariantProductBasisBlock` (`original_mace=True`) | `molrep.interaction.EquivariantProductBasis` | matched | -| Non-linear readout | `NonLinearBiasReadoutBlock` | `molrep.readout.NonLinearBiasReadout` | matched | +| Non-linear readout | `NonLinearBiasReadoutBlock` | `molrep.readout.mace.NonLinearBiasReadout` (re-exported as `molrep.readout.NonLinearBiasReadout`) | matched | | Per-element E0 | `AtomicEnergiesBlock` | `molpot.heads.AtomicReferenceEnergy` | matched | | Scale/shift | `ScaleShiftBlock` | `molpot.heads.GlobalRescale` | matched | | Charge/spin embed | `GenericJointEmbedding` | `molrep.embedding.JointFeatureEmbedding` | matched | -| Forces | `autograd.grad` | `molpot.derivation.ForceDerivation` (`torch.func.grad`) | adapted | -| Weight import | e3nn `state_dict` | `load_omol_state_dict` (after `convert_e3nn_cueq`) | adapted | +| Forces | `autograd.grad` | `molpot.derivation.kernels.grad_force_pass` / `molpot.derivation.force.autograd_forces_from_energy` (both `torch.autograd.grad`; cuEq's fused ops are legacy `autograd.Function`s that `torch.func.grad` rejects) | matched | +| Weight import | e3nn `state_dict` | `molzoo.mace.checkpoint.OMOL_REMAP` via `load_omol_state_dict` / `MACEPotential.from_checkpoint` (after `convert_e3nn_cueq`) | adapted | ## 6. MolNex Adaptations @@ -174,10 +182,10 @@ $$ |----|------------|--------|------|------------| | A1 | PolynomialCutoff + trainable un-normalised Bessel (`eps=0`) | OMOL variant vs standard MACE | low | `scripts/omol_port/verify_radial.py` (7e-15) | | A2 | charge/spin via `JointFeatureEmbedding` added to node feats + into E0 | OMOL conditioning | low | `scripts/omol_port/verify_joint_embed.py` (0) | -| A3 | cue `"O3"` group everywhere (no `O3_e3nn`) | weights are converted into the cue-O3 twin, so O3 is the native target; O3 vs O3_e3nn CG differ only ~1.4e-8/op and the O3 twin already matches e3nn to 1.5e-8 — O3_e3nn would not reduce the residual and would add an `e3nn` dep | low | residual 7e-7 eV / 4.3e-6 eV·Å vs official (reimplementation accumulation, not a convention diff) — inside the 1e-4 bar (`mace-omol-port-02` ac-003) | -| A4 | TensorDict `forward` forces via `ForceDerivation` (`func.grad`); `energy_forces` via `autograd.grad` | compile-friendly molnex contract | low | `tests/test_molzoo/test_mace_omol.py` (1e-8 vs autograd) | -| A5 | edge convention `v=pos[t]-pos[s]`, `edge_index (E,2)→(2,E)` | MolNex collate schema | low | `tests/test_molzoo/test_mace_omol.py` | -| A6 | `RadialMLP` honors `config.ftype` | fp64-via-config without `.double()` | low | `tests/test_molzoo/test_mace_omol.py` (fp64) | +| A3 | cue `"O3"` group everywhere (no `O3_e3nn`) | weights are converted into the cue-O3 twin, so O3 is the native target; O3 vs O3_e3nn CG differ only ~1.4e-8/op and the O3 twin already matches e3nn to 1.5e-8 — O3_e3nn would not reduce the residual and would add an `e3nn` dep | low | residual 7e-7 eV / 4.3e-6 eV/Å vs official (reimplementation accumulation, not a convention diff) — inside the 1e-4 bar (`mace-omol-port-02` ac-003) | +| A4 | TensorDict `forward` forces via `molpot.derivation.kernels.grad_force_pass`; `energy_forces` via `molpot.derivation.force.autograd_forces_from_energy`. Both are `torch.autograd.grad`; the earlier `torch.func.grad` plan was dropped because cuEquivariance's fused ops register legacy `autograd.Function`s without `setup_context` | compile-friendly molnex contract | low | `tests/test_molzoo/test_mace/test_variants.py::TestMACEOMol::test_forward_matches_energy_forces` (1e-8 vs autograd) | +| A5 | edge convention `v=pos[t]-pos[s]`, `edge_index (E,2)` end to end (upstream's `(2,E)` transposed away at the port boundary) | MolNex collate schema | low | `tests/test_molzoo/test_mace/test_geometry.py::TestEdgeVectors`, `tests/test_molzoo/test_mace/test_variants.py::TestMACEOMol` | +| A6 | `RadialMLP` honors `config.ftype` | fp64-via-config without `.double()` | low | `tests/test_molzoo/test_mace/test_variants.py::TestMACEOMol` (fp64 fixtures) | | A7 | per-layer irreps + edge-mid (128) hardcoded to OMOL dims | faithful OMOL weight load | medium | non-OMOL `l_max≥2`+small `num_features` unsupported; tracked in `mace-omol-port-02` | ## 7. Validation Contract @@ -185,25 +193,54 @@ $$ ### 7.1 Research Reproduction The accepted accuracy bar is E/F within **1e-4** of official OMOL (operator -decision, 2026-06-21). The full model with official OMOL weights reproduces the -official cueq OMOL twin on a charged molecule to **7.0e-7 eV / 4.3e-6 eV·Å** -(`scripts/omol_port/verify_e2e.py`, RESULT: PASS) — three to four orders inside -the bar; the cueq twin itself matches e3nn OMOL to 1.5e-8 eV / 3.2e-8 eV·Å -(`scripts/omol_port/verify_omol_cueq_equiv.py`). The 7e-7 residual is molnex's -own reimplementation accumulation, **not** a CG-convention difference: O3 vs -O3_e3nn Clebsch-Gordan differ only ~1.4e-8/op and the O3 twin already aligns -with e3nn to 1.5e-8, so the e3nn-convention group is neither used nor needed -(A3). +decision, 2026-06-21). + +**Historical record (2026-06-21; oracles deleted in `b85d12f`, not +reproducible in-tree — Appendix A).** Full model with official weights vs the +official cueq OMOL twin, charged molecule: **7.0e-7 eV / 4.3e-6 eV/Å** +(`verify_e2e.py`, RESULT: PASS); the cueq twin vs e3nn OMOL: 1.5e-8 eV / +3.2e-8 eV/Å (`verify_omol_cueq_equiv.py`). Three to four orders inside the +bar. The run predates the 2026-08-07 loader fix (official `bessel_weights` +silently dropped, `bessel.freqs` left at init), so the 7e-7 eV includes that +~2.2e-7 Å⁻¹ perturbation; it is otherwise molnex's own reimplementation +accumulation, **not** a CG-convention difference (O3 vs O3_e3nn CG +~1.4e-8/op, A3). Re-measuring upstream parity needs an out-of-tree oracle +(route per Appendix A, 2026-08-09). + +**Current in-tree verification (2026-08-09).** `MOLNEX_MACE_WEIGHTS_DIR`-gated +`tests/test_molzoo/test_mace/test_checkpoint.py::TestOfficialOMolWeights`: +strict 104-parameter load through `OMOL_REMAP` plus E/F stability goldens on a +five-atom cluster; the weights dump is regenerated offline by +`scripts/omol_port/convert_omol_to_cueq_state.py` (no `mace`/`e3nn`). This is +a stability lock on this machine's own output — not an upstream parity claim. + +**Bessel frequencies (measured 2026-08-09).** The official `bessel_weights` +are bit-for-bit the fp32 evaluation of the analytic init `nπ/r_max` (upcast to +fp64); the offset from the fp64 analytic values (max 2.2120e-7 Å⁻¹ at n=7, +≤ 1 fp32 ulp per entry) is fp32 rounding of an untrained parameter, **not** +fitted drift. Doctrine unchanged: `bessel.freqs` is an `nn.Parameter` and must +be filled from the checkpoint — bit-exactness against the official surface +requires the checkpoint's fp32-rounded values, not the fp64 re-derivation. ### 7.2 Symmetry and Shape Tests +Variant-level claims live in +`tests/test_molzoo/test_mace/test_variants.py::TestMACEOMol` (the `::…` rows +below are relative to it); the shared pipeline it aliases is covered by +`tests/test_molzoo/test_mace/test_potential.py::TestMACEPotential`. + | Claim | Test path | Tolerance | |-------|-----------|-----------| -| `forward(td)` energy == `energy_forces` | `tests/test_molzoo/test_mace_omol.py::test_forward_matches_energy_forces` | 1e-9 eV | -| `forward(td)` forces == `energy_forces` | same | 1e-8 eV·Å | -| net force ≈ 0 (translation invariance) | `::test_forces_translation_invariant` | 1e-7 | +| `forward(td)` energy == `energy_forces` | `tests/test_molzoo/test_mace/test_variants.py::TestMACEOMol::test_forward_matches_energy_forces` | 1e-9 eV | +| `forward(td)` forces == `energy_forces` | same | 1e-8 eV/Å | +| net force ≈ 0 (translation invariance) | `::test_net_force_vanishes_on_an_isolated_molecule` | 1e-7 | | neutral default when `graphs.*` absent | `::test_missing_charge_spin_defaults_to_neutral` | 1e-9 | +| default spin is the closed-shell singlet `1`, not `0` | `tests/test_molzoo/test_mace/test_potential.py::TestMACEPotential::test_omol_defaults_to_a_neutral_closed_shell_singlet` | exact | | per-graph batching | `::test_batched_graphs` | 1e-9 | +| force loss reaches parameters (eval mode / through `forward`) | `::test_force_loss_reaches_parameters_in_eval_mode`, `::test_force_loss_through_forward_reaches_parameters` | exact | +| alien checkpoint rejected | `::test_load_omol_state_dict_refuses_an_alien_checkpoint` | raises | +| `OMOL_REMAP` roundtrip restores every parameter | `tests/test_molzoo/test_mace/test_checkpoint.py::TestCheckpointRemap::test_roundtrip_restores_every_parameter_omol` | exact | +| edge vectors / lengths under PBC shifts | `tests/test_molzoo/test_mace/test_geometry.py` | exact | | block-level vs cueq (radial/mlp/e0/joint/interaction/product/readout) | `scripts/omol_port/verify_*.py` | 0–7e-15 | ### 7.3 Engineering Benchmark @@ -222,12 +259,15 @@ for this spec (tracked in `scripts/omol_port/SPEC.md`). | Concern | Owner | Contract | |---------|-------|----------| -| Neighbor list / edges | collate / `NeighborList` | populates `edges.edge_index` `(E,2)` | +| Neighbor list / edges | collate / `molix.data.tasks.NeighborList` | populates `edges.edge_index` `(E,2)` | +| Edge displacements | `molzoo.mace.geometry` | `edge_vectors` / `edge_lengths`; differentiable w.r.t. `pos` | | Charge/spin inputs | dataset / collate | `graphs.total_charge`, `graphs.total_spin` (optional) | -| Building blocks | `molrep` / `molpot` | reused; not owned here | -| Forces | `molpot.derivation.ForceDerivation` | `F=-∂E/∂pos` from energy closure | -| Weight import | `load_omol_state_dict` + `mace.cli.convert_e3nn_cueq` | cueq `state_dict` → `MACEOMol` | -| Lazy export | `molzoo/__init__` | PEP 562 `__getattr__`; no eager cueq import | +| Building blocks | `molrep.readout.mace` (+ generic `molrep` / `molpot`) | reused; not owned here | +| Model graph | `molzoo.mace.encoder.MACEEncoder` | blocks + wiring; no energy, no forces | +| Energy / forces | `molzoo.mace.potential.MACEPotential` | `energy_core` (public, compile seam) + `molpot.derivation.kernels.grad_force_pass` | +| Configuration | `molzoo.mace.spec.MACEOMolSpec` | torch-free pydantic preset | +| Weight import | `molzoo.mace.checkpoint.OMOL_REMAP` via `load_omol_state_dict` / `MACEPotential.from_checkpoint`, after `mace.cli.convert_e3nn_cueq` | cueq `state_dict` → `MACEOMol`; strict on missing `nn.Parameter`s since 2026-08-07 | +| Lazy export | `molzoo/__init__` + `molzoo/mace/__init__` | PEP 562 `__getattr__`; no eager cueq import | ## 9. Version Pinning @@ -237,7 +277,8 @@ for this spec (tracked in `scripts/omol_port/SPEC.md`). | Reference repository | `ACEsuit/mace` | | Reference commit | not pinned to sha; `mace==0.3.16` (PyPI) used for conversion + verify. Follow-up audit to pin exact sha. | | Dependencies | `torch==2.12.1`, `cuequivariance==0.10.0`, `cuequivariance_torch==0.10.0`, `tensordict==0.13.0` | -| Public docs mirror | `docs/molzoo/specs/mace_omol.md` | +| Module relocation | `mace-subpackage-restructure` chain, commits `1ddd5ff..e825a51` (merged 2026-08-09): `src/molzoo/mace_omol.py` was retired into the `src/molzoo/mace/` package — config in `spec.py`, blocks in `encoder.py`, energy/forces in `potential.py`, key remap in `checkpoint.py`, the `MACEOMol` alias in `variants.py`. MACE-only `molrep` blocks moved to `molrep/interaction/mace/{conv,block,density}.py`, `molrep/readout/mace.py`, `molrep/embedding/mace.py`. Tests moved to `tests/test_molzoo/test_mace/`. Weights, hyper-parameters and numerics unchanged (§7.1 not re-run). | +| Public docs mirror | `docs/molzoo/specs/mace_omol.md` — byte-identical copy of this file; re-sync both halves on every edit | ## 10. Drift Policy @@ -254,8 +295,51 @@ rows. `mace-omol-port-01/02` implementation (status draft → partial). §5 rows `matched` per `scripts/omol_port/verify_*.py`. - 2026-06-21: accuracy bar set to 1e-4 (operator); cue O3 meets it at - 7e-7 eV / 4.3e-6 eV·Å. `mace-omol-port-02` ac-003 verified; chain done. + 7e-7 eV / 4.3e-6 eV/Å. `mace-omol-port-02` ac-003 verified; chain done. - 2026-06-21: measured O3 vs O3_e3nn CG = 1.4e-8/op → the 7e-7 residual is reimplementation accumulation, not a convention diff. Dropped the O3_e3nn pursuit entirely and removed the dead `MACEOMol(group=)` hook from MACEOMol / ResidualInteraction / EquivariantProductBasis (always cue O3). +- 2026-08-09: Anchors re-pointed for the `mace-subpackage-restructure` chain + (`1ddd5ff..e825a51`) — header, §1 ownership, §2.2 / §3.1 / §3.5 code anchors, + §5 crosswalk, §6 A4/A5/A6 verification paths, §7.2 test paths, §8 boundary, + §9 pinning row. Three content corrections found while re-pointing: + (a) `graphs.total_spin` absent defaults to **`1`** (closed-shell singlet + multiplicity `2S+1 = 1`), not `0` — `MACEPotential._condition_charge_spin` + fills `torch.ones(...)`, and `0` would index an untrained embedding row; + (b) force units written `eV·Å` now read `eV/Å` (§6 A3, §7.1, §7.2 and this + log); (c) forces are `torch.autograd.grad` + (`molpot.derivation.kernels.grad_force_pass` / + `molpot.derivation.force.autograd_forces_from_energy`), never + `ForceDerivation(method="functorch")` / `torch.func.grad` — cuEquivariance's + fused ops are legacy `autograd.Function`s that `torch.func.grad` rejects. + No section added, removed or renamed; §7.4 rows untouched; no numerical + claim changed. `docs/molzoo/specs/mace_omol.md` re-synced from this file + (the two copies had drifted on §6 A5, §7.1 and §8). +- 2026-08-09: **Dangling anchors, not fixed here.** Every + `scripts/omol_port/verify_*.py` cited by §5 (`matched` source of truth), §6 + A1/A2, §7.1, §7.2, §7.3 and §10 was deleted in commit `b85d12f`; only + `README.md` and `SPEC.md` remain in that directory. The recorded numbers + (7.0e-7 eV / 4.3e-6 eV/Å, the 0–7e-15 block-level residuals) are therefore + no longer reproducible in-tree, and §10's drift trigger (b) — "`verify_e2e.py` + E/F residual regresses > 10×" — cannot fire. Restoring the oracles (or + re-homing them under `regressions/` with hard-coded goldens) is out of scope + for `mace-subpackage-restructure-07-cleanup`, which is anchor-refresh only. + Route: `/mol:fix` or a follow-up spec. Combined with the 2026-08-07 caveat + in §7.1 (the parity run predates the trainable-Bessel loader fix), §7.1 + should be treated as **stale, pending re-measurement**. +- 2026-08-09: §7.1 rewritten (molzoo-auditor, operator-directed). (a) The + mace-torch parity figures are now labelled a dated **historical record** + (oracles deleted in `b85d12f`), and the current in-tree surface is named: + `MOLNEX_MACE_WEIGHTS_DIR`-gated `TestOfficialOMolWeights` (strict 104-param + load + E/F stability goldens) with the dump regenerable via + `scripts/omol_port/convert_omol_to_cueq_state.py`. (b) The "fitted drift" + reading of `bessel.freqs` is corrected to measurement: the official + `bessel_weights` are **bit-for-bit** the fp32 evaluation of the analytic + `nπ/r_max` init upcast to fp64 (max 2.2120e-7 Å⁻¹ from the fp64 values at + n=7, ≤ 1 fp32 ulp per entry) — storage rounding of an untrained parameter, + not training drift; the strict-loading doctrine is unchanged. A ⚠️ was + printed (not applied) against `src/molzoo/mace/checkpoint.py`'s docstring + ("fitted like any other weight" / "the fitted frequencies were dropped"). + No section added, removed or renamed; §7.4 untouched; §6 A1/A2 and §7.2's + dangling `verify_*.py` anchors remain covered by the entry above. diff --git a/docs/molzoo/tutorials/index.md b/docs/molzoo/tutorials/index.md index f720ad1..c8e4515 100644 --- a/docs/molzoo/tutorials/index.md +++ b/docs/molzoo/tutorials/index.md @@ -26,7 +26,7 @@ The complete project shape is: ```text source samples -> molix.data pipeline - -> GraphBatch + -> collated nested TensorDict -> molzoo encoder -> learned features -> molpot readout / potential head @@ -37,7 +37,7 @@ source samples For Allegro specifically: ```text -GraphBatch +collated nested TensorDict -> molzoo.Allegro -> ("edges", "edge_features") -> molpot.heads.EdgeEnergyHead @@ -46,7 +46,8 @@ GraphBatch ## 2. Prepare the Batch -MolZoo encoders expect a post-collate `GraphBatch`, not a raw sample dict. +MolZoo encoders expect a post-collate batch — the nested `TensorDict` with +`atoms` / `edges` / `graphs` namespaces — not a raw sample dict. For Allegro the required fields are: ```text @@ -68,14 +69,22 @@ pipe = ( .build() ) +dag = pipe.cache(source, base_dir="./cache") +train_ds, val_ds = dag.dataset(mmap=True).split(ratio=0.9, seed=42) + dm = DataModule( - source=source, - pipeline=pipe, + train_ds, + val_ds, + batch_nodes=pipe.batch_nodes, batch_size=32, ) dm.setup("fit") ``` +`DataModule` consumes pre-built datasets; the source → pipeline → cache → +dataset chain runs before it. See +[Data Modules](../../molix/user-guide/data-modules.md) for the full walkthrough. + The edge convention is fixed: ```text diff --git a/docs/molzoo/user-guide/allegro.md b/docs/molzoo/user-guide/allegro.md index 20cbaa7..45ef645 100644 --- a/docs/molzoo/user-guide/allegro.md +++ b/docs/molzoo/user-guide/allegro.md @@ -417,7 +417,7 @@ This section constructs a tiny directed graph by hand. In production, let ```python import torch -from molix.data.types import AtomData, EdgeData, GraphBatch, GraphData +from tensordict import TensorDict pos = torch.tensor( [ @@ -444,20 +444,20 @@ dst = edge_index[:, 1] edge_diff = pos[dst] - pos[src] edge_dist = edge_diff.norm(dim=-1) -batch = GraphBatch( - atoms=AtomData( +batch = TensorDict( + atoms=TensorDict( Z=Z, pos=pos, batch=torch.zeros(len(Z), dtype=torch.long), batch_size=[len(Z)], ), - edges=EdgeData( + edges=TensorDict( edge_index=edge_index, edge_diff=edge_diff, edge_dist=edge_dist, batch_size=[len(edge_index)], ), - graphs=GraphData( + graphs=TensorDict( num_atoms=torch.tensor([len(Z)]), batch_size=[1], ), @@ -465,6 +465,10 @@ batch = GraphBatch( ) ``` +Each namespace is a plain `TensorDict` with its own `batch_size` — 3 atoms, +4 directed edges, 1 graph. This is the same object `collate_molecules` would +hand you; nothing here is a molecule-specific subclass. + The directed average neighbor count for this toy batch is: ```python diff --git a/pyproject.toml b/pyproject.toml index d67dbd0..7d28c31 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -30,7 +30,7 @@ dependencies = [ "tensordict", "cuequivariance", "cuequivariance-torch", - "molcrafts-molpy>=0.3.0", + "molcrafts-molpy>=0.13.0", "python-multipart", "molcrafts-mollog>=1.0.0", "molcrafts-molcfg>=1.0.0", @@ -45,8 +45,14 @@ dependencies = [ ] [project.optional-dependencies] +# cuEquivariance fused kernels. Without one of these, cuEq silently degrades +# to its pure-torch path (~30x slower on MACE-class models) with only a +# UserWarning — install the extra matching the local CUDA major version. +cueq-cu12 = ["cuequivariance-ops-torch-cu12"] +cueq-cu13 = ["cuequivariance-ops-torch-cu13"] dev = [ "pytest>=7.0", + "pytest-xdist>=3.5", "pytest-cov>=4.0", "pytest-benchmark>=4.0", "ruff>=0.15", @@ -54,7 +60,8 @@ dev = [ "tox>=4.23", ] docs = [ - "zensical>=0.0.27", + "zensical>=0.0.51", + "molcrafts-zensical-theme>=0.2.3", "mkdocstrings[python]>=0.30", ] bench = ["pytest-benchmark>=4.0"] @@ -173,5 +180,5 @@ commands_pre = [ ], ] commands = [ - ["{envpython}", "-m", "pytest", "tests/", "-q", "-p", "no:cacheprovider"], + ["{envpython}", "-m", "pytest", "tests/", "-q", "-p", "no:cacheprovider", "-n", "auto", "--dist", "worksteal"], ] diff --git a/regressions/README.md b/regressions/README.md new file mode 100644 index 0000000..75e7e31 --- /dev/null +++ b/regressions/README.md @@ -0,0 +1,16 @@ +# regressions/ + +Standalone, public-API scenario scripts. Each file exercises molnex the way a +user would (construct → configure → one concern → read result) and asserts +against **hard-coded golden literals** captured once, offline, from a named +commit — never from a live third-party oracle, network call, or subprocess. + +These are **not** collected by pytest (`tests/` holds single-function unit +tests only; nothing here is imported by the suite). Run one directly: + +```bash +PYTHONPATH=src python regressions/.py # prints OK, exits 0; exits 1 on drift +``` + +Every script's header comment must record the capture command, the commit sha +the goldens came from, the torch version, the date, and the device/precision. diff --git a/regressions/dataset-profiler-salvage.py b/regressions/dataset-profiler-salvage.py new file mode 100644 index 0000000..b3d00cb --- /dev/null +++ b/regressions/dataset-profiler-salvage.py @@ -0,0 +1,513 @@ +"""Public-API scenario for `DatasetProfiler` (spec `dataset-profiler-salvage`, ac-007). + +The scenario a user actually performs, in one process, with nothing mocked: + + five sample dicts → PackedCache(tmp).save(...) → CachedDataset(sink) + → DatasetProfiler(n_samples=5).run(ds) + +and then every non-timing field of the returned `DatasetResult` is compared to +a literal computed by hand from the five samples written out below. The five +samples are **literals in this file**, not a fixture, a generator or an RNG +draw, so each golden is arithmetic anyone can redo on the page: 3 + 4 + 5 + 6 + +2 = 20 atoms over 5 records is a mean of 4.00, and 4 + 6 + 8 + 10 + 2 = 30 +edges over those 20 atoms is `avg_num_neighbors` = 1.50, exactly. + +What is pinned, and why those things + +* **Sizes** — `n_total`, `atom_stats.mean`, `max_atoms`, `max_edges`, + `avg_num_neighbors`. These are the claim the salvage rests on: with a + packed cache behind it the profiler reads `atom_ptr` / `edge_ptr` and + reports **all five** records exactly, rather than extrapolating from the + sampled ones. `counts_exact is True` and an empty `warnings` list are pinned + alongside them, because a silent fall-back to the sampled path would still + produce these same numbers here (`n_samples=5` of 5) and must not pass. +* **Field layout** — the sorted key list with each key's axis, dtype and + trailing shape, straight off the packed `payload["schema"]` + (`fields_exact is True`). Pins the atom / edge / graph classification of a + realistic key set, including that `targets.U0` `(1,)` is *not* mistaken for + a per-atom or per-edge column. +* **Targets** — `targets.U0` mean / min / max. Every literal target value is + dyadic (exactly representable in float32 **and** float64), so the mean is + exact arithmetic rather than a captured measurement: (-1.5) + (-0.5) + 0.25 + + 2.0 + 3.75 = 4.0, over 5 records, is 0.80. +* **Footprint** — `sample_bytes.mean` and `est_total_mb`. Not required by + ac-007, but they are analytic (Σ numel × element_size over the leaves) and + they are the one golden here that notices a dtype change: widening `pos` to + float64 moves 204.0 B/record and nothing else in this file. +* **Package re-export** — `from molix.profiler import DatasetProfiler, + DatasetResult` resolves to the very objects `molix.profiler.dataset` + defines, and `molix.profiler.__all__` lists them and stays alphabetised. + The unit tests import from the submodule, so this file is the only thing + holding the documented package-level import path. + +What is deliberately **not** pinned: anything timed. `access_ms`, +`cold_access_ms` and the whole Access section of the report are wall-clock +measurements, non-deterministic by construction; asserting a magnitude on them +would make this file fail on a busy node rather than on a regression. They are +exercised (the sampled path runs) and then ignored. + +Goldens +------- + capture command : none. There is no oracle and nothing was captured — every + literal below is arithmetic over the five sample dicts in + this file, and the derivation is written next to each one. + Confirmed against the implementation by running + `PYTHONPATH=src python regressions/dataset-profiler-salvage.py`. + commit : e94c35e (e94c35eb02a0bf961d0676addc5467e3ccd62623), with + the uncommitted `dataset-profiler-salvage` working tree on + top of it (`src/molix/profiler/dataset.py` is new there). + torch : 2.12.1+cpu (python 3.14.5, numpy via molix.profiler._utils) + date : 2026-08-09 + device / dtype : CPU throughout. `pos` / `edge_dist` / `targets.U0` are + float32 and `Z` / `edge_index` are int64, fixed by the + literals below rather than by `molix.config` — this file + never touches the global precision singleton, so it is + insensitive to it. + oracle : none. No third-party package, no network, no subprocess, + no RNG, no wall-clock value. The only filesystem use is a + `tempfile.TemporaryDirectory` that is deleted on exit. + tolerance : `math.isclose(rel_tol=1e-12)` on floats, exact equality on + counts, dtypes, axes and keys. 1e-12 is slack, not need: + every float golden here is a dyadic rational or a + correctly-rounded quotient of two exact integers, so the + observed values match to the last bit. + +Run: + PYTHONPATH=src python regressions/dataset-profiler-salvage.py +""" + +from __future__ import annotations + +import io +import math +import sys +import tempfile +from contextlib import redirect_stdout +from pathlib import Path + +import torch + +from molix import profiler as profiler_package +from molix.data.cache import PackedCache +from molix.data.dataset import CachedDataset +from molix.profiler import DatasetProfiler, DatasetResult +from molix.profiler import dataset as dataset_module + +# --------------------------------------------------------------------------- +# The five samples, written out in full. +# +# Each is a linear chain: atoms on the x axis, bidirectional nearest-neighbour +# edges (source, target), so `edge_dist` is the spacing and every literal is a +# dyadic rational — exact in float32, exact in float64, and hand-checkable +# against `pos`. Chain of n atoms → n-1 bonds → 2(n-1) edges. The `edge_index` +# rows follow the repo convention: column 0 is the source, column 1 the target. +# +# The atom counts 3 / 4 / 5 / 6 / 2 are deliberately not monotone and the last +# record's atom count (2) coincides with its edge count, which is exactly the +# ambiguity `molix.data.cache._infer_schema_across` resolves by scanning all +# records: a single-record inference could classify `edge_index` as per-atom. +# --------------------------------------------------------------------------- + +SAMPLES: list[dict[str, object]] = [ + { # 0 — 3 atoms, spacings 1.0 / 1.5 → 4 edges + "Z": torch.tensor([8, 1, 1], dtype=torch.long), + "pos": torch.tensor( + [[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [2.5, 0.0, 0.0]], dtype=torch.float32 + ), + "edge_index": torch.tensor([[0, 1], [1, 0], [1, 2], [2, 1]], dtype=torch.long), + "edge_dist": torch.tensor([1.0, 1.0, 1.5, 1.5], dtype=torch.float32), + "targets": {"U0": torch.tensor([-1.5], dtype=torch.float32)}, + }, + { # 1 — 4 atoms, spacings 1.0 / 1.0 / 1.5 → 6 edges + "Z": torch.tensor([6, 1, 1, 1], dtype=torch.long), + "pos": torch.tensor( + [[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [2.0, 0.0, 0.0], [3.5, 0.0, 0.0]], + dtype=torch.float32, + ), + "edge_index": torch.tensor( + [[0, 1], [1, 0], [1, 2], [2, 1], [2, 3], [3, 2]], dtype=torch.long + ), + "edge_dist": torch.tensor([1.0, 1.0, 1.0, 1.0, 1.5, 1.5], dtype=torch.float32), + "targets": {"U0": torch.tensor([-0.5], dtype=torch.float32)}, + }, + { # 2 — 5 atoms, spacings 1.5 / 1.0 / 1.0 / 1.5 → 8 edges + "Z": torch.tensor([6, 6, 1, 1, 1], dtype=torch.long), + "pos": torch.tensor( + [ + [0.0, 0.0, 0.0], + [1.5, 0.0, 0.0], + [2.5, 0.0, 0.0], + [3.5, 0.0, 0.0], + [5.0, 0.0, 0.0], + ], + dtype=torch.float32, + ), + "edge_index": torch.tensor( + [[0, 1], [1, 0], [1, 2], [2, 1], [2, 3], [3, 2], [3, 4], [4, 3]], dtype=torch.long + ), + "edge_dist": torch.tensor([1.5, 1.5, 1.0, 1.0, 1.0, 1.0, 1.5, 1.5], dtype=torch.float32), + "targets": {"U0": torch.tensor([0.25], dtype=torch.float32)}, + }, + { # 3 — 6 atoms, spacings 1.0 ×4 / 1.5 → 10 edges (the largest record) + "Z": torch.tensor([6, 6, 6, 1, 1, 1], dtype=torch.long), + "pos": torch.tensor( + [ + [0.0, 0.0, 0.0], + [1.0, 0.0, 0.0], + [2.0, 0.0, 0.0], + [3.0, 0.0, 0.0], + [4.0, 0.0, 0.0], + [5.5, 0.0, 0.0], + ], + dtype=torch.float32, + ), + "edge_index": torch.tensor( + [ + [0, 1], + [1, 0], + [1, 2], + [2, 1], + [2, 3], + [3, 2], + [3, 4], + [4, 3], + [4, 5], + [5, 4], + ], + dtype=torch.long, + ), + "edge_dist": torch.tensor( + [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.5, 1.5], dtype=torch.float32 + ), + "targets": {"U0": torch.tensor([2.0], dtype=torch.float32)}, + }, + { # 4 — 2 atoms, spacing 1.25 → 2 edges (n_edges == n_atoms here) + "Z": torch.tensor([1, 1], dtype=torch.long), + "pos": torch.tensor([[0.0, 0.0, 0.0], [1.25, 0.0, 0.0]], dtype=torch.float32), + "edge_index": torch.tensor([[0, 1], [1, 0]], dtype=torch.long), + "edge_dist": torch.tensor([1.25, 1.25], dtype=torch.float32), + "targets": {"U0": torch.tensor([3.75], dtype=torch.float32)}, + }, +] + +#: How many records the sampled slow path may read. Equal to `len(SAMPLES)`, so +#: the sampled and exact paths cover the same records — which is why +#: `counts_exact` has to be asserted separately from the numbers themselves. +N_SAMPLES = 5 + +#: Relative tolerance for the float goldens. See the docstring: slack, not need. +REL_TOL = 1e-12 + +# --------------------------------------------------------------------------- +# Goldens — analytic, derived above each literal. +# --------------------------------------------------------------------------- + +#: `len(ds)`. +N_TOTAL = 5 + +#: (3 + 4 + 5 + 6 + 2) / 5 = 20 / 5. +ATOM_MEAN = 4.0 + +#: Population std (numpy default, ddof=0) of 3, 4, 5, 6, 2 about the mean 4: +#: (1 + 0 + 1 + 4 + 4) / 5 = 2, so √2. Pins the ddof convention of `ValueStat`. +ATOM_STD = 1.4142135623730951 + +#: (4 + 6 + 8 + 10 + 2) / 5 = 30 / 5. Non-`None` here: the cache has an +#: `edge_ptr`, which is the branch the "no NeighborList" warning path forgoes. +EDGE_MEAN = 6.0 + +#: Record 3, the 6-atom chain. +MAX_ATOMS = 6 + +#: Record 3 again: 2 × (6 - 1). +MAX_EDGES = 10 + +#: 30 total edges / 20 total atoms. Bidirectional edges, so this is the mean +#: neighbour count per atom — the MACE / Allegro normalisation constant. +AVG_NUM_NEIGHBORS = 1.5 + +#: Σ numel × element_size per record, over the five records: +#: n=3, e=4 → 3·8 + 3·3·4 + 4·2·8 + 4·4 + 1·4 = 144 B +#: n=4, e=6 → 32 + 48 + 96 + 24 + 4 = 204 B +#: n=5, e=8 → 40 + 60 + 128 + 32 + 4 = 264 B +#: n=6, e=10 → 48 + 72 + 160 + 40 + 4 = 324 B +#: n=2, e=2 → 16 + 24 + 32 + 8 + 4 = 84 B +#: 1020 B over 5 records. +SAMPLE_BYTES_MEAN = 204.0 + +#: 204.0 B/record × 5 records / 1e6. +EST_TOTAL_MB = 0.00102 + +#: The packed `payload["schema"]`, as `(key, axis, dtype, extra_shape)` sorted +#: by key. `extra_shape` is the shape after the packing axis for atom / edge +#: fields — `(3,)` for `pos` `(N, 3)`, `()` for `Z` `(N,)` — and the full +#: per-record shape for a graph field, hence `(1,)` for `targets.U0`. +FIELDS: tuple[tuple[str, str, str, tuple[int, ...]], ...] = ( + ("Z", "atom", "int64", ()), + ("edge_dist", "edge", "float32", ()), + ("edge_index", "edge", "int64", (2,)), + ("pos", "atom", "float32", (3,)), + ("targets.U0", "graph", "float32", (1,)), +) + +#: The one label column: `targets.U0`, dotted exactly as the packed schema +#: names it (nested `{"targets": {"U0": ...}}` in the raw sample). +TARGET_KEY = "targets.U0" + +#: (-1.5) + (-0.5) + 0.25 + 2.0 + 3.75 = 4.0, over 5 records. +TARGET_MEAN = 0.8 + +#: Record 0 / record 4. +TARGET_MIN = -1.5 +TARGET_MAX = 3.75 + +#: No `nan` / `inf` among the five literals, so the non-finite warning path +#: stays silent and `warnings` is empty overall. +TARGET_NONFINITE = 0 + +#: Section labels `DatasetResult.print_report` must emit. The atom-count skew +#: is p95/p50 = 5.8/4 = 1.45× here, well under the 3× warning threshold, so a +#: clean run prints these and no `[WARN]` line at all. +REPORT_SECTIONS = ("Size", "Access", "Footprint", "Fields", "Targets") + + +# --------------------------------------------------------------------------- +# Checking +# --------------------------------------------------------------------------- + + +class Checker: + """Collects every deviation so one run reports all failures, not the first.""" + + def __init__(self) -> None: + self.failures: list[str] = [] + + def _row(self, name: str, got: object, ok: bool) -> None: + print(f" {name:<38} {got!s:<34} {'ok' if ok else 'FAILED'}") + + def exact(self, name: str, got: object, want: object) -> None: + """Assert a count, string, dtype or key list — no tolerance applies.""" + ok = got == want + if not ok: + self.failures.append(f"{name}: got {got!r}, want {want!r}") + self._row(name, got, ok) + + def close(self, name: str, got: float, want: float) -> None: + """Assert a float golden within :data:`REL_TOL`.""" + ok = math.isclose(got, want, rel_tol=REL_TOL, abs_tol=0.0) + if not ok: + self.failures.append(f"{name}: got {got!r}, want {want!r} (rel_tol={REL_TOL})") + self._row(name, got, ok) + + def truth(self, name: str, holds: bool, message: str) -> None: + """Assert a boolean contract (a path was taken, a name is re-exported, ...).""" + if not holds: + self.failures.append(f"{name}: {message}") + self._row(name, holds, holds) + + +def profile_samples(directory: Path) -> DatasetResult: + """Run the whole user-facing scenario inside *directory*. + + Args: + directory: Scratch directory for the cache file; caller owns its + lifetime. Nothing about the path enters a golden. + + Returns: + The `DatasetResult` for the five literal samples, read back through a + real `CachedDataset` — i.e. after a genuine pack / save / load round + trip, not from the in-memory dicts. + """ + sink = directory / "dataset-profiler-salvage.pt" + PackedCache(sink).save(SAMPLES) + dataset = CachedDataset(sink) + return DatasetProfiler(n_samples=N_SAMPLES).run(dataset) + + +def check_reexport(checker: Checker, result: DatasetResult) -> None: + """The documented package-level import path resolves to the real objects. + + `molix.profiler.__init__` re-exports both names; the unit suite imports + from `molix.profiler.dataset` directly, so nothing but this file would + notice the package-level path breaking. + + Args: + checker: Failure collector. + result: A result produced by the package-level profiler class. + """ + print("Package re-export (from molix.profiler import DatasetProfiler, DatasetResult)") + checker.truth( + "reexport.profiler_is_submodule_class", + DatasetProfiler is dataset_module.DatasetProfiler, + "molix.profiler.DatasetProfiler is not molix.profiler.dataset.DatasetProfiler", + ) + checker.truth( + "reexport.result_is_submodule_class", + DatasetResult is dataset_module.DatasetResult, + "molix.profiler.DatasetResult is not molix.profiler.dataset.DatasetResult", + ) + checker.truth( + "reexport.run_returns_that_result", + isinstance(result, DatasetResult), + f"run() returned {type(result).__name__}, not the re-exported DatasetResult", + ) + exported = list(profiler_package.__all__) + checker.truth( + "reexport.both_names_in_all", + {"DatasetProfiler", "DatasetResult"} <= set(exported), + f"molix.profiler.__all__ is missing " + f"{sorted({'DatasetProfiler', 'DatasetResult'} - set(exported))}", + ) + checker.truth( + "reexport.all_is_alphabetised", + exported == sorted(exported), + f"molix.profiler.__all__ is out of order: {exported}", + ) + + +def check_sizes(checker: Checker, result: DatasetResult) -> None: + """Size statistics, and that they came from the packed pointers. + + Args: + checker: Failure collector. + result: The profiled result. + """ + print("\nSizes (exact packed-pointer fast path)") + checker.exact("size.n_total", result.n_total, N_TOTAL) + checker.exact("size.n_sampled", result.n_sampled, N_SAMPLES) + checker.truth( + "size.counts_exact", + result.counts_exact, + "size stats fell back to the sampled path — atom_ptr / edge_ptr were " + "not read, so the numbers below cover only the sampled records", + ) + checker.close("size.atom_stats.mean", result.atom_stats.mean, ATOM_MEAN) + checker.close("size.atom_stats.std", result.atom_stats.std, ATOM_STD) + checker.truth( + "size.edge_stats_present", + result.edge_stats is not None, + "edge_stats is None — the cache lost its edge_ptr", + ) + if result.edge_stats is not None: + checker.close("size.edge_stats.mean", result.edge_stats.mean, EDGE_MEAN) + checker.exact("size.max_atoms", result.max_atoms, MAX_ATOMS) + checker.exact("size.max_edges", result.max_edges, MAX_EDGES) + checker.close("size.avg_num_neighbors", result.avg_num_neighbors, AVG_NUM_NEIGHBORS) + + +def check_footprint(checker: Checker, result: DatasetResult) -> None: + """Per-record byte footprint and the full-materialisation extrapolation. + + Args: + checker: Failure collector. + result: The profiled result. + """ + print("\nFootprint (leaf tensors, analytic — the dtype tripwire)") + checker.close("footprint.sample_bytes.mean", result.sample_bytes.mean, SAMPLE_BYTES_MEAN) + checker.close("footprint.est_total_mb", result.est_total_mb, EST_TOTAL_MB) + + +def check_fields(checker: Checker, result: DatasetResult) -> None: + """Field layout, read off the packed schema rather than inferred. + + Args: + checker: Failure collector. + result: The profiled result. + """ + print("\nFields (packed payload['schema'], sorted by key)") + checker.truth( + "fields.fields_exact", + result.fields_exact, + "field layout was inferred from the sampled records — the packed " + "schema was not reachable, so axes and trailing shapes are guesses", + ) + observed = sorted( + (f.key, f.axis, str(f.dtype).removeprefix("torch."), tuple(f.extra_shape)) + for f in result.fields + ) + checker.exact("fields.keys", tuple(row[0] for row in observed), tuple(f[0] for f in FIELDS)) + for want in FIELDS: + got = next((row for row in observed if row[0] == want[0]), None) + checker.exact(f"fields.{want[0]}", got, want) + + +def check_targets(checker: Checker, result: DatasetResult) -> None: + """Label statistics for the single `targets.U0` column. + + Args: + checker: Failure collector. + result: The profiled result. + """ + print("\nTargets (targets.U0)") + checker.exact("targets.keys", tuple(t.key for t in result.targets), (TARGET_KEY,)) + target = next((t for t in result.targets if t.key == TARGET_KEY), None) + if target is None: + checker.truth("targets.U0_present", False, f"no TargetStat for {TARGET_KEY!r}") + return + checker.close("targets.U0.mean", target.stat.mean, TARGET_MEAN) + checker.close("targets.U0.min", target.min, TARGET_MIN) + checker.close("targets.U0.max", target.max, TARGET_MAX) + checker.exact("targets.U0.n_nonfinite", target.n_nonfinite, TARGET_NONFINITE) + + +def check_report(checker: Checker, result: DatasetResult) -> None: + """`print_report` emits every section, and this clean run warns about nothing. + + The report text is captured rather than printed: it carries wall-clock + latencies, and this file's own output has to stay byte-identical between + runs. Only the section labels and the absence of `[WARN]` are asserted — + both deterministic. + + Args: + checker: Failure collector. + result: The profiled result. + """ + print("\nReport (captured, not printed — it carries timings)") + buffer = io.StringIO() + with redirect_stdout(buffer): + result.print_report() + text = buffer.getvalue() + missing = [section for section in REPORT_SECTIONS if f" {section}" not in text] + checker.truth( + "report.sections_present", + not missing, + f"print_report omitted section(s) {missing}", + ) + checker.truth( + "report.rule_present", + "─" * 72 in text, + "print_report lost its 72-char section rule", + ) + checker.exact("report.warnings", tuple(result.warnings), ()) + checker.truth( + "report.no_warn_line", + "[WARN]" not in text, + "print_report emitted a [WARN] line on a cache that is exact in every " + "section and has no non-finite label", + ) + + +def main() -> int: + """Profile the five literal samples and compare every non-timing field.""" + checker = Checker() + with tempfile.TemporaryDirectory(prefix="molnex-dataset-profiler-") as directory: + result = profile_samples(Path(directory)) + + check_reexport(checker, result) + check_sizes(checker, result) + check_footprint(checker, result) + check_fields(checker, result) + check_targets(checker, result) + check_report(checker, result) + + if checker.failures: + print("\nFAILED — DatasetProfiler no longer matches the analytic goldens:") + for failure in checker.failures: + print(f" {failure}") + return 1 + print("\nOK") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/regressions/learnable-classical-ff-01-ir-kernels.py b/regressions/learnable-classical-ff-01-ir-kernels.py new file mode 100644 index 0000000..a9191e9 --- /dev/null +++ b/regressions/learnable-classical-ff-01-ir-kernels.py @@ -0,0 +1,154 @@ +"""Public-API regression for Class-I Potential IR + torsion kernels. + +Spec: `learnable-classical-ff-01-ir-kernels`. + +Hard-coded goldens only — no OpenMM / third-party oracle. Pins: + +1. CLASS_I_CANONICAL units (kcal/mol, angstrom, e, radian) +2. NonbondedScaling 1-2/1-3/1-4 defaults (AMBER/GAFF Class-I) +3. ProperTorsionPeriodic cis identity: + E = (k/s)[1 + cos(n*phi - gamma)] + k=1, s=1, n=1, gamma=0, phi=0 → E = 2.0 kcal/mol +4. Multi-term sum at cis: k=(1, 0.5), n=(1, 2) → E = 3.0 +5. ImproperHarmonic: k=2, chi=pi/6, chi0=0 → E = (pi/6)^2 +6. BondHarmonic sanity reuse: k=2, r=1.5, r0=1 → E = 0.25 + +Provenance +---------- + formula : OpenMM User Guide §19.4 / SMIRNOFF proper torsion; + E = (k/idivf)[1 + cos(n*phi - gamma)] + geometry : analytic cis fixture + i=(0,1,0), j=(0,0,0), k=(1,0,0), l=(1,1,0) → phi=0 + capture command : PYTHONPATH=src python regressions/learnable-classical-ff-01-ir-kernels.py + date : 2026-08-10 + note : Spec domain-basis text said k=2 → E=2.0, which contradicts + the written formula (2k/s = 4 for k=2). Goldens follow the + formula with k=1 → E=2.0 at cis. +""" + +from __future__ import annotations + +import math +import sys + +import torch + + +def main() -> int: + from molpot.ir import CLASS_I_CANONICAL, NonbondedScaling, PotentialIR + from molpot.potentials import ( + BondHarmonic, + ImproperHarmonic, + ImproperPeriodic, + ProperTorsionPeriodic, + ) + + # --- 1. Units --- + assert CLASS_I_CANONICAL["energy"] == "kcal/mol" + assert CLASS_I_CANONICAL["length"] in ("angstrom", "Å", "A") + assert CLASS_I_CANONICAL["charge"] == "e" + assert CLASS_I_CANONICAL["angle"] in ("radian", "rad") + + # --- 2. Nonbonded scaling defaults --- + scaling = NonbondedScaling() + assert float(scaling.scale_q_12) == 0.0 + assert float(scaling.scale_q_13) == 0.0 + assert math.isclose(float(scaling.scale_q_14), 5.0 / 6.0, abs_tol=1e-12) + assert float(scaling.scale_lj_12) == 0.0 + assert float(scaling.scale_lj_13) == 0.0 + assert math.isclose(float(scaling.scale_lj_14), 0.5, abs_tol=1e-12) + + # Empty IR is valid + _ = PotentialIR() + + # --- Geometry: cis phi=0 --- + cis = torch.tensor( + [ + [0.0, 1.0, 0.0], + [0.0, 0.0, 0.0], + [1.0, 0.0, 0.0], + [1.0, 1.0, 0.0], + ], + dtype=torch.float64, + ) + # proper torsion i-j-k-l on atoms 0-1-2-3 + proper_idx = torch.tensor([[0], [1], [2], [3]], dtype=torch.long) + # improper molrs center-first [center, i, j, k] with center = atom 1 + improper_idx = torch.tensor([[1], [0], [2], [3]], dtype=torch.long) + types0 = torch.tensor([0], dtype=torch.long) + + # --- 3. Proper torsion cis golden E=2.0 (k=1) --- + proper = ProperTorsionPeriodic( + k=torch.tensor([[1.0]], dtype=torch.float64), + periodicity=torch.tensor([1], dtype=torch.long), + phase=torch.tensor([[0.0]], dtype=torch.float64), + idivf=torch.tensor([1.0], dtype=torch.float64), + ) + e_proper = float(proper(pos=cis, proper_index=proper_idx, proper_types=types0)) + assert math.isclose(e_proper, 2.0, abs_tol=1e-10), f"proper cis E={e_proper}" + + # --- 4. Multi-term sum --- + proper_mt = ProperTorsionPeriodic( + k=torch.tensor([[1.0, 0.5]], dtype=torch.float64), + periodicity=torch.tensor([1, 2], dtype=torch.long), + phase=torch.tensor([[0.0, 0.0]], dtype=torch.float64), + idivf=torch.tensor([1.0], dtype=torch.float64), + ) + e_mt = float(proper_mt(pos=cis, proper_index=proper_idx, proper_types=types0)) + assert math.isclose(e_mt, 3.0, abs_tol=1e-10), f"multi-term E={e_mt}" + + # --- Improper periodic same cis golden (center-first layout) --- + improper = ImproperPeriodic( + k=torch.tensor([[1.0]], dtype=torch.float64), + periodicity=torch.tensor([1], dtype=torch.long), + phase=torch.tensor([[0.0]], dtype=torch.float64), + idivf=torch.tensor([1.0], dtype=torch.float64), + ) + e_imp = float(improper(pos=cis, improper_index=improper_idx, improper_types=types0)) + assert math.isclose(e_imp, 2.0, abs_tol=1e-10), f"improper cis E={e_imp}" + + # --- 5. Improper harmonic chi=pi/6 --- + pi6_pos = torch.tensor( + [ + [0.0, 1.0, 0.0], + [0.0, 0.0, 0.0], + [1.0, 0.0, 0.0], + [1.0, math.cos(math.pi / 6.0), math.sin(math.pi / 6.0)], + ], + dtype=torch.float64, + ) + imp_h = ImproperHarmonic( + k=torch.tensor([2.0], dtype=torch.float64), + chi0=torch.tensor([0.0], dtype=torch.float64), + ) + e_ih = float(imp_h(pos=pi6_pos, improper_index=improper_idx, improper_types=types0)) + expected_ih = (math.pi / 6.0) ** 2 + assert math.isclose(e_ih, expected_ih, abs_tol=1e-8), f"improper harmonic E={e_ih}" + + # --- 6. BondHarmonic reuse sanity --- + bond = BondHarmonic(k=torch.tensor([2.0]), r0=torch.tensor([1.0])) + e_bond = float( + bond( + pos=torch.tensor([[0.0, 0.0, 0.0], [1.5, 0.0, 0.0]]), + bond_index=torch.tensor([[0], [1]], dtype=torch.long), + bond_types=torch.tensor([0], dtype=torch.long), + ) + ) + # 0.5 * 2.0 * (1.5 - 1.0)^2 = 0.25 + assert math.isclose(e_bond, 0.25, abs_tol=1e-10), f"bond E={e_bond}" + + print("learnable-classical-ff-01-ir-kernels: all hard-coded goldens OK") + print(f" proper cis E = {e_proper}") + print(f" proper multi-term E = {e_mt}") + print(f" improper cis E = {e_imp}") + print(f" improper harmonic E = {e_ih}") + print(f" bond harmonic E = {e_bond}") + return 0 + + +if __name__ == "__main__": + try: + raise SystemExit(main()) + except Exception as exc: # noqa: BLE001 — standalone regression script + print(f"FAIL: {exc}", file=sys.stderr) + raise diff --git a/regressions/learnable-classical-ff-02-valence-topology.py b/regressions/learnable-classical-ff-02-valence-topology.py new file mode 100644 index 0000000..74c315c --- /dev/null +++ b/regressions/learnable-classical-ff-02-valence-topology.py @@ -0,0 +1,156 @@ +"""Public-API regression for valence topology collate namespaces. + +Spec: `learnable-classical-ff-02-valence-topology`. + +Hard-coded goldens only — no third-party oracles. Pins: + +1. Nested pre-collate → nested TensorDict columns post-collate +2. Atom-offset rebase: second molecule indices = local + n_atoms_0 +3. Improper center-first (atomi = center, molrs) +4. Optional stack helpers build COO [arity, N] without making COO the schema +5. Bonds collate still works alongside angles +6. PackedCache round-trip + collate_packed leaf equality + +Provenance +---------- + capture command : PYTHONPATH=src python regressions/learnable-classical-ff-02-valence-topology.py + date : 2026-08-10 +""" + +from __future__ import annotations + +import sys +import tempfile +from pathlib import Path + +import torch + + +def main() -> int: + from molix.data.cache import PackedCache + from molix.data.collate import collate_molecules, collate_packed + from molix.data.dataset import MmapDataset + from molix.datasets._valence_columns import ( + stack_angle_index, + stack_improper_index, + stack_proper_index, + ) + + # --- Fixtures: m1 (3 atoms, 1 angle), m2 (4 atoms, 2 angles) --- + m1 = { + "Z": torch.ones(3, dtype=torch.long), + "pos": torch.zeros(3, 3), + "bond_index": torch.tensor([[0, 1], [1, 2]], dtype=torch.long), + "bond_types": torch.tensor([0, 1], dtype=torch.long), + "angles": { + "atomi": torch.tensor([0], dtype=torch.long), + "atomj": torch.tensor([1], dtype=torch.long), + "atomk": torch.tensor([2], dtype=torch.long), + "type": torch.tensor([0], dtype=torch.long), + }, + "propers": { + "atomi": torch.tensor([0], dtype=torch.long), + "atomj": torch.tensor([1], dtype=torch.long), + "atomk": torch.tensor([2], dtype=torch.long), + "atoml": torch.tensor([0], dtype=torch.long), # reuse atom 0 for fixture + "type": torch.tensor([0], dtype=torch.long), + }, + "impropers": { + "atomi": torch.tensor([1], dtype=torch.long), # center + "atomj": torch.tensor([0], dtype=torch.long), + "atomk": torch.tensor([2], dtype=torch.long), + "atoml": torch.tensor([0], dtype=torch.long), + }, + } + m2 = { + "Z": torch.ones(4, dtype=torch.long), + "pos": torch.zeros(4, 3), + "bond_index": torch.tensor([[0, 1, 2], [1, 2, 3]], dtype=torch.long), + "bond_types": torch.tensor([0, 0, 1], dtype=torch.long), + "angles": { + "atomi": torch.tensor([0, 1], dtype=torch.long), + "atomj": torch.tensor([1, 2], dtype=torch.long), + "atomk": torch.tensor([2, 3], dtype=torch.long), + "type": torch.tensor([1, 0], dtype=torch.long), + }, + "propers": { + "atomi": torch.tensor([0], dtype=torch.long), + "atomj": torch.tensor([1], dtype=torch.long), + "atomk": torch.tensor([2], dtype=torch.long), + "atoml": torch.tensor([3], dtype=torch.long), + "type": torch.tensor([1], dtype=torch.long), + }, + "impropers": { + "atomi": torch.tensor([2], dtype=torch.long), # center + "atomj": torch.tensor([0], dtype=torch.long), + "atomk": torch.tensor([1], dtype=torch.long), + "atoml": torch.tensor([3], dtype=torch.long), + }, + } + + batch = collate_molecules([m1, m2]) + + # --- 1. Nested column access + rebase (m2 offset = 3) --- + assert batch["angles"]["atomi"].tolist() == [0, 3, 4] + assert batch["angles"]["atomj"].tolist() == [1, 4, 5] + assert batch["angles"]["atomk"].tolist() == [2, 5, 6] + assert batch["angles"]["type"].tolist() == [0, 1, 0] + assert list(batch["angles"].batch_size) == [3] + assert "angle_index" not in batch["angles"].keys() + + # --- 2. Propers rebased --- + assert batch["propers"]["atomi"].tolist() == [0, 3] + assert batch["propers"]["atoml"].tolist() == [0, 6] + assert batch["propers"]["type"].tolist() == [0, 1] + + # --- 3. Impropers: atomi = center, rebased --- + assert batch["impropers"]["atomi"].tolist() == [1, 5] # 1, 2+3 + + # --- 4. Bonds still green --- + assert batch["bonds", "bond_index"].tolist() == [ + [0, 1, 3, 4, 5], + [1, 2, 4, 5, 6], + ] + assert batch["bonds", "bond_types"].tolist() == [0, 1, 0, 0, 1] + + # --- 5. Stack helpers (kernel-local only) --- + a_idx = stack_angle_index(batch["angles"]) + assert a_idx.shape == (3, 3) + assert a_idx.tolist() == [[0, 3, 4], [1, 4, 5], [2, 5, 6]] + p_idx = stack_proper_index(batch["propers"]) + assert p_idx.shape == (4, 2) + i_idx = stack_improper_index(batch["impropers"]) + assert i_idx[0].tolist() == [1, 5] # centers at row 0 + + # --- 6. PackedCache round-trip + collate_packed --- + with tempfile.TemporaryDirectory() as td: + sink = Path(td) / "valence.pt" + PackedCache(sink).save([m1, m2]) + assert sink.is_file() + ds = MmapDataset(sink) + u1 = ds[1] + assert torch.equal(u1["angles"]["atomi"], m2["angles"]["atomi"]) + assert torch.equal(u1["impropers"]["atomi"], m2["impropers"]["atomi"]) + + fast = collate_packed(ds.packed_view(), [0, 1]) + oracle = collate_molecules([ds[0], ds[1]]) + for fam in ("angles", "propers", "impropers"): + for col in oracle[fam].keys(): + assert torch.equal(fast[fam][col], oracle[fam][col]), f"{fam}.{col}" + assert torch.equal(fast["bonds", "bond_index"], oracle["bonds", "bond_index"]) + + # --- 7. No molpy as batch store in collate/cache --- + import molix.data.cache as cache_mod + import molix.data.collate as collate_mod + + for mod in (collate_mod, cache_mod): + text = Path(mod.__file__).read_text() + assert "from molpy" not in text + assert "import molrs" not in text + + print("learnable-classical-ff-02-valence-topology: OK") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/regressions/learnable-classical-ff-03-mm-heads.py b/regressions/learnable-classical-ff-03-mm-heads.py new file mode 100644 index 0000000..8e13238 --- /dev/null +++ b/regressions/learnable-classical-ff-03-mm-heads.py @@ -0,0 +1,142 @@ +"""Public-API regression for continuous MM heads + ClassicalMMComposer. + +Spec: `learnable-classical-ff-03-mm-heads`. + +Hard-coded goldens only — no third-party oracle. Pins: + +1. BondParamHead / AngleParamHead positivity + theta0 ∈ (0, π) +2. ProperTorsionParamHead multi-term shapes (k ≥ 0) +3. MultiHead(LJ + Charge) merge + neutrality +4. ClassicalMMComposer.parameterize → CLASS_I PotentialIR +5. Bond harmonic energy golden: k=2, r=1.5, r0=1 → E = 0.25 kcal/mol +6. No molzoo / molrep.chem imports in classical_mm / mm_heads + +Provenance +---------- + formula : E = ½ k (r − r₀)² (OpenMM §19 / Class-I) + geometry : two atoms on x-axis at distance r + capture command : PYTHONPATH=src python regressions/learnable-classical-ff-03-mm-heads.py + date : 2026-08-10 +""" + +from __future__ import annotations + +import ast +import math +import sys +from pathlib import Path + +import torch +from tensordict import TensorDict + + +def main() -> int: + from molpot import ( + AngleParamHead, + BondParamHead, + ChargeHead, + ClassicalMMComposer, + ImproperParamHead, + LJParameterHead, + MultiHead, + PotentialIR, + ProperTorsionParamHead, + ) + from molpot.ir import CLASS_I_CANONICAL + + # --- 1. Head positivity --- + bond_head = BondParamHead(feature_dim=8, hidden_dim=16) + bout = bond_head(torch.randn(4, 8)) + assert bout["k"].shape == (4,) and torch.all(bout["k"] > 0) + assert bout["r0"].shape == (4,) and torch.all(bout["r0"] > 0) + + angle_head = AngleParamHead(feature_dim=8, hidden_dim=16) + aout = angle_head(torch.randn(3, 8)) + assert torch.all(aout["k"] > 0) + assert torch.all(aout["theta0"] > 0) and torch.all(aout["theta0"] < math.pi) + + proper_head = ProperTorsionParamHead( + feature_dim=8, n_terms=2, periodicity=(1, 2) + ) + pout = proper_head(torch.randn(2, 8)) + assert pout["k"].shape == (2, 2) and torch.all(pout["k"] >= 0) + assert pout["phase"].shape == (2, 2) + + _ = ImproperParamHead(feature_dim=8, include_harmonic=True, include_periodic=False) + + # --- 2. MultiHead reuse --- + multi = MultiHead( + { + "lj": LJParameterHead(feature_dim=8, hidden_dim=16), + "q": ChargeHead(feature_dim=8, hidden_dim=16, total_charge=0.0), + } + ) + batch_idx = torch.tensor([0, 0, 1, 1], dtype=torch.long) + mout = multi(torch.randn(4, 8), batch=batch_idx) + assert set(mout) == {"epsilon", "sigma", "charge"} + assert mout["charge"][:2].sum().abs() < 1e-5 + + # --- 3. Composer IR + golden energy --- + class _ConstBond(torch.nn.Module): + def forward(self, features: torch.Tensor) -> dict[str, torch.Tensor]: + n = features.shape[0] + return { + "k": torch.full((n,), 2.0, dtype=features.dtype), + "r0": torch.full((n,), 1.0, dtype=features.dtype), + } + + composer = ClassicalMMComposer(bond_head=_ConstBond()) + pos = torch.tensor([[0.0, 0.0, 0.0], [1.5, 0.0, 0.0]], dtype=torch.float64) + batch = TensorDict( + { + "atoms": TensorDict( + { + "pos": pos, + "Z": torch.tensor([1, 1]), + "batch": torch.zeros(2, dtype=torch.long), + }, + batch_size=[2], + ), + "bonds": TensorDict( + { + "atomi": torch.tensor([0], dtype=torch.long), + "atomj": torch.tensor([1], dtype=torch.long), + }, + batch_size=[1], + ), + }, + batch_size=[], + ) + features = {"bonds": torch.zeros(1, 1, dtype=torch.float64)} + ir = composer.parameterize(features, batch) + assert isinstance(ir, PotentialIR) + assert ir.unit_system == "class_i_canonical" + assert CLASS_I_CANONICAL["energy"] == "kcal/mol" + e = float(composer.energy(ir, batch, pos=pos)) + # 0.5 * 2.0 * (1.5 - 1.0)^2 = 0.25 + assert math.isclose(e, 0.25, abs_tol=1e-8), f"bond golden E={e}" + + # --- 4. Import boundary --- + root = Path(__file__).resolve().parents[1] / "src/molpot/composition" + for name in ("classical_mm.py", "mm_heads.py"): + tree = ast.parse((root / name).read_text()) + for node in ast.walk(tree): + if isinstance(node, ast.Import): + for alias in node.names: + assert not alias.name.startswith("molzoo") + assert not alias.name.startswith("molrep.chem") + elif isinstance(node, ast.ImportFrom) and node.module: + assert not node.module.startswith("molzoo") + assert not node.module.startswith("molrep.chem") + + print("learnable-classical-ff-03-mm-heads: all hard-coded goldens OK") + print(f" bond harmonic E = {e}") + return 0 + + +if __name__ == "__main__": + try: + raise SystemExit(main()) + except Exception as exc: # noqa: BLE001 — standalone regression script + print(f"FAIL: {exc}", file=sys.stderr) + raise diff --git a/regressions/learnable-classical-ff-04-chem-encoder.py b/regressions/learnable-classical-ff-04-chem-encoder.py new file mode 100644 index 0000000..5d1a61e --- /dev/null +++ b/regressions/learnable-classical-ff-04-chem-encoder.py @@ -0,0 +1,192 @@ +"""Public-API regression for continuous chemical perception encoder. + +Spec: `learnable-classical-ff-04-chem-encoder`. + +Hard-coded goldens only — no third-party oracle. Pins: + +1. AtomChemEmbedding Z → (N, D_a); reuses JointEmbedding +2. Bond / angle / proper reverse symmetries; improper outer-swap (center fixed) +3. ChemEncoder write_batch keys + ChemEmbeddings counts +4. molzoo ChemPerception recipe forwards without energy keys +5. No molpot imports under molrep.chem / molzoo.chem + +Provenance +---------- + capture command : PYTHONPATH=src python regressions/learnable-classical-ff-04-chem-encoder.py + date : 2026-08-10 +""" + +from __future__ import annotations + +import ast +import sys +from pathlib import Path + +import torch +from tensordict import TensorDict + + +def main() -> int: + from molrep.chem import ( + AngleContext, + AtomChemEmbedding, + BondChemEmbedding, + ChemEmbeddings, + ChemEncoder, + ImproperContext, + ProperContext, + ) + from molrep.embedding.node import JointEmbedding + from molzoo.chem import ChemPerception, ChemPerceptionSpec + + torch.manual_seed(0) + + # --- 1. Atom embedding --- + atom_emb = AtomChemEmbedding(atom_dim=8, num_elements=20) + assert isinstance(atom_emb.joint, JointEmbedding) + z = torch.tensor([8, 1, 1, 6], dtype=torch.long) + h = atom_emb(z) + assert h.shape == (4, 8) + + # --- 2. Bond symmetry --- + bond_emb = BondChemEmbedding(atom_dim=8, bond_dim=6) + atomi = torch.tensor([0, 0, 1], dtype=torch.long) + atomj = torch.tensor([1, 2, 3], dtype=torch.long) + assert torch.allclose( + bond_emb(h, atomi, atomj), + bond_emb(h, atomj, atomi), + rtol=1e-5, + atol=1e-6, + ) + + # --- 3. Angle reverse --- + ang = AngleContext(atom_dim=8, angle_dim=6) + ai = torch.tensor([1], dtype=torch.long) + aj = torch.tensor([0], dtype=torch.long) + ak = torch.tensor([2], dtype=torch.long) + assert torch.allclose(ang(h, ai, aj, ak), ang(h, ak, aj, ai), rtol=1e-5, atol=1e-6) + + # --- 4. Proper reverse --- + prop = ProperContext(atom_dim=8, proper_dim=6) + pi = torch.tensor([1], dtype=torch.long) + pj = torch.tensor([0], dtype=torch.long) + pk = torch.tensor([3], dtype=torch.long) + pl = torch.tensor([2], dtype=torch.long) + assert torch.allclose( + prop(h, pi, pj, pk, pl), + prop(h, pl, pk, pj, pi), + rtol=1e-5, + atol=1e-6, + ) + + # --- 5. Improper outer swap (center fixed at atomi) --- + imp = ImproperContext(atom_dim=8, improper_dim=6) + center = torch.tensor([0], dtype=torch.long) + j = torch.tensor([1], dtype=torch.long) + k = torch.tensor([2], dtype=torch.long) + l = torch.tensor([3], dtype=torch.long) + base = imp(h, center, j, k, l) + assert torch.allclose(base, imp(h, center, k, j, l), rtol=1e-5, atol=1e-6) + assert torch.allclose(base, imp(h, center, j, l, k), rtol=1e-5, atol=1e-6) + + # --- 6. ChemEncoder I/O --- + batch = TensorDict( + { + "atoms": TensorDict( + { + "Z": z, + "batch": torch.zeros(4, dtype=torch.long), + }, + batch_size=[4], + ), + "bonds": TensorDict( + { + "atomi": torch.tensor([0, 0, 0], dtype=torch.long), + "atomj": torch.tensor([1, 2, 3], dtype=torch.long), + }, + batch_size=[3], + ), + "angles": TensorDict( + { + "atomi": torch.tensor([1, 1], dtype=torch.long), + "atomj": torch.tensor([0, 0], dtype=torch.long), + "atomk": torch.tensor([2, 3], dtype=torch.long), + }, + batch_size=[2], + ), + "propers": TensorDict( + { + "atomi": pi, + "atomj": pj, + "atomk": pk, + "atoml": pl, + }, + batch_size=[1], + ), + "impropers": TensorDict( + { + "atomi": center, + "atomj": j, + "atomk": k, + "atoml": l, + }, + batch_size=[1], + ), + }, + batch_size=[], + ) + enc = ChemEncoder( + atom_dim=8, + bond_dim=6, + angle_dim=6, + proper_dim=6, + improper_dim=6, + num_elements=20, + ) + out = enc(batch) + emb = enc.embeddings(out) + assert isinstance(emb, ChemEmbeddings) + assert emb.atom.shape == (4, 8) + assert emb.bond.shape == (3, 6) + assert emb.angle.shape == (2, 6) + assert emb.proper.shape == (1, 6) + assert emb.improper.shape == (1, 6) + + # --- 7. molzoo recipe --- + spec = ChemPerceptionSpec( + atom_dim=8, + bond_dim=6, + angle_dim=6, + proper_dim=6, + improper_dim=6, + num_elements=20, + ) + model = ChemPerception(spec=spec) + out2 = model(batch) + assert out2["atoms", "chem_features"].shape == (4, 8) + nested = {str(k) for k in out2.keys(include_nested=True)} + assert not any("energy" in k for k in nested) + + # --- 8. Import boundary --- + root = Path(__file__).resolve().parents[1] / "src" + for rel in ("molrep/chem", "molzoo/chem"): + for path in (root / rel).rglob("*.py"): + tree = ast.parse(path.read_text()) + for node in ast.walk(tree): + if isinstance(node, ast.Import): + for alias in node.names: + assert not alias.name.startswith("molpot"), path + elif isinstance(node, ast.ImportFrom) and node.module: + assert not node.module.startswith("molpot"), path + + print("learnable-classical-ff-04-chem-encoder: all hard-coded goldens OK") + print(f" atom features {tuple(emb.atom.shape)}, bond {tuple(emb.bond.shape)}") + return 0 + + +if __name__ == "__main__": + try: + raise SystemExit(main()) + except Exception as exc: # noqa: BLE001 — standalone regression script + print(f"FAIL: {exc}", file=sys.stderr) + raise diff --git a/regressions/learnable-classical-ff-05-neural-parameterizer.py b/regressions/learnable-classical-ff-05-neural-parameterizer.py new file mode 100644 index 0000000..3dc82b8 --- /dev/null +++ b/regressions/learnable-classical-ff-05-neural-parameterizer.py @@ -0,0 +1,151 @@ +"""Public-API regression for ClassicalMMParameterizer (encoder → IR → E/F). + +Spec: `learnable-classical-ff-05-neural-parameterizer`. + +Hard-coded goldens only — no third-party oracle. Pins: + +1. FakeEncoder satisfies Protocol; no molzoo import in parameterizer +2. parameterize → CLASS_I IR (kcal/mol) +3. Bond harmonic energy golden: k=2, r=1.5, r0=1 → E = 0.25 kcal/mol +4. compute_forces=True → forces (N,3) matching analytic F = -∂E/∂r +5. Optional eV conversion at boundary only (IR stays class_i_canonical) + +Provenance +---------- + formula : E = ½ k (r − r₀)² (OpenMM §19 / Class-I) + geometry : two atoms on x-axis at distance r + capture command : PYTHONPATH=src python \ + regressions/learnable-classical-ff-05-neural-parameterizer.py + date : 2026-08-10 +""" + +from __future__ import annotations + +import ast +import math +import sys +from pathlib import Path +from typing import Mapping + +import torch +import torch.nn as nn +from tensordict import TensorDict + + +class _FakeEmbeddings: + def __init__(self, features: Mapping[str, torch.Tensor]) -> None: + self._features = dict(features) + + def interaction_dict(self) -> dict[str, torch.Tensor]: + return dict(self._features) + + +class FakeEncoder(nn.Module): + def __init__(self, features: Mapping[str, torch.Tensor]) -> None: + super().__init__() + self._features = dict(features) + self._dummy = nn.Parameter(torch.zeros(1)) + + def forward(self, td: TensorDict) -> TensorDict: + return td + + def embeddings(self, td: TensorDict) -> _FakeEmbeddings: + return _FakeEmbeddings(self._features) + + +class _ConstBond(nn.Module): + def forward(self, features: torch.Tensor) -> dict[str, torch.Tensor]: + n = features.shape[0] + return { + "k": torch.full((n,), 2.0, dtype=features.dtype), + "r0": torch.full((n,), 1.0, dtype=features.dtype), + } + + +def main() -> int: + from molpot import ( + KCAL_MOL_TO_EV, + ClassicalMMComposer, + ClassicalMMParameterizer, + PotentialIR, + energy_kcal_to_ev, + ) + from molpot.composition.parameterizer import ChemEncoderProtocol + from molpot.ir import CLASS_I_CANONICAL + + torch.manual_seed(0) + + # --- 1. Protocol + import boundary --- + features = {"bonds": torch.zeros(1, 4, dtype=torch.float64)} + encoder = FakeEncoder(features) + assert isinstance(encoder, ChemEncoderProtocol) + + path = Path(__file__).resolve().parents[1] / "src/molpot/composition/parameterizer.py" + tree = ast.parse(path.read_text()) + for node in ast.walk(tree): + if isinstance(node, ast.Import): + for alias in node.names: + assert not alias.name.startswith("molzoo") + elif isinstance(node, ast.ImportFrom) and node.module: + assert not node.module.startswith("molzoo") + + # --- 2. Parameterize IR --- + composer = ClassicalMMComposer(bond_head=_ConstBond()) + param = ClassicalMMParameterizer(encoder=encoder, composer=composer) + pos = torch.tensor([[0.0, 0.0, 0.0], [1.5, 0.0, 0.0]], dtype=torch.float64) + batch = TensorDict( + { + "atoms": TensorDict( + { + "pos": pos, + "Z": torch.tensor([1, 1]), + "batch": torch.zeros(2, dtype=torch.long), + }, + batch_size=[2], + ), + "bonds": TensorDict( + { + "atomi": torch.tensor([0], dtype=torch.long), + "atomj": torch.tensor([1], dtype=torch.long), + }, + batch_size=[1], + ), + }, + batch_size=[], + ) + ir = param.parameterize(batch) + assert isinstance(ir, PotentialIR) + assert ir.unit_system == "class_i_canonical" + assert CLASS_I_CANONICAL["energy"] == "kcal/mol" + + # --- 3. Energy golden --- + e = float(param.energy(batch, ir=ir, pos=pos)) + assert math.isclose(e, 0.25, abs_tol=1e-8), f"bond golden E={e}" + + # --- 4. Forces --- + out = param.forward(batch, compute_forces=True) + forces = out["forces"] + assert forces.shape == (2, 3) + # F0_x = +k*(r-r0) = 1.0, F1_x = -1.0 + assert math.isclose(float(forces[0, 0]), 1.0, abs_tol=1e-5) + assert math.isclose(float(forces[1, 0]), -1.0, abs_tol=1e-5) + assert "energy" in batch["graphs"] + assert "forces" in batch["atoms"] + + # --- 5. Units boundary --- + e_ev = float(energy_kcal_to_ev(torch.tensor(e, dtype=torch.float64))) + assert math.isclose(e_ev, e * KCAL_MOL_TO_EV, abs_tol=1e-9) + assert ir.unit_system == "class_i_canonical" + + print("learnable-classical-ff-05-neural-parameterizer: all hard-coded goldens OK") + print(f" bond harmonic E = {e} kcal/mol") + print(f" F0_x = {float(forces[0, 0])}") + return 0 + + +if __name__ == "__main__": + try: + raise SystemExit(main()) + except Exception as exc: # noqa: BLE001 — standalone regression script + print(f"FAIL: {exc}", file=sys.stderr) + raise diff --git a/regressions/learnable-classical-ff-06-condensation.py b/regressions/learnable-classical-ff-06-condensation.py new file mode 100644 index 0000000..86e262a --- /dev/null +++ b/regressions/learnable-classical-ff-06-condensation.py @@ -0,0 +1,173 @@ +"""Public-API regression for physics-aware chemical class condensation. + +Spec: `learnable-classical-ff-06-condensation`. + +Hard-coded goldens only — no third-party oracle. Pins: + +1. MergeCriterion bond budgets (Å / relative k) accept/reject +2. Condenser: identical params → 1 type; far params → 2 types +3. Multi-system merge shares global type ids +4. TypeSystem.assign soft-fails out-of-budget rows (UNMATCHED_TYPE_ID) +5. TypeSystemLabeler satisfies Labeler Protocol +6. MultiTypeHead additive; TypeHead atom API unchanged +7. No SMARTS emitters under molrep.condensation; physical_eval is injected + +Provenance +---------- + capture command : PYTHONPATH=src python \ + regressions/learnable-classical-ff-06-condensation.py + date : 2026-08-10 +""" + +from __future__ import annotations + +import ast +import sys +from pathlib import Path + +import torch + + +def main() -> int: + from molrep.condensation import ( + ClassAssignment, + Condenser, + InteractionClass, + PhysicalErrorMetrics, + TypeSystem, + TypeSystemLabeler, + UNMATCHED_TYPE_ID, + bond_default_criterion, + ) + from molrep.heads import Labeler, MultiTypeHead, TypeHead + + # --- 1. MergeCriterion --- + crit = bond_default_criterion() + assert crit.interaction is InteractionClass.BOND + assert crit.abs_tol["r0"] == 0.01 + assert crit.accepts( + {"k": torch.tensor(300.0), "r0": torch.tensor(1.09)}, + {"k": torch.tensor(310.0), "r0": torch.tensor(1.095)}, + ) + assert not crit.accepts( + {"k": torch.tensor(300.0), "r0": torch.tensor(1.09)}, + {"k": torch.tensor(300.0), "r0": torch.tensor(1.12)}, + ) + + # --- 2. Identical → 1 type --- + condenser = Condenser() + identical = { + "k": torch.tensor([300.0, 300.0, 300.0]), + "r0": torch.tensor([1.09, 1.09, 1.09]), + } + r_id = condenser.merge( + [identical], + interaction=InteractionClass.BOND, + criterion=crit, + ) + assert r_id.type_system.n_types == 1 + assert r_id.type_system.get(0).member_count == 3 + assert isinstance(r_id.assignment, ClassAssignment) + assert isinstance(r_id.metrics, PhysicalErrorMetrics) + + # --- 3. Far params → 2 types (stable ids) --- + far = { + "k": torch.tensor([300.0, 300.0]), + "r0": torch.tensor([1.09, 1.5]), + } + r_far = condenser.greedy_merge( + [far], + interaction=InteractionClass.BOND, + criterion=crit, + ) + assert r_far.type_system.n_types == 2 + assert r_far.assignment.type_ids.tolist() == [0, 1] + + # --- 4. Multi-system global ids --- + sys_a = {"k": torch.tensor([300.0]), "r0": torch.tensor([1.09])} + sys_b = { + "k": torch.tensor([301.0, 500.0]), + "r0": torch.tensor([1.091, 1.5]), + } + r_ms = condenser.merge( + [sys_a, sys_b], + interaction=InteractionClass.BOND, + criterion=crit, + ) + assert r_ms.type_system.n_types == 2 + assert r_ms.assignment.for_system(0).tolist() == [0] + assert r_ms.assignment.for_system(1).tolist() == [0, 1] + + # --- 5. assign soft-fail --- + ts = TypeSystem.from_prototypes( + InteractionClass.BOND, + [{"k": 300.0, "r0": 1.09}], + criterion=crit, + ) + assert ts.assign({"k": 300.0, "r0": 1.09}) == 0 + assert ts.assign({"k": 300.0, "r0": 1.5}) == UNMATCHED_TYPE_ID + + # --- 6. Labeler Protocol --- + labeler = TypeSystemLabeler(r_ms.type_system) + assert isinstance(labeler, Labeler) + assert labeler.num_types == 2 + assert set(labeler.type_map) == {0, 1} + ids = labeler.label( + { + "k": torch.tensor([300.0, 500.0]), + "r0": torch.tensor([1.09, 1.5]), + } + ) + assert ids.tolist() == [0, 1] + + # --- 7. TypeHead / MultiTypeHead --- + head = TypeHead(hidden_dim=4, num_types=5) + assert head(torch.ones(2, 4)).shape == (2, 5) + multi = MultiTypeHead.from_type_systems( + 4, + {"bond": r_ms.type_system}, + ) + assert multi.num_types["bond"] == 2 + assert multi({"bond": torch.randn(3, 4)})["bond"].shape == (3, 2) + + # --- 8. Physics gate via injected callable (no molpot energy in package) --- + def hot(_p, _c, _i): + return 5.0 + + r_phys = condenser.merge( + [ + { + "k": torch.tensor([300.0, 300.0]), + "r0": torch.tensor([1.09, 1.09]), + } + ], + interaction=InteractionClass.BOND, + criterion=crit, + physical_eval=hot, + energy_tol=0.1, + ) + assert r_phys.type_system.n_types == 2 + assert r_phys.metrics.rejected_by_physics == 1 + + # --- 9. No SMARTS emitters; no molpot energy imports in condensation --- + pkg = Path(__file__).resolve().parents[1] / "src/molrep/condensation" + for path in sorted(pkg.glob("*.py")): + text = path.read_text() + lower = text.lower() + assert "def to_smarts" not in lower + assert "def emit_smarts" not in lower + assert "def to_smirks" not in lower + tree = ast.parse(text) + for node in ast.walk(tree): + if isinstance(node, ast.Import): + for alias in node.names: + assert not alias.name.startswith("molpot") + elif isinstance(node, ast.ImportFrom) and node.module: + assert not node.module.startswith("molpot") + + print("learnable-classical-ff-06-condensation: OK") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/regressions/learnable-classical-ff-07-smarts.py b/regressions/learnable-classical-ff-07-smarts.py new file mode 100644 index 0000000..7361ece --- /dev/null +++ b/regressions/learnable-classical-ff-07-smarts.py @@ -0,0 +1,150 @@ +"""Public-API regression for symbolic SMARTS perception + SymbolicForceField. + +Spec: `learnable-classical-ff-07-smarts`. + +Hard-coded goldens only — no third-party oracle. Pins: + +1. SymbolicPattern validates non-empty pattern + arity in {1,2,3,4} +2. FakeSmartsMatcher returns configured [arity, K] hits +3. ClassPatternRegistry bind/get + conflict detection + reverse lookup +4. SymbolicForceField.records pairs prototypes with SMARTS +5. match_molecule assigns type ids via FakeSmartsMatcher (no energy) +6. SmartsMatcher Protocol satisfied by Fake (and Molpy when available) +7. molrep.perception has zero molrs / molpot imports + +Provenance +---------- + capture command : PYTHONPATH=src python regressions/learnable-classical-ff-07-smarts.py + date : 2026-08-10 +""" + +from __future__ import annotations + +import ast +import sys +from pathlib import Path + +import torch + +ROOT = Path(__file__).resolve().parents[1] +PERCEPTION = ROOT / "src" / "molrep" / "perception" + + +def _assert_no_forbidden_imports(package_dir: Path, forbidden: tuple[str, ...]) -> None: + for path in sorted(package_dir.glob("*.py")): + tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path)) + for node in ast.walk(tree): + if isinstance(node, ast.Import): + for alias in node.names: + for bad in forbidden: + assert not alias.name.startswith(bad), ( + f"{path.name} imports {alias.name}" + ) + elif isinstance(node, ast.ImportFrom): + mod = node.module or "" + for bad in forbidden: + assert not mod.startswith(bad), f"{path.name} imports from {mod}" + + +def main() -> int: + from molrep.condensation import InteractionClass, TypeRecord, TypeSystem + from molrep.perception import ( + ClassPatternRegistry, + DiscreteClassRecord, + FakeSmartsMatcher, + SmartsMatcher, + SymbolicForceField, + SymbolicPattern, + ) + + # --- 1. SymbolicPattern validation --- + ok = SymbolicPattern("[#6]-[#6]", arity=2) + assert ok.arity == 2 + try: + SymbolicPattern("", arity=2) + raise AssertionError("empty pattern should raise") + except ValueError: + pass + try: + SymbolicPattern("[#6]", arity=0) + raise AssertionError("arity 0 should raise") + except ValueError: + pass + try: + SymbolicPattern("[#6]", arity=5) + raise AssertionError("arity 5 should raise") + except ValueError: + pass + + # --- 2. FakeSmartsMatcher --- + hits = torch.tensor([[0, 2], [1, 3]], dtype=torch.long) + matcher = FakeSmartsMatcher({"[#6]-[#6]": hits}) + assert isinstance(matcher, SmartsMatcher) + out = matcher.match(None, ok) + assert out.shape == (2, 2) + assert torch.equal(out, hits) + empty = matcher.match(None, SymbolicPattern("[#8]", arity=1)) + assert empty.shape == (1, 0) + + # --- 3. Registry --- + reg = ClassPatternRegistry() + pat_cc = SymbolicPattern("[#6]-[#6]", arity=2) + pat_co = SymbolicPattern("[#6]-[#8]", arity=2) + reg.bind(InteractionClass.BOND, 0, pat_cc) + reg.bind(InteractionClass.BOND, 1, pat_co) + assert reg.get(InteractionClass.BOND, 0).pattern == "[#6]-[#6]" + assert reg.reverse_lookup("[#6]-[#8]") == (InteractionClass.BOND, 1) + try: + reg.bind(InteractionClass.BOND, 0, SymbolicPattern("[#7]-[#7]", arity=2)) + raise AssertionError("conflicting bind should raise") + except ValueError: + pass + + # --- 4–5. SymbolicForceField --- + ts = TypeSystem( + InteractionClass.BOND, + [ + TypeRecord(type_id=0, prototype={"k": 300.0, "r0": 1.09}, label="CT-CT"), + TypeRecord(type_id=1, prototype={"k": 320.0, "r0": 1.41}, label="CT-OH"), + ], + ) + ff = SymbolicForceField({InteractionClass.BOND: ts}, reg) + recs = ff.records() + assert len(recs) == 2 + assert all(isinstance(r, DiscreteClassRecord) for r in recs) + by_id = {r.type_id: r for r in recs} + assert by_id[0].smarts == "[#6]-[#6]" + assert by_id[0].prototype == {"k": 300.0, "r0": 1.09} + assert by_id[1].smarts == "[#6]-[#8]" + + fake = FakeSmartsMatcher( + { + "[#6]-[#6]": torch.tensor([[0], [1]], dtype=torch.long), + "[#6]-[#8]": torch.tensor([[1], [2]], dtype=torch.long), + } + ) + assigned = ff.match_molecule(mol=object(), matcher=fake) + bond = assigned[InteractionClass.BOND] + assert bond["matches"].shape == (2, 2) + assert bond["type_ids"].tolist() == [0, 1] + assert bond["matches"][:, 0].tolist() == [0, 1] + assert bond["matches"][:, 1].tolist() == [1, 2] + + # --- 6. Molpy matcher protocol (optional smoke) --- + try: + from molrep.perception import MolpySmartsMatcher + + molpy_matcher = MolpySmartsMatcher() + assert isinstance(molpy_matcher, SmartsMatcher) + except ImportError: + pass + + # --- 7. Import hard rules --- + _assert_no_forbidden_imports(PERCEPTION, ("molrs", "molpot")) + + print("learnable-classical-ff-07-smarts: all hard-coded goldens OK") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/regressions/learnable-classical-ff-08-ff-export.py b/regressions/learnable-classical-ff-08-ff-export.py new file mode 100644 index 0000000..0393dd5 --- /dev/null +++ b/regressions/learnable-classical-ff-08-ff-export.py @@ -0,0 +1,117 @@ +"""Public-API regression for Potential IR → OpenMM force-spec export. + +Spec: `learnable-classical-ff-08-ff-export`. + +Hard-coded goldens only — no live OpenMM / third-party oracle. Pins: + +1. Bond force constant: k=100 kcal mol⁻¹ Å⁻² → 41840 kJ mol⁻¹ nm⁻² +2. AMBER torsion Vn=2 kcal/mol → OpenMM PeriodicTorsion k=4.184 kJ/mol +3. Class-I NonbondedScaling 1–4 defaults flow into ForceSpec.scaling +4. Unsupported improper_harmonic raises UnsupportedTermError (no silent drop) +5. ForceSpec.to_dict is JSON-serializable + +Provenance +---------- + formula : OpenMM User Guide §19; k_omm_bond = k_ir * 4.184 / 0.01; + AMBER E=(Vn/2)[1+cos] → OpenMM k=(Vn/2)*4.184 + capture command : PYTHONPATH=src python regressions/learnable-classical-ff-08-ff-export.py + date : 2026-08-10 + note : Pure unit/form translation; no dynamics / energy compare. +""" + +from __future__ import annotations + +import json +import math +import sys + +import torch + + +def main() -> int: + from molix.ff_export import ( + ForceFieldCompiler, + TranslationCase, + UnsupportedTermError, + scale_amber_vn, + scale_bond_k, + ) + from molpot.ir import ( + BondBag, + ChargeBag, + ImproperHarmonicBag, + LJBag, + NonbondedScaling, + PotentialIR, + ProperTorsionBag, + ) + + # --- 1. Bond k golden --- + assert scale_bond_k(100.0) == 41840.0, f"bond k golden: {scale_bond_k(100.0)}" + + # --- 2. Torsion Vn=2 → k_omm=4.184 --- + assert scale_amber_vn(2.0) == 4.184, f"Vn golden: {scale_amber_vn(2.0)}" + + ir = PotentialIR( + bonds=BondBag( + k=torch.tensor([100.0], dtype=torch.float64), + r0=torch.tensor([1.5], dtype=torch.float64), + ), + propers=ProperTorsionBag( + # IR half-barrier for AMBER Vn=2 (E=(k/s)[1+cos], s=1, k=1) + k=torch.tensor([[1.0]], dtype=torch.float64), + periodicity=torch.tensor([2], dtype=torch.long), + phase=torch.tensor([[0.0]], dtype=torch.float64), + idivf=torch.tensor([1.0], dtype=torch.float64), + ), + lj=LJBag( + epsilon=torch.tensor([0.1], dtype=torch.float64), + sigma=torch.tensor([3.5], dtype=torch.float64), + ), + charges=ChargeBag(q=torch.tensor([0.5, -0.5], dtype=torch.float64)), + scaling=NonbondedScaling(), + ) + + compiler = ForceFieldCompiler("openmm") + spec = compiler.compile(ir) + + bond = next(f for f in spec.forces if f["type"] == "HarmonicBondForce") + assert bond["parameters"][0]["k"] == 41840.0 + assert math.isclose(bond["parameters"][0]["r0"], 0.15, abs_tol=1e-15) + assert bond["case"] == TranslationCase.DIRECT_UNIT_SCALE.value + + torsion = next(f for f in spec.forces if f["type"] == "PeriodicTorsionForce") + assert torsion["parameters"][0]["k"] == 4.184 + assert torsion["parameters"][0]["periodicity"] == 2 + + # --- 3. 1–4 scales --- + assert spec.scaling is not None + assert math.isclose(spec.scaling["scale_q_14"], 5.0 / 6.0, abs_tol=1e-15) + assert math.isclose(spec.scaling["scale_lj_14"], 0.5, abs_tol=1e-15) + + # --- 4. Unsupported improper harmonic --- + bad = PotentialIR( + impropers_harmonic=ImproperHarmonicBag( + k=torch.tensor([1.0], dtype=torch.float64), + chi0=torch.tensor([0.0], dtype=torch.float64), + ) + ) + try: + compiler.compile(bad) + except UnsupportedTermError as err: + assert err.term == "improper_harmonic" + assert err.case is TranslationCase.UNSUPPORTED + else: + raise AssertionError("expected UnsupportedTermError for improper_harmonic") + + # --- 5. JSON serializable --- + payload = json.dumps(spec.to_dict()) + assert "HarmonicBondForce" in payload + assert "41840" in payload + + print("OK learnable-classical-ff-08-ff-export") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/regressions/learnable-classical-ff-09-provenance.py b/regressions/learnable-classical-ff-09-provenance.py new file mode 100644 index 0000000..ed47b1a --- /dev/null +++ b/regressions/learnable-classical-ff-09-provenance.py @@ -0,0 +1,172 @@ +"""Public-API regression for confidence / coverage / provenance surfaces. + +Spec: `learnable-classical-ff-09-provenance`. + +Hard-coded goldens only — no third-party oracle. Pins: + +1. ChemicalSupportIndex kNN L2 membership + coverage_fraction +2. from_type_ids exact set membership +3. SupportClassifier regime policy (in / near / extrapolating / unknown) +4. TypeHead.decode_with_confidence → classify (no local softmax) +5. ParameterProvenance frozen + as_dict +6. No molpot import under molrep.embedding.support; no AL APIs + +Provenance +---------- + capture command : PYTHONPATH=src python \ + regressions/learnable-classical-ff-09-provenance.py + date : 2026-08-10 +""" + +from __future__ import annotations + +import ast +import math +import sys +from dataclasses import FrozenInstanceError +from pathlib import Path + +import torch + + +def main() -> int: + from molpot.heads.provenance import ( + CoverageRegime, + ParameterProvenance, + SupportClassifier, + attach_provenance, + ) + from molpot.heads.type import TypeHead + from molrep.embedding.support import ChemicalSupportIndex + + torch.manual_seed(0) + + # --- 1. Continuous kNN L2 bank --- + bank = torch.tensor( + [ + [0.0, 0.0], + [1.0, 0.0], + [0.0, 1.0], + ], + dtype=torch.float64, + ) + index = ChemicalSupportIndex(bank, radius=0.5, k=1) + query = torch.tensor( + [ + [0.0, 0.0], # in + [0.3, 0.0], # in + [2.0, 2.0], # out + ], + dtype=torch.float64, + ) + mask = index.contains(query) + assert mask.tolist() == [True, True, False], mask + cov = index.coverage_fraction(query) + assert math.isclose(cov, 2 / 3, rel_tol=0, abs_tol=1e-6), cov + dist, nn_idx = index.knn(query[:1]) + assert int(nn_idx[0, 0]) == 0 + assert math.isclose(float(dist[0, 0]), 0.0, abs_tol=1e-6) + + # --- 2. Discrete type-id bank --- + tindex = ChemicalSupportIndex.from_type_ids({0, 2, 5}, radius=0.0) + ids = torch.tensor([0, 1, 2, 5, 7], dtype=torch.long) + assert tindex.contains_type_ids(ids).tolist() == [True, False, True, True, False] + assert math.isclose(tindex.coverage_fraction_type_ids(ids), 0.6, abs_tol=1e-6) + + # --- 3. Classifier policy --- + clf = SupportClassifier(tindex, conf_in=0.8, conf_near=0.5) + regimes = clf.classify( + torch.tensor([0, 0, 0, 99, 99], dtype=torch.long), + torch.tensor([0.95, 0.65, 0.20, 0.99, 0.10]), + ) + assert regimes == [ + CoverageRegime.IN_SUPPORT, + CoverageRegime.NEAR_SUPPORT, + CoverageRegime.UNKNOWN, + CoverageRegime.EXTRAPOLATING, + CoverageRegime.UNKNOWN, + ], regimes + + # confidence-only (no support) + clf_free = SupportClassifier(None, conf_in=0.8, conf_near=0.5) + assert clf_free.classify( + torch.tensor([1], dtype=torch.long), + torch.tensor([0.9]), + ) == [CoverageRegime.IN_SUPPORT] + + # --- 4. TypeHead.decode_with_confidence reuse --- + head = TypeHead(hidden_dim=4, num_types=3) + logits = torch.tensor( + [ + [10.0, 0.0, 0.0], + [0.0, 3.0, 0.0], + ], + dtype=torch.float64, + ) + indices, confidence = head.decode_with_confidence(logits) + assert int(indices[0]) == 0 + assert float(confidence[0]) > 0.8 + regimes2 = clf.classify(indices, confidence) + assert regimes2[0] is CoverageRegime.IN_SUPPORT + + # provenance package must not reimplement softmax + prov_root = ( + Path(__file__).resolve().parents[1] + / "src" + / "molpot" + / "heads" + / "provenance" + ) + for py in prov_root.glob("*.py"): + text = py.read_text() + assert "torch.softmax" not in text and "F.softmax" not in text, py + + # --- 5. ParameterProvenance --- + rec = ParameterProvenance( + interaction="bond", + type_id=int(indices[0]), + confidence=float(confidence[0]), + regime=regimes2[0], + source="condensed_type", + pattern=None, + ) + d = rec.as_dict() + assert d["regime"] == "IN_SUPPORT" + assert d["ir_units"] == "class_i_canonical" + try: + rec.type_id = 99 # type: ignore[misc] + raise AssertionError("ParameterProvenance must be frozen") + except (FrozenInstanceError, AttributeError): + pass + + meta: dict = {} + attach_provenance(meta, [rec]) + assert meta["provenance"][0]["source"] == "condensed_type" + + # --- 6. Import / AL boundaries --- + support_src = ( + Path(__file__).resolve().parents[1] + / "src" + / "molrep" + / "embedding" + / "support.py" + ) + tree = ast.parse(support_src.read_text()) + for node in ast.walk(tree): + if isinstance(node, ast.Import): + for alias in node.names: + assert not alias.name.startswith("molpot"), alias.name + if isinstance(node, ast.ImportFrom) and node.module: + assert not node.module.startswith("molpot"), node.module + + for py in list(prov_root.glob("*.py")) + [support_src]: + text = py.read_text().lower() + for bad in ("active_learning", "acquisition_function", "query_selector"): + assert bad not in text, (py, bad) + + print("learnable-classical-ff-09-provenance: all hard-coded goldens OK") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/regressions/mace-subpackage-restructure-01-blocks.py b/regressions/mace-subpackage-restructure-01-blocks.py new file mode 100644 index 0000000..4042f4a --- /dev/null +++ b/regressions/mace-subpackage-restructure-01-blocks.py @@ -0,0 +1,387 @@ +"""Numerics-neutrality smoke for chain step mace-subpackage-restructure-01. + +Builds the eight MACE blocks that the step relocates — through their **new** +public import paths — on a fixed 4-node / 5-edge graph with deterministic +``linspace`` inputs, and asserts every block's output ``sum()`` / ``abs().max()`` +and ``sorted(state_dict().keys())`` against literals captured on the pre-move +parent commit. A pure move must reproduce them bit-for-bit; any drift here means +the "verbatim relocation" claim is false. + +Goldens +------- +Captured by running this script's construction sequence with the OLD import +paths (``molrep.interaction.density`` / ``molrep.interaction.product`` / +``molrep.readout.scalar`` / ``molrep.readout.product`` / ``molzoo.mace``) and +dumping JSON instead of asserting — same seeds, same inputs, same ``.double()`` +calls. To re-capture, swap the imports below back and print the statistics. + + capture command : PYTHONPATH=src python capture_01_goldens.py (old-path variant) + parent commit : cf60f99c23b8b322eed14d8c9a6a37e833482787 + torch : 2.12.1+cpu + date : 2026-08-08 + device / dtype : CPU, float64 (``config.set_precision("fp64")`` + ``.double()``) + oracle : self-baseline (molnex at cf60f99) — no third-party oracle + + goldens re-captured 2026-08-09 at e8d6595 + working-tree dtype/init fixes: + ``_ScalarO3Linear`` N(0,1) init + ``config.ftype`` at construction (see + commit message); previous values captured at cf60f99 (2026-08-08), which + reproduces them bit-for-bit. Three blocks moved, each for one reason: + + * ``InteractionBlock`` — ``radial_mlp`` linears were built at the torch + default fp32 and cast by ``.double()`` afterwards; they are now built at + ``config.ftype`` directly, so ``kaiming_uniform_`` draws a different + (fp64) stream from the same ``manual_seed(0)``. + * ``NonLinearBiasReadout`` — ``_ScalarO3Linear`` now draws its weight from + ``N(0, 1)`` instead of zero, so the untrained readout is no longer + identically zero. The bias is still zero-initialised. + * ``EmbeddingBlock`` — same fp32→fp64 init-stream shift in the node + embedding. Only the ``node_features`` term of the composite checksum + moved (−7.378921712535659 → 18.411832194168696); the deterministic + ``edge_angular`` (5.0) and ``edge_radial`` (3.168506132872447) terms are + bit-identical, and the weight std stays ≈1 (1.032 → 0.995), so the large + swing in the *total* is cancellation bookkeeping, not a scale change. + + Unmoved, and re-verified: every block's ``sorted(state_dict().keys())``, + and the five blocks whose parameters were already fp64 at construction. + +Run: + PYTHONPATH=src python regressions/mace-subpackage-restructure-01-blocks.py +""" + +from __future__ import annotations + +import sys +from typing import NamedTuple + +import torch + +from molix import config + +config.set_precision("fp64") + +from molrep.embedding.mace import EmbeddingBlock # noqa: E402 +from molrep.embedding.node import DiscreteEmbeddingSpec # noqa: E402 +from molrep.interaction.mace.block import InteractionBlock # noqa: E402 +from molrep.interaction.mace.conv import ConvTP # noqa: E402 +from molrep.interaction.mace.density import ( # noqa: E402 + DensityInteraction, + DensityResidualInteraction, +) +from molrep.readout.mace import ( # noqa: E402 + LinearReadout, + NonLinearBiasReadout, + NonLinearReadout, + ProductHead, +) + +RTOL = 1e-12 +ATOL = 1e-15 # absolute floor so a golden that lands near zero stays comparable + +N, E, F, NUM_RADIAL, L_MAX = 4, 5, 8, 5, 1 +SH = "1x0e+1x1o" +TARGET = f"{F}x0e+{F}x1o" +HIDDEN_DIM = 32 # cue.Irreps("O3", irreps_from_l_max(1, 8)).dim = 8*1 + 8*3 + + +class Golden(NamedTuple): + """One block's captured signature: output statistics + state_dict key list.""" + + total: float + absmax: float + keys: list[str] + + +GOLDENS: dict[str, Golden] = { + "ConvTP": Golden( + total=-3.8099331628376785, + absmax=0.125, + keys=[ + "cue_tp.f.m.graphs.0.graph.c0", + "cue_tp.f.m.graphs.0.graph.c1", + ], + ), + "InteractionBlock": Golden( + total=0.07905200345001925, + absmax=0.07585597226756544, + keys=[ + "conv_tp.cue_tp.f.m.graphs.0.graph.c0", + "conv_tp.cue_tp.f.m.graphs.0.graph.c1", + "linear.f.m.graphs.0.graph.c0", + "linear.f.m.graphs.0.graph.c1", + "linear.weight", + "node_linear.f.m.graphs.0.graph.c0", + "node_linear.weight", + "radial_mlp.mlp.0.bias", + "radial_mlp.mlp.0.weight", + "radial_mlp.mlp.2.bias", + "radial_mlp.mlp.2.weight", + "radial_mlp.mlp.4.bias", + "radial_mlp.mlp.4.weight", + ], + ), + "DensityInteraction": Golden( + total=-0.08267766358519069, + absmax=0.04609357899752259, + keys=[ + "conv_tp.f.m.graphs.0.graph.c0", + "conv_tp.f.m.graphs.0.graph.c1", + "conv_tp_weights.layer0.weight", + "conv_tp_weights.layer1.weight", + "density_fn.layer0.weight", + "linear.f.m.graphs.0.graph.c0", + "linear.f.m.graphs.0.graph.c1", + "linear.weight", + "linear_up.f.m.graphs.0.graph.c0", + "linear_up.weight", + "skip_tp.f.m.graphs.0.graph.c0", + "skip_tp.f.m.graphs.0.graph.c1", + "skip_tp.weight", + ], + ), + "DensityResidualInteraction": Golden( + total=0.052798278602837645, + absmax=0.09087984372039665, + keys=[ + "conv_tp.f.m.graphs.0.graph.c0", + "conv_tp.f.m.graphs.0.graph.c1", + "conv_tp_weights.layer0.weight", + "conv_tp_weights.layer1.weight", + "density_fn.layer0.weight", + "linear.f.m.graphs.0.graph.c0", + "linear.f.m.graphs.0.graph.c1", + "linear.weight", + "linear_up.f.m.graphs.0.graph.c0", + "linear_up.weight", + "skip_tp.f.m.graphs.0.graph.c0", + "skip_tp.weight", + ], + ), + "LinearReadout": Golden( + total=-0.017091029094867177, + absmax=0.2238684806385588, + keys=[ + "linear.f.m.graphs.0.graph.c0", + "linear.weight", + ], + ), + "NonLinearReadout": Golden( + total=-0.1969420757337682, + absmax=0.4829924729613344, + keys=[ + "linear_1.f.m.graphs.0.graph.c0", + "linear_1.weight", + "linear_2.f.m.graphs.0.graph.c0", + "linear_2.weight", + ], + ), + "NonLinearBiasReadout": Golden( + # _ScalarO3Linear draws its weight from N(0, 1) and zero-initialises its + # bias, so the untrained OMOL-variant readout is non-trivial. It used to + # be identically zero (zero weight *and* zero bias), which silently + # zeroed every downstream force — that is what the init fix removed. + total=-0.04351653087296714, + absmax=0.04499837160305635, + keys=[ + "linear_1.f.m.graphs.0.graph.c0", + "linear_1.weight", + "linear_2.bias", + "linear_2.weight", + "linear_mid.bias", + "linear_mid.weight", + ], + ), + "ProductHead": Golden( + total=1.29665946085198, + absmax=0.6573909647520318, + keys=[ + "linear.bias", + "linear.weight", + "symmetric_contraction.symmetric_contraction.f.m.graphs.0.graph.c0", + "symmetric_contraction.symmetric_contraction.f.m.graphs.1.graph.c0", + "symmetric_contraction.symmetric_contraction.f.m.graphs.1.graph.c1", + "symmetric_contraction.symmetric_contraction.f.m.graphs.1.graph.c2", + "symmetric_contraction.symmetric_contraction.f.m.graphs.1.graph.c3", + "symmetric_contraction.symmetric_contraction.projection", + "symmetric_contraction.symmetric_contraction.weight", + ], + ), + "EmbeddingBlock": Golden( + total=26.580338327041144, + absmax=2.854573509905571, + keys=[ + "node_embedding.embedders.0.weight", + "node_embedding.project.0.f.m.graphs.0.graph.c0", + "node_embedding.project.0.weight", + ], + ), +} + + +def lin(*shape: int, lo: float = -0.5, hi: float = 0.5) -> torch.Tensor: + """Deterministic ramp input — no RNG, so the goldens are reproducible.""" + n = 1 + for d in shape: + n *= d + return torch.linspace(lo, hi, n, dtype=torch.float64).reshape(*shape) + + +class Checker: + """Collects per-block deviations against the embedded goldens.""" + + def __init__(self) -> None: + self.failures: list[str] = [] + + def check(self, name: str, module: torch.nn.Module, stats: tuple[float, float]) -> None: + """Compare one block's (sum, absmax) and state_dict keys to its golden.""" + golden = GOLDENS[name] + worst = 0.0 + for label, actual, expected in ( + ("sum", stats[0], golden.total), + ("absmax", stats[1], golden.absmax), + ): + dev = abs(actual - expected) + worst = max(worst, dev) + if dev > RTOL * abs(expected) + ATOL: + self.failures.append( + f"{name}.{label}: got {actual!r}, want {expected!r} (deviation {dev:.3e})" + ) + keys = sorted(module.state_dict().keys()) + if keys != golden.keys: + self.failures.append( + f"{name}.state_dict keys drifted:\n got {keys}\n want {golden.keys}" + ) + print(f" {name:<28} max deviation {worst:.3e}") + + def tensor( + self, + name: str, + module: torch.nn.Module, + result: torch.Tensor | tuple[torch.Tensor | None, ...], + ) -> None: + """Record a block whose output is a tensor (or a tuple whose first item is).""" + t = result if isinstance(result, torch.Tensor) else result[0] + assert isinstance(t, torch.Tensor) + t = t.detach() + self.check(name, module, (float(t.sum()), float(t.abs().max()))) + + +def main() -> int: + """Rebuild every relocated block through its new path and verify the goldens.""" + checker = Checker() + + edge_index = torch.tensor([[0, 1, 2, 3, 0], [1, 2, 3, 0, 2]]).t().contiguous() # (E, 2) + node_attrs = torch.zeros(N, 3, dtype=torch.float64) + node_attrs[torch.arange(N), torch.tensor([0, 1, 2, 0])] = 1.0 + node_feats = lin(N, F) + edge_attrs = lin(E, 4) + edge_feats = lin(E, NUM_RADIAL) + Z = torch.tensor([1, 6, 8, 1]) + edge_dist = torch.linspace(0.8, 2.5, E, dtype=torch.float64) + edge_diff = lin(E, 3, lo=-1.0, hi=1.0) + + torch.manual_seed(0) + conv = ConvTP(in_irreps=f"{F}x0e", out_irreps=TARGET, sh_irreps=SH).double() + checker.tensor( + "ConvTP", + conv, + conv( + node_features=lin(N, F), + edge_angular=edge_attrs, + edge_index=edge_index, + tp_weights=lin(E, conv.weight_numel), + ), + ) + + torch.manual_seed(0) + blk = InteractionBlock( + num_features=F, num_bessel=NUM_RADIAL, l_max=L_MAX, avg_num_neighbors=2.0 + ).double() + checker.tensor( + "InteractionBlock", + blk, + blk( + node_feats=node_feats, + edge_attrs=edge_attrs, + edge_feats=edge_feats, + edge_index=edge_index, + ), + ) + + common = dict( + node_attrs_irreps="3x0e", + edge_attrs_irreps=SH, + edge_feats_irreps=f"{NUM_RADIAL}x0e", + target_irreps=TARGET, + radial_mlp=[8], + ) + + torch.manual_seed(0) + di = DensityInteraction(node_feats_irreps=f"{F}x0e", edge_irreps=f"{F}x0e", **common).double() + checker.tensor( + "DensityInteraction", + di, + di(node_attrs, node_feats, edge_attrs, edge_feats, edge_index), + ) + + torch.manual_seed(0) + dri = DensityResidualInteraction( + node_feats_irreps=f"{F}x0e", edge_irreps=f"{F}x0e", hidden_irreps=f"{F}x0e", **common + ).double() + checker.tensor( + "DensityResidualInteraction", + dri, + dri(node_attrs, node_feats, edge_attrs, edge_feats, edge_index), + ) + + torch.manual_seed(0) + lr = LinearReadout(irreps_in=f"{F}x0e").double() + checker.tensor("LinearReadout", lr, lr(node_feats)) + + torch.manual_seed(0) + nlr = NonLinearReadout(irreps_in=f"{F}x0e", mlp_dim=4).double() + checker.tensor("NonLinearReadout", nlr, nlr(node_feats)) + + torch.manual_seed(0) + nlbr = NonLinearBiasReadout(irreps_in=f"{F}x0e", mlp_dim=4).double() + checker.tensor("NonLinearBiasReadout", nlbr, nlbr(node_feats)) + + torch.manual_seed(0) + ph = ProductHead( + hidden_dim=HIDDEN_DIM, + out_dim=F, + num_radial=NUM_RADIAL, + l_max=L_MAX, + max_body_order=2, + num_species=9, + ).double() + checker.tensor("ProductHead", ph, ph(node_features=lin(N, HIDDEN_DIM), atom_types=Z)) + + torch.manual_seed(0) + emb = EmbeddingBlock( + node_attr_specs=[DiscreteEmbeddingSpec(input_key="Z", num_classes=119, emb_dim=F)], + num_features=F, + r_max=5.0, + num_bessel=NUM_RADIAL, + l_max=L_MAX, + ).double() + nf, ea, ef = emb(Z=Z, edge_dist=edge_dist, edge_diff=edge_diff) + checker.check( + "EmbeddingBlock", + emb, + ( + float(nf.sum()) + float(ea.sum()) + float(ef.sum()), + max(float(nf.abs().max()), float(ea.abs().max()), float(ef.abs().max())), + ), + ) + + if checker.failures: + print("\nFAILED — the relocation is not numerics-neutral:") + for failure in checker.failures: + print(f" {failure}") + return 1 + print("OK") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/regressions/mace-subpackage-restructure-02-core.py b/regressions/mace-subpackage-restructure-02-core.py new file mode 100644 index 0000000..cd6b9c5 --- /dev/null +++ b/regressions/mace-subpackage-restructure-02-core.py @@ -0,0 +1,374 @@ +"""Public-API parity scenario for chain step mace-subpackage-restructure-02. + +Builds both MACE foundation variants through the **new** core layer only — +``MACEMatpesSpec`` / ``MACEOMolSpec`` → ``MACEEncoder`` → +``molzoo.mace.geometry`` → the encoder's primitives — and composes the total +energy in this script, the way a caller has to (the library exposes primitives; +there is no ``compute_all`` façade). Every number is then checked against +literals captured from the flat ``MACEMatpes`` / ``MACEOMol`` models before the +restructure. Any deviation means the spec-driven backbone is not the same model +the flat variants were. + +Goldens +------- +All literals in ``MATPES`` / ``OMOL`` below were captured from the **flat** +``molzoo.mace_matpes.MACEMatpes`` and ``molzoo.mace_omol.MACEOMol`` — the +in-repo self-oracle, deleted in chain step 07, hence hard-coded here and never +imported at run time. The capture script re-implemented each flat model's +``_compute_energy`` line by line to expose the intermediate scalars, and +asserted the recomposed total against the flat model's own public +``energy_forces(..., compute_forces=False)["energy"]`` (bit-equal) before +emitting the table. No third-party oracle, no network, no RNG-derived weights. + +To re-capture while the flat variants still exist: copy this file, construct +``MACEMatpes`` / ``MACEOMol`` with the same kwargs (passing ``atomic_energies`` +as a **float64** tensor — a float32 one shifts ``E0`` by ~1e-6 eV), apply +:func:`deterministic_weights`, inline each flat ``_compute_energy`` in place of +the ``MACEEncoder`` composition below, and print the scalars instead of +asserting. + + capture script : ad-hoc, per the recipe above (flat-variant decomposition) + capture command: PYTHONPATH=src python capture_02_goldens.py + commit : 1ddd5ff (1ddd5ffe9772d490cbd48fd352d59456bc5b5683) + torch : 2.12.1+cpu + date : 2026-08-08 + device / dtype : CPU, float64 (``config.set_precision("fp64")``) + oracle : molzoo.mace_matpes.MACEMatpes / molzoo.mace_omol.MACEOMol + at 1ddd5ff — in-repo self-oracle, no third party + observed : new-vs-flat deviation 0.0 (bit-identical) on every entry, + and ``sorted(state_dict())`` equal for both variants + + goldens re-captured 2026-08-09 at e8d6595 + working-tree dtype/init fixes: + ``_ScalarO3Linear`` N(0,1) init + ``config.ftype`` at construction (see + commit message); previous values captured at 1ddd5ff (2026-08-08), which + reproduces them bit-for-bit. + + Only :data:`OMOL` moved, and not because of the init change — every + parameter here is overwritten by :func:`deterministic_weights`, so this + scenario draws no RNG at all. Three OMOL parameters + (``joint_embedding.embedders.total_charge.weight``, + ``joint_embedding.embedders.total_spin.weight``, + ``joint_embedding.project.0.weight``) were built at the torch default fp32 + before the fix: the fp64 ``linspace`` ramp was rounded into them on the way + in, and the conditioning was then contracted in fp32. Both now stay fp64, + which is why the shifts sit at the fp32 epsilon scale (≤4.7e-8 relative on + the per-layer checksums, ~9.8e-10 eV on the energy) and why the new numbers + are the strictly higher-precision ones. :data:`MATPES` is bit-identical + across the fix — its only fp32 leak was the ``cutoff_fn.r_cut`` buffer, and + ``r_max=5.0`` is exact in both precisions. + +Why the energy total is not the only assertion +---------------------------------------------- +Under the ``linspace(-0.1, 0.1)`` weight fill the MatPES interaction stack is +numerically tiny (node features ~5e-10, readout energies ~7e-12 eV), so its +total energy is dominated by ``E0`` and the ZBL term: a 1e-9 eV check on the +total alone would pass even with the whole interaction stack zeroed. The +per-layer feature sums-of-squares and per-layer readout energies are therefore +asserted **relatively** (1e-10) — those are the load-bearing assertions on the +interaction/product stack. They are ordinary products of small numbers, not +cancellations, so their relative accuracy is full double precision (verified +stable across ``OMP_NUM_THREADS`` 1/2/4 at capture time). + +Run: + PYTHONPATH=src python regressions/mace-subpackage-restructure-02-core.py +""" + +from __future__ import annotations + +import sys +from typing import NamedTuple + +import torch + +from molix import config + +config.set_precision("fp64") # before any module construction: cuEq bakes in dtype + +from molix.F.scatter import scatter_sum_compile_safe as scatter_sum # noqa: E402 +from molzoo.mace.encoder import MACEEncoder # noqa: E402 +from molzoo.mace.geometry import edge_lengths, edge_vectors # noqa: E402 +from molzoo.mace.spec import MACEMatpesSpec, MACEOMolSpec # noqa: E402 + +#: Absolute tolerance on a total energy, in eV. +ENERGY_ATOL = 1e-9 +#: Relative tolerance on the dimensionless per-layer checksums. +CHECKSUM_RTOL = 1e-10 + +# --------------------------------------------------------------------------- +# System: a fixed 5-atom cluster, all 20 ordered pairs as directed edges. +# Literal coordinates in Å — no RNG, so the goldens are reproducible anywhere. +# --------------------------------------------------------------------------- +ATOMIC_NUMBERS: list[int] = [1, 6, 8] +ATOMIC_ENERGIES: list[float] = [-13.6, -1029.0, -2041.0] # eV/atom + +POSITIONS = torch.tensor( + [ + [0.00, 0.00, 0.00], + [1.09, 0.00, 0.00], + [1.70, 1.15, 0.00], + [-0.40, 0.95, 0.30], + [2.10, -0.85, -0.50], + ], + dtype=torch.float64, +) +Z = torch.tensor([1, 6, 8, 1, 1]) +BATCH = torch.zeros(5, dtype=torch.long) +NUM_GRAPHS = 1 +EDGE_INDEX = torch.tensor( # (E, 2): [:, 0] = source, [:, 1] = target + [(i, j) for i in range(5) for j in range(5) if i != j], dtype=torch.long +) + +#: Tiny MatPES hyper-parameters (l_max=1, 16 channels) — CPU-fast, structurally +#: identical to the shipped model. +MATPES_KWARGS = dict( + r_max=5.0, + num_bessel=4, + num_polynomial_cutoff=5, + l_max=1, + num_features=16, + max_hidden_l=1, + num_interactions=2, + correlation=2, + mlp_dim=8, + radial_mlp=[8], + use_fallback=True, # CPU: the fused kernels need a GPU + the ops wheel +) +#: Tiny OMOL hyper-parameters. The flat ``MACEOMol`` had no ``use_fallback`` +#: argument and hard-coded the fused path; on CPU without +#: ``cuequivariance-ops-torch`` that degrades to the same naive contraction, so +#: ``use_fallback=True`` here reproduces the flat oracle bit-for-bit (verified +#: at capture time against both settings). +OMOL_KWARGS = dict( + r_max=5.0, + num_bessel=4, + num_polynomial_cutoff=5, + l_max=1, + num_features=16, + num_interactions=2, + correlation=2, + mlp_dim=8, + edge_channels=8, + use_fallback=True, +) + +# --------------------------------------------------------------------------- +# Goldens — flat-variant capture, see the module docstring. +# --------------------------------------------------------------------------- + + +class MatpesGolden(NamedTuple): + """Flat ``MACEMatpes`` reference scalars for the system above.""" + + energy: float # eV, total + e0: float # eV, Σ E0[Z] + zbl: float # eV, Σ_i ZBL_i + sumsq: tuple[float, ...] # per-layer Σ h², dimensionless + readouts: tuple[float, ...] # eV, per-layer Σ_i readout_l(h_l)_i + + +class OMolGolden(NamedTuple): + """Flat ``MACEOMol`` reference scalars for the system above.""" + + energy: float # eV, total + e0: float # eV, Σ E0[Z] + emb: float # eV, charge/spin embedding readout + sumsq: tuple[float, ...] # per-layer Σ h², dimensionless + readout: float # eV, Σ_i readout(h_last)_i (pre scale_shift) + + +MATPES = MatpesGolden( + energy=-3110.794351427106, + e0=-3110.7999999999997, + zbl=0.0056485728867009056, + sumsq=(1.6644954580828804e-17, 1.4248167631489696e-21), + readouts=(6.852674614296589e-12, 3.8986176107825314e-14), +) +OMOL = OMolGolden( + energy=-3111.2905857505443, + e0=-3110.7999999999997, + emb=-0.04147213575986683, + sumsq=(0.00489757656783223, 2.020471377543396e-06), + readout=-0.4491136147850327, +) + + +def deterministic_weights(model: torch.nn.Module) -> None: + """Overwrite every parameter with a fixed ramp — no RNG anywhere. + + Applied identically to the flat oracle at capture time and to + :class:`~molzoo.mace.encoder.MACEEncoder` here; ``state_dict`` key parity is + what makes "identically" well defined. Buffers are left alone (they carry + the element table and the cuEquivariance graph constants). + + Args: + model: Module whose parameters are replaced in place. + """ + with torch.no_grad(): + for _, parameter in sorted(model.named_parameters()): + parameter.data.copy_( + torch.linspace(-0.1, 0.1, parameter.numel(), dtype=torch.float64).reshape( + parameter.shape + ) + ) + + +class Checker: + """Collects deviations of observed scalars from the embedded goldens.""" + + def __init__(self) -> None: + self.failures: list[str] = [] + + def absolute(self, name: str, observed: float, golden: float, atol: float) -> None: + """Assert ``|observed - golden| <= atol`` (energies, in eV).""" + deviation = abs(observed - golden) + if deviation > atol: + self.failures.append( + f"{name}: got {observed!r}, want {golden!r} " + f"(deviation {deviation:.3e} > atol {atol:.1e})" + ) + print(f" {name:<26} deviation {deviation:.3e}") + + def relative(self, name: str, observed: float, golden: float, rtol: float) -> None: + """Assert ``|observed - golden| <= rtol * |golden|`` (checksums).""" + deviation = abs(observed - golden) + if deviation > rtol * abs(golden): + self.failures.append( + f"{name}: got {observed!r}, want {golden!r} " + f"(relative deviation {deviation / abs(golden):.3e} > rtol {rtol:.1e})" + ) + print(f" {name:<26} relative deviation {deviation / abs(golden):.3e}") + + +def check_matpes(checker: Checker) -> None: + """Compose the MatPES total energy from primitives and verify it. + + Mirrors ``MACEMatpes._compute_energy``: + ``E0 + scale_shift(ZBL + Σ_layers readout_l(h_l))``, scattered per graph. + """ + print("MACEMatpes (density interactions, Agnesi transform, ZBL, per-layer readouts)") + spec = MACEMatpesSpec( + atomic_numbers=ATOMIC_NUMBERS, atomic_energies=ATOMIC_ENERGIES, **MATPES_KWARGS + ) + torch.manual_seed(0) # module construction may draw; the fill below overwrites it + encoder = MACEEncoder(spec).eval() + deterministic_weights(encoder) + encoder.validate_elements(Z) + + vectors = edge_vectors(POSITIONS, EDGE_INDEX) + lengths = edge_lengths(vectors) + node_attrs = encoder.node_attrs(Z, POSITIONS.dtype) + edge_attrs = encoder.angular_features(vectors) + edge_feats, cutoff = encoder.radial_features(lengths, Z, EDGE_INDEX) + + e0 = scatter_sum(encoder.atomic_energies(Z), BATCH, NUM_GRAPHS) + node_energies = encoder.pair_repulsion(lengths, Z, EDGE_INDEX) + checker.absolute("matpes.e0", float(e0[0]), MATPES.e0, ENERGY_ATOL) + checker.absolute("matpes.zbl", float(node_energies.sum()), MATPES.zbl, ENERGY_ATOL) + + layers = encoder.layer_features( + node_feats=encoder.initial_node_features(node_attrs), + node_attrs=node_attrs, + edge_attrs=edge_attrs, + edge_feats=edge_feats, + edge_index=EDGE_INDEX, + cutoff=cutoff, + ) + if len(layers) != len(MATPES.sumsq): + checker.failures.append( + f"matpes.layer_features: got {len(layers)} layers, want {len(MATPES.sumsq)}" + ) + return + for i, features in enumerate(layers): + checker.relative( + f"matpes.sumsq[{i}]", float((features**2).sum()), MATPES.sumsq[i], CHECKSUM_RTOL + ) + head = encoder.readouts[i](features).squeeze(-1) + checker.relative( + f"matpes.readout[{i}]", float(head.sum()), MATPES.readouts[i], CHECKSUM_RTOL + ) + node_energies = node_energies + head + + energy = e0 + scatter_sum(encoder.scale_shift(node_energies), BATCH, NUM_GRAPHS) + checker.absolute("matpes.energy", float(energy[0]), MATPES.energy, ENERGY_ATOL) + + +def check_omol(checker: Checker) -> None: + """Compose the OMOL total energy from primitives and verify it. + + Mirrors ``MACEOMol._compute_energy``: ``E0 + Σ embedding_readout(h_0 + + conditioning) + scale_shift(readout(h_last))``, scattered per graph. + """ + print("\nMACEOMol (residual interactions, charge/spin conditioning, final readout)") + spec = MACEOMolSpec( + atomic_numbers=ATOMIC_NUMBERS, atomic_energies=ATOMIC_ENERGIES, **OMOL_KWARGS + ) + torch.manual_seed(0) + encoder = MACEEncoder(spec).eval() + deterministic_weights(encoder) + encoder.validate_elements(Z) + + total_charge = torch.zeros(NUM_GRAPHS, dtype=torch.long) + total_spin = torch.ones(NUM_GRAPHS, dtype=torch.long) # OMOL: 1 = closed-shell singlet + + vectors = edge_vectors(POSITIONS, EDGE_INDEX) + lengths = edge_lengths(vectors, keepdim=True) # OMOL's own (E, 1) shape + node_attrs = encoder.node_attrs(Z, POSITIONS.dtype) + edge_attrs = encoder.angular_features(vectors) + edge_feats, cutoff = encoder.radial_features(lengths.squeeze(-1), Z, EDGE_INDEX) + + e0 = scatter_sum(encoder.atomic_energies(Z), BATCH, NUM_GRAPHS) + checker.absolute("omol.e0", float(e0[0]), OMOL.e0, ENERGY_ATOL) + + node_feats = encoder.initial_node_features(node_attrs) + encoder.conditioning( + BATCH, total_spin=total_spin, total_charge=total_charge + ) + embedding_energy = scatter_sum( + encoder.embedding_readout(node_feats).squeeze(-1), BATCH, NUM_GRAPHS + ) + checker.absolute("omol.emb", float(embedding_energy[0]), OMOL.emb, ENERGY_ATOL) + + layers = encoder.layer_features( + node_feats=node_feats, + node_attrs=node_attrs, + edge_attrs=edge_attrs, + edge_feats=edge_feats, + edge_index=EDGE_INDEX, + cutoff=cutoff, + ) + if len(layers) != len(OMOL.sumsq): + checker.failures.append( + f"omol.layer_features: got {len(layers)} layers, want {len(OMOL.sumsq)}" + ) + return + for i, features in enumerate(layers): + checker.relative( + f"omol.sumsq[{i}]", float((features**2).sum()), OMOL.sumsq[i], CHECKSUM_RTOL + ) + + node_energies = encoder.readout(layers[-1]).squeeze(-1) + checker.relative("omol.readout", float(node_energies.sum()), OMOL.readout, CHECKSUM_RTOL) + + energy = ( + e0 + embedding_energy + scatter_sum(encoder.scale_shift(node_energies), BATCH, NUM_GRAPHS) + ) + checker.absolute("omol.energy", float(energy[0]), OMOL.energy, ENERGY_ATOL) + + +def main() -> int: + """Rebuild both variants on the new core layer and verify every golden.""" + checker = Checker() + with torch.no_grad(): + check_matpes(checker) + check_omol(checker) + + if checker.failures: + print("\nFAILED — the spec-driven MACEEncoder is not the flat variants' model:") + for failure in checker.failures: + print(f" {failure}") + return 1 + print("\nOK") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/regressions/mace-subpackage-restructure-03-kernels.py b/regressions/mace-subpackage-restructure-03-kernels.py new file mode 100644 index 0000000..7a87122 --- /dev/null +++ b/regressions/mace-subpackage-restructure-03-kernels.py @@ -0,0 +1,218 @@ +"""Public-API force-pass parity scenario for chain step mace-subpackage-restructure-03. + +Chain step 03 extracts the two batch-level force passes duplicated in +``molzoo/pinet/potential.py`` and ``molpot/derivation/modes/`` into the shared +``molpot.derivation.kernels`` (``grad_force_pass`` / ``func_force_pass``). The +spec declares that refactor **zero-behaviour-change**, so this script pins the +only thing that can prove it end to end: a fixed-seed ``PiNetPotential`` must +keep producing bit-identical energies and forces through both public force +methods after the rebinding. + +Scenario (public API only): construct ``PiNetPotential(..., compute_forces=True, +method="func")`` and the same model with ``method="grad"``, feed one fixed +4-atom / 8-edge batch, and compare ``graphs.energy`` (eV) and ``atoms.forces`` +(eV/Å) against the hard-coded literals below at ``atol=1e-12, rtol=0``. A third +``compute_forces=False`` instance (bit-identical ``state_dict``) supplies the +energies for a central-difference check of the forces — the domain leg, since +"the kernels still return ``F = -∂E/∂r``" is a physics claim, not a diff claim. + +Goldens +------- +``GOLDEN_ENERGY`` / ``GOLDEN_FORCES`` were captured from the **pre-rebinding** +``PiNetPotential`` (its own local ``_pipeline_ef_func`` / ``_pipeline_ef_grad`` +copies) — an in-repo self-oracle at the commit below, never re-imported at run +time. ``method="func"`` and ``method="grad"`` produced bit-identical values at +capture, so one table covers both and the script additionally asserts the two +methods agree exactly. + + capture script : scratchpad ``capture_03_goldens.py`` — same construction, + seed and batch as this file, printing ``.tolist()`` + capture command: PYTHONPATH=src:. python capture_03_goldens.py + commit : 03c0e85 (03c0e85b91f64cc09a1576139eb5576fe7cc4ab2) + torch : 2.12.1+cpu + date : 2026-08-08 + device / dtype : CPU, float64 (``config.set_precision("fp64")``) + oracle : molzoo.pinet.PiNetPotential at 03c0e85 — in-repo + self-oracle; no ASE / e3nn / mace-torch, no network, + no subprocess + observed : func-vs-grad deviation 0.0 (bit-identical); autograd-vs + -central-difference deviation 1.3e-11 eV/Å at h = 1e-5 Å + +The batch is built inline (the same ``edge_diff = pos[dst] - pos[src]``, +``edge_dist = ‖edge_diff‖`` post-collate schema that ``tests/conftest.py`` +produces) so the script stays standalone — ``PYTHONPATH=src`` is enough, the +repo root is not needed on the path. + +Run: + PYTHONPATH=src python regressions/mace-subpackage-restructure-03-kernels.py +""" + +from __future__ import annotations + +import sys +from typing import Literal + +import torch + +from molix import config + +config.set_precision("fp64") # before any module construction + +from tensordict import TensorDict # noqa: E402 + +from molzoo.pinet import PiNetPotential # noqa: E402 + +# --------------------------------------------------------------------------- # +# Fixed system: 4 atoms (H, C, N, O), 8 directed edges, one graph. +# --------------------------------------------------------------------------- # +POS = torch.tensor( + [ + [0.0, 0.0, 0.0], + [1.1, 0.1, 0.0], + [0.3, 1.2, 0.2], + [1.4, 1.1, -0.1], + ], + dtype=torch.float64, +) +Z = torch.tensor([1, 6, 7, 8], dtype=torch.long) +EDGE_INDEX = torch.tensor( + [[0, 1], [1, 0], [0, 2], [2, 0], [1, 3], [3, 1], [2, 3], [3, 2]], + dtype=torch.long, +) +BATCH = torch.zeros(4, dtype=torch.long) + +# --------------------------------------------------------------------------- # +# Hard-coded goldens (see module docstring for provenance). +# --------------------------------------------------------------------------- # +GOLDEN_ENERGY: list[float] = [0.8332936477597297] # eV +GOLDEN_FORCES: list[list[float]] = [ # eV/Å + [-0.005663236275122078, -0.005354983795618726, -0.0003581090284556928], + [0.0035887033693461586, -0.004020850505374732, -0.0002506004328096707], + [-0.003822775289657374, 0.005858775873778405, 0.0017552896792324322], + [0.005897308195433293, 0.0035170584272150533, -0.0011465802179670686], +] + +ATOL = 1e-12 # exact band: the refactor must not move a single bit +FD_ATOL = 1e-6 # numerical band for the central-difference domain check +FD_STEP = 1e-5 # Å + + +def make_batch(pos: torch.Tensor) -> TensorDict: + """Post-collate batch for ``pos`` ``(4, 3)`` Å (atoms / edges / graphs).""" + edge_diff = pos[EDGE_INDEX[:, 1]] - pos[EDGE_INDEX[:, 0]] + edge_dist = edge_diff.norm(dim=-1).clamp(min=1e-6) + num_atoms = torch.tensor([pos.shape[0]], dtype=torch.long) + return TensorDict( + atoms=TensorDict(Z=Z, pos=pos, batch=BATCH, batch_size=[pos.shape[0]]), + edges=TensorDict( + edge_index=EDGE_INDEX, + edge_diff=edge_diff, + edge_dist=edge_dist, + batch_size=[EDGE_INDEX.shape[0]], + ), + graphs=TensorDict(num_atoms=num_atoms, batch_size=[1]), + batch_size=[], + ) + + +def build(method: Literal["func", "grad"], *, compute_forces: bool = True) -> PiNetPotential: + """Seed-0 PiNet potential — identical weights for every ``method``.""" + torch.manual_seed(0) + return PiNetPotential( + atom_types=[1, 6, 7, 8], + r_max=4.0, + n_basis=3, + pp_nodes=[8, 8], + pi_nodes=[8, 8], + ii_nodes=[8, 8], + depth=2, + rank=3, + hidden_dim=16, + compute_forces=compute_forces, + method=method, + ).eval() + + +class Checker: + """Collect deviations so one run reports every failure, not just the first.""" + + def __init__(self) -> None: + self.failures: list[str] = [] + + def close(self, name: str, got: torch.Tensor, want: torch.Tensor, atol: float) -> None: + if got.shape != want.shape: + self.failures.append(f"{name}: shape {tuple(got.shape)} != {tuple(want.shape)}") + return + deviation = (got - want).abs().max().item() + if not deviation <= atol: + self.failures.append(f"{name}: max|Δ| = {deviation:.3e} > atol {atol:.1e}") + + +def check_method(checker: Checker, method: Literal["func", "grad"]) -> TensorDict: + """Run one public force method and compare energy + forces to the goldens.""" + result = build(method)(make_batch(POS.clone())) + checker.close( + f"{method}.energy", + result["graphs", "energy"].detach(), + torch.tensor(GOLDEN_ENERGY, dtype=torch.float64), + ATOL, + ) + checker.close( + f"{method}.forces", + result["atoms", "forces"].detach(), + torch.tensor(GOLDEN_FORCES, dtype=torch.float64), + ATOL, + ) + return result + + +def check_finite_differences(checker: Checker, forces: torch.Tensor) -> None: + """Domain leg: ``F = -dE/dr`` by central differences at ``h = 1e-5`` Å.""" + energy_model = build("func", compute_forces=False) + numerical = torch.zeros_like(POS) + with torch.no_grad(): + for atom in range(POS.shape[0]): + for axis in range(POS.shape[1]): + plus = POS.clone() + plus[atom, axis] += FD_STEP + minus = POS.clone() + minus[atom, axis] -= FD_STEP + e_plus = energy_model(make_batch(plus))["graphs", "energy"].sum() + e_minus = energy_model(make_batch(minus))["graphs", "energy"].sum() + numerical[atom, axis] = -(e_plus - e_minus) / (2.0 * FD_STEP) + checker.close("finite_difference.forces", forces, numerical, FD_ATOL) + + +def main() -> int: + """Verify both force methods against the goldens and against the physics.""" + checker = Checker() + + func_result = check_method(checker, "func") + grad_result = check_method(checker, "grad") + + checker.close( + "func_vs_grad.energy", + func_result["graphs", "energy"].detach(), + grad_result["graphs", "energy"].detach(), + 0.0, + ) + checker.close( + "func_vs_grad.forces", + func_result["atoms", "forces"].detach(), + grad_result["atoms", "forces"].detach(), + 0.0, + ) + + check_finite_differences(checker, func_result["atoms", "forces"].detach()) + + if checker.failures: + print("FAILED — PiNet force passes moved after the kernel rebinding:") + for failure in checker.failures: + print(f" {failure}") + return 1 + print("OK") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/regressions/mace-subpackage-restructure-04-potential.py b/regressions/mace-subpackage-restructure-04-potential.py new file mode 100644 index 0000000..25eb94d --- /dev/null +++ b/regressions/mace-subpackage-restructure-04-potential.py @@ -0,0 +1,523 @@ +"""Public-API energy/force scenario for chain step mace-subpackage-restructure-04. + +Chain step 04 merges the two flat foundation models +(``molzoo.mace_matpes.MACEMatpes`` / ``molzoo.mace_omol.MACEOMol``) into the one +spec-driven :class:`molzoo.mace.MACEPotential`. Step 07 then **deletes** the flat +modules, at which point this file becomes the only frozen memory of what those +models computed — so it must land before 07 and it must never import them. + +Scenario (public API only): ``MACEMatpesSpec`` / ``MACEOMolSpec`` → +``MACEPotential`` → the four public seams a caller actually has, on one fixed +5-atom cluster at fp64 on CPU: + +1. ``potential(batch)`` — energy ``(B,)`` eV + forces ``(N, 3)`` eV/Å; +2. ``potential(batch)`` with ``edges.shifts`` — the periodic / MD path; +3. ``potential.energy_core(...)`` — the flat, compilable seam (raw tensors in, + ``(B,)`` out), including the OMOL ``total_charge`` / ``total_spin`` + conditioning and the MatPES refusal of it; +4. ``MACEPotential(spec, compute_forces=False)`` and + ``molpot.derivation.protocol.call_energy`` — the two energy-only entries; + "energy only" is a *construction*, never a per-call flag. + +Physics leg: ``Σ_i F_i = 0`` on an isolated cluster (the energy depends on the +positions only through ``r_ij``), asserted for both variants. + +Goldens +------- +Every literal in :data:`MATPES` / :data:`OMOL` was captured by running this +file's own builders at the commit below and printing ``.tolist()``. At capture +time each number was **also** cross-checked against the flat model it replaces: +a fresh ``MACEMatpes`` / ``MACEOMol`` built with the same hyper-parameters, +handed the potential's ``state_dict()`` under ``load_state_dict(strict=True)`` +("All keys matched successfully" for both), then called through its public +``energy_forces(...)``. The flat modules are deliberately **not** imported here +— after 07 removes them this script must still run unchanged — so that parity +lives in this comment only: + + capture script : scratchpad ``capture_04_goldens.py`` — this file's + builders plus the flat-model cross-check described above + capture command : PYTHONPATH=src python capture_04_goldens.py + commit : c2ccd6e (c2ccd6e57b4d05ab7e177154e0e8bb63831801d0) + torch : 2.12.1+cpu + date : 2026-08-08 + device / dtype : CPU, float64 (``config.set_precision("fp64")``) + oracle : molzoo.mace_matpes.MACEMatpes / molzoo.mace_omol.MACEOMol + at c2ccd6e — in-repo self-oracle. No ASE / e3nn / + mace-torch, no network, no subprocess, no RNG beyond the + two fixed seeds below. + observed : energies bit-identical to the flat models (deviation + exactly 0.0) for MatPES free, MatPES periodic and OMOL; + OMOL forces bit-identical; MatPES forces agree to + 7.3e-17 eV/Å (free) and 1.2e-16 eV/Å (periodic) — the + backward runs over an equal-valued but not + operation-identical graph, ~2 ulp at these magnitudes, + far inside the 1e-12 band the spec asks for. + + goldens re-captured 2026-08-09 at e8d6595 + working-tree dtype/init fixes: + ``_ScalarO3Linear`` N(0,1) init + ``config.ftype`` at construction (see + commit message); previous values captured at c2ccd6e (2026-08-08), which + reproduces them bit-for-bit. + + Only :data:`OMOL` and :data:`OMOL_CATION_ENERGY` moved. Both variants are + seeded models, so the shift is an init-stream shift, not a physics change: + OMOL's ``joint_embedding`` conditioning tables were built at the torch + default fp32 and are now built at ``config.ftype``, and drawing the same + number of elements at fp64 consumes the global RNG differently — every + parameter constructed after the embedding therefore differs. Decomposed, + ``Σ E0[Z]`` is bit-identical (-3110.7999999999997), the embedding readout + stays O(1) eV (-1.084109101152535 → -1.454179831707279) and the readout + energy moves with the new draw (-31.130204741718686 → 1.102275717507076). + + Unmoved, and re-verified: every MatPES golden (free and periodic energies, + the full force table and its checksum) is bit-identical — MatPES's only + fp32 leak was the ``cutoff_fn.r_cut`` buffer, and ``r_max=5.0`` is exact in + both precisions. The seam-equality legs still hold at ``SEAM_ATOL = 0.0``, + and ``Σ_i F_i`` stays at 1.0e-17 (MatPES) / 3.5e-18 (OMOL) eV/Å. + +Determinism: ``torch.manual_seed(0)`` immediately before each construction +(module init is the only RNG consumer), plus a private +``torch.Generator().manual_seed(0)`` for the OMOL readout — see +:func:`wake_readout`. Two consecutive runs were byte-identical at +capture time. + +Run: + PYTHONPATH=src python regressions/mace-subpackage-restructure-04-potential.py +""" + +from __future__ import annotations + +import sys +from typing import NamedTuple + +import torch + +from molix import config + +config.set_precision("fp64") # before any module construction: cuEq bakes in dtype + +from tensordict import TensorDict # noqa: E402 + +from molpot.derivation.protocol import call_energy # noqa: E402 +from molzoo.mace import MACEMatpesSpec, MACEOMolSpec, MACEPotential # noqa: E402 + +#: Golden band on energies (eV) and forces (eV/Å). The spec's numerical +#: invariant for the move is ``max|Δ| ≤ 1e-12``; measured deviations at capture +#: were 0.0 (energies) and ≤ 1.2e-16 (forces). +GOLDEN_ATOL = 1e-12 + +#: Cross-seam band: ``energy_core`` / ``forward`` / ``compute_forces=False`` / +#: ``call_energy`` must return the *same* energy, not a close one. Measured +#: bit-identical at capture, so the band is exactly zero. +SEAM_ATOL = 0.0 + +#: Newton's third law on an isolated cluster (measured ≤ 3.2e-17 eV/Å). +NET_FORCE_ATOL = 1e-10 + +# --------------------------------------------------------------------------- +# System: a fixed 5-atom cluster, all 20 ordered pairs as directed edges. +# Literal coordinates in Å — no RNG, so the geometry is reproducible anywhere. +# Same cluster as regressions/mace-subpackage-restructure-02-core.py. +# --------------------------------------------------------------------------- +ATOMIC_NUMBERS: list[int] = [1, 6, 8] +ATOMIC_ENERGIES: list[float] = [-13.6, -1029.0, -2041.0] # eV/atom + +POSITIONS = torch.tensor( + [ + [0.00, 0.00, 0.00], + [1.09, 0.00, 0.00], + [1.70, 1.15, 0.00], + [-0.40, 0.95, 0.30], + [2.10, -0.85, -0.50], + ], + dtype=torch.float64, +) +Z = torch.tensor([1, 6, 8, 1, 1], dtype=torch.long) +BATCH = torch.zeros(5, dtype=torch.long) +NUM_GRAPHS = 1 +EDGE_INDEX = torch.tensor( # (E, 2): [:, 0] = source, [:, 1] = target + [(i, j) for i in range(5) for j in range(5) if i != j], dtype=torch.long +) + +#: ``unit_shifts @ cell`` on the first two (mutually reverse) edges, in Å: the +#: minimal periodic image that still exercises the ``edges.shifts`` seam. +SHIFTS = torch.zeros(EDGE_INDEX.shape[0], 3, dtype=torch.float64) +SHIFTS[0, 0] = 2.0 +SHIFTS[1, 0] = -2.0 + +#: Neutral closed-shell singlet — the OMOL defaults. Spin ``0`` would index an +#: untrained embedding row, so the conditioning is always spelled out. +TOTAL_CHARGE = torch.zeros(NUM_GRAPHS, dtype=torch.long) +TOTAL_SPIN = torch.ones(NUM_GRAPHS, dtype=torch.long) +CATION_CHARGE = torch.ones(NUM_GRAPHS, dtype=torch.long) + +#: Tiny MatPES hyper-parameters (``l_max=1``, 16 channels) — CPU-fast and +#: structurally identical to the shipped model. +MATPES_KWARGS = dict( + r_max=5.0, + num_bessel=4, + num_polynomial_cutoff=5, + l_max=1, + num_features=16, + max_hidden_l=1, + num_interactions=2, + correlation=2, + mlp_dim=8, + radial_mlp=[8], +) +#: Tiny OMOL hyper-parameters. ``use_fallback`` is passed to the *constructor* +#: rather than the spec (the flat ``MACEOMol`` hard-coded the fused path; on CPU +#: without ``cuequivariance-ops-torch`` that degrades to the same naive +#: contraction, so the two agree bit-for-bit). +OMOL_KWARGS = dict( + r_max=5.0, + num_bessel=4, + num_polynomial_cutoff=5, + l_max=1, + num_features=16, + num_interactions=2, + correlation=2, + mlp_dim=8, + edge_channels=8, +) + +#: ``molrep.readout.mace._ScalarO3Linear`` draws its weight from ``N(0, 1)`` on +#: the **global** RNG and zero-initialises its bias. These four ``state_dict`` +#: entries are refilled from a *private* generator (see +#: :func:`wake_readout`) so the OMOL goldens below depend only on that +#: generator, not on how much global RNG the rest of the model happens to +#: consume ahead of the readout. +#: Fill order is load-bearing — the private generator is consumed entry by +#: entry, so reordering this tuple silently re-rolls every OMOL golden. +READOUT_WAKE_KEYS = ( + "readout.linear_mid.weight", + "readout.linear_mid.bias", + "readout.linear_2.weight", + "readout.linear_2.bias", +) + + +# --------------------------------------------------------------------------- +# Goldens — see the module docstring for provenance. +# --------------------------------------------------------------------------- + + +class Golden(NamedTuple): + """One variant's captured energy/force signature for this cluster.""" + + energy: list[float] # eV, (B,) + forces: list[list[float]] # eV/Å, (N, 3) + force_abs_sum: float # eV/Å, Σ|F| — a single-number checksum + + +MATPES = Golden( + energy=[-3110.9265191001523], + forces=[ + [0.019021511971043834, 0.06740306449657524, 0.0042655930165905525], + [-0.03884100764534142, -0.14330968688519968, 0.03275558451993726], + [-0.025188575897160226, 0.0664617371393798, -0.004729317224920349], + [-0.005148555478052416, 0.02124111057338372, -0.005410511456735142], + [0.05015662704951024, -0.011796225324139067, -0.02688134885487232], + ], + force_abs_sum=0.5226104575328413, +) + +#: The same MatPES potential with ``edges.shifts`` on the first two edges. Only +#: the total energy and the force checksum are pinned: the point of this leg is +#: that the shift *reaches* the edge vectors (it moves the energy by 1.0e-2 eV), +#: not a second full force table. +MATPES_PERIODIC_ENERGY: list[float] = [-3110.936741536619] +MATPES_PERIODIC_FORCE_ABS_SUM = 0.6035586909434021 + +OMOL = Golden( + energy=[-3111.1519041142], + forces=[ + [0.0016186653588899126, 0.02170414821495261, 0.006063378516186664], + [0.0010368535618123821, 0.005192545579351304, -0.0021008220476804373], + [-0.014156038043983701, -0.017889811992860463, -0.0005845605451978166], + [0.019110658287980815, -0.020804622960958453, -0.00796628298659983], + [-0.007610139164699409, 0.011797741159515, 0.004588287063291421], + ], + force_abs_sum=0.1422245554839602, +) + +#: The same OMOL potential and geometry at total charge +1 (spin still 1): the +#: conditioning must reach the joint embedding, and it moves the energy by +#: 0.47 eV. +OMOL_CATION_ENERGY: list[float] = [-3110.6811189024315] + + +# --------------------------------------------------------------------------- +# Construction +# --------------------------------------------------------------------------- + + +def wake_readout(potential: MACEPotential) -> None: + """Pin the OMOL readout to a private generator, off the global RNG. + + ``_ScalarO3Linear`` draws its weight from ``N(0, 1)`` on the global RNG and + zero-initialises its bias. Refilling all four scalar-linear entries from a + private, fixed generator keeps the goldens below stable against changes in + how much global RNG the rest of the model consumes before the readout is + built (write-back through the public ``state_dict`` / ``load_state_dict`` + pair, no private attribute touched). Checkpoint use is unaffected: official + weights overwrite these entries. + + Both halves of the init are asserted on the way through, because either one + silently invalidates the goldens: a weight back at zero would make the + untrained readout emit a constant per-atom energy and zero every OMOL force + (the state this fills in for historically), and a non-zero bias would mean + the init changed again. + + Args: + potential: OMOL potential to modify in place. + + Raises: + RuntimeError: If a key is missing (the readout was renamed), a weight is + zero, or a bias is non-zero — in each case the goldens below no + longer describe the model that produced them. + """ + state = potential.state_dict() + generator = torch.Generator().manual_seed(0) + for key in READOUT_WAKE_KEYS: + if key not in state: + raise RuntimeError(f"OMOL readout key {key!r} is gone — re-capture the goldens") + magnitude = float(state[key].abs().max()) + if key.endswith(".weight") and magnitude == 0.0: + raise RuntimeError( + f"OMOL readout weight {key!r} is zero-init again — every force would be zero" + ) + if key.endswith(".bias") and magnitude != 0.0: + raise RuntimeError(f"OMOL readout bias {key!r} is no longer zero-init at construction") + state[key] = torch.empty_like(state[key]).normal_(generator=generator) + potential.load_state_dict(state, strict=True) + + +def matpes_potential(*, compute_forces: bool = True) -> MACEPotential: + """Seed-0 MatPES potential (``use_fallback=True``: CPU has no ops wheel).""" + spec = MACEMatpesSpec( + atomic_numbers=ATOMIC_NUMBERS, atomic_energies=ATOMIC_ENERGIES, **MATPES_KWARGS + ) + torch.manual_seed(0) + return MACEPotential(spec, compute_forces=compute_forces, use_fallback=True).eval() + + +def omol_potential(*, compute_forces: bool = True) -> MACEPotential: + """Seed-0 OMOL potential, with the readout pinned to its private generator.""" + spec = MACEOMolSpec( + atomic_numbers=ATOMIC_NUMBERS, atomic_energies=ATOMIC_ENERGIES, **OMOL_KWARGS + ) + torch.manual_seed(0) + potential = MACEPotential(spec, compute_forces=compute_forces, use_fallback=True).eval() + wake_readout(potential) + return potential + + +def make_batch( + *, + shifts: torch.Tensor | None = None, + graphs: dict[str, torch.Tensor] | None = None, +) -> TensorDict: + """Post-collate batch for :data:`POSITIONS` (atoms / edges / graphs). + + A fresh object every call: ``forward`` writes in place and swaps + ``atoms.pos`` for its own differentiation leaf. + + Args: + shifts: Optional periodic shift vectors ``(E, 3)`` in Å, folded into + ``edge_diff`` and published as ``edges.shifts``. + graphs: Optional per-graph fields (``total_charge`` / ``total_spin``). + + Returns: + A ``TensorDict`` with ``edge_diff = pos[target] - pos[source] (+ S)`` + and ``edge_dist = ‖edge_diff‖``. + """ + pos = POSITIONS.clone() + edge_diff = pos[EDGE_INDEX[:, 1]] - pos[EDGE_INDEX[:, 0]] + if shifts is not None: + edge_diff = edge_diff + shifts + edges = TensorDict( + edge_index=EDGE_INDEX, + edge_diff=edge_diff, + edge_dist=edge_diff.norm(dim=-1).clamp(min=1e-6), + batch_size=[EDGE_INDEX.shape[0]], + ) + if shifts is not None: + edges["shifts"] = shifts + graph_data = TensorDict( + num_atoms=torch.tensor([pos.shape[0]], dtype=torch.long), batch_size=[NUM_GRAPHS] + ) + for key, value in (graphs or {}).items(): + graph_data[key] = value + return TensorDict( + atoms=TensorDict(Z=Z, pos=pos, batch=BATCH, batch_size=[pos.shape[0]]), + edges=edges, + graphs=graph_data, + batch_size=[], + ) + + +# --------------------------------------------------------------------------- +# Checking +# --------------------------------------------------------------------------- + + +class Checker: + """Collects every deviation so one run reports all failures, not the first.""" + + def __init__(self) -> None: + self.failures: list[str] = [] + + def close(self, name: str, got: torch.Tensor, want: list, atol: float) -> None: + """Assert ``max|got - want| <= atol`` against a literal nested list.""" + expected = torch.tensor(want, dtype=torch.float64) + if got.shape != expected.shape: + self.failures.append(f"{name}: shape {tuple(got.shape)} != {tuple(expected.shape)}") + return + deviation = float((got - expected).abs().max()) + if not deviation <= atol: + self.failures.append(f"{name}: max|Δ| = {deviation:.6e} > atol {atol:.1e}") + print(f" {name:<34} max|Δ| {deviation:.3e}") + + def scalar(self, name: str, got: float, want: float, atol: float) -> None: + """Assert ``|got - want| <= atol`` on a single number.""" + deviation = abs(got - want) + if not deviation <= atol: + self.failures.append( + f"{name}: got {got!r}, want {want!r} (|Δ| = {deviation:.6e} > atol {atol:.1e})" + ) + print(f" {name:<34} |Δ| {deviation:.3e}") + + def truth(self, name: str, holds: bool, message: str) -> None: + """Assert a boolean contract (key present / absent, error raised).""" + if not holds: + self.failures.append(f"{name}: {message}") + print(f" {name:<34} {'ok' if holds else 'FAILED'}") + + +def check_matpes(checker: Checker) -> None: + """MatPES: energy + forces, the periodic seam, and Newton's third law.""" + print("MACEMatpes variant (density interactions, ZBL, per-layer readouts)") + potential = matpes_potential() + + out = potential(make_batch()) + energy = out["graphs", "energy"].detach() + forces = out["atoms", "forces"].detach() + checker.close("matpes.energy", energy, MATPES.energy, GOLDEN_ATOL) + checker.close("matpes.forces", forces, MATPES.forces, GOLDEN_ATOL) + checker.scalar( + "matpes.force_abs_sum", float(forces.abs().sum()), MATPES.force_abs_sum, GOLDEN_ATOL + ) + checker.scalar("matpes.net_force", float(forces.sum(0).abs().max()), 0.0, NET_FORCE_ATOL) + + periodic = potential(make_batch(shifts=SHIFTS)) + periodic_forces = periodic["atoms", "forces"].detach() + checker.close( + "matpes.periodic.energy", + periodic["graphs", "energy"].detach(), + MATPES_PERIODIC_ENERGY, + GOLDEN_ATOL, + ) + checker.scalar( + "matpes.periodic.force_abs_sum", + float(periodic_forces.abs().sum()), + MATPES_PERIODIC_FORCE_ABS_SUM, + GOLDEN_ATOL, + ) + + +def check_matpes_energy_seams(checker: Checker) -> None: + """The three energy-only entries must return the *same* energy, not a close one.""" + print("\nEnergy-only seams (energy_core / compute_forces=False / call_energy)") + potential = matpes_potential() + + with torch.no_grad(): + core = potential.energy_core(POSITIONS.clone(), Z, EDGE_INDEX, BATCH, NUM_GRAPHS) + checker.close("energy_core.energy", core, MATPES.energy, SEAM_ATOL) + + energy_only = matpes_potential(compute_forces=False)(make_batch()) + checker.close( + "compute_forces=False.energy", + energy_only["graphs", "energy"].detach(), + MATPES.energy, + SEAM_ATOL, + ) + checker.truth( + "compute_forces=False.no_forces", + ("atoms", "forces") not in energy_only.keys(include_nested=True), + "an energy-only potential must not write atoms.forces", + ) + + protocol = call_energy(potential, make_batch()) + checker.close( + "call_energy.energy", protocol["graphs", "energy"].detach(), MATPES.energy, SEAM_ATOL + ) + checker.truth( + "call_energy.no_forces", + ("atoms", "forces") not in protocol.keys(include_nested=True), + "call_energy must reach the energy core, never the force pipeline", + ) + + rejected = False + try: + potential.energy_core( + POSITIONS.clone(), Z, EDGE_INDEX, BATCH, NUM_GRAPHS, total_charge=TOTAL_CHARGE + ) + except ValueError: + rejected = True + checker.truth( + "matpes.rejects_total_charge", + rejected, + "a MatPES potential silently ignoring total_charge returns a wrong energy", + ) + + +def check_omol(checker: Checker) -> None: + """OMOL: energy + forces under charge/spin conditioning, and the cation shift.""" + print("\nMACEOMol variant (residual interactions, charge/spin conditioning)") + potential = omol_potential() + conditioning = {"total_charge": TOTAL_CHARGE, "total_spin": TOTAL_SPIN} + + out = potential(make_batch(graphs=conditioning)) + energy = out["graphs", "energy"].detach() + forces = out["atoms", "forces"].detach() + checker.close("omol.energy", energy, OMOL.energy, GOLDEN_ATOL) + checker.close("omol.forces", forces, OMOL.forces, GOLDEN_ATOL) + checker.scalar("omol.force_abs_sum", float(forces.abs().sum()), OMOL.force_abs_sum, GOLDEN_ATOL) + checker.scalar("omol.net_force", float(forces.sum(0).abs().max()), 0.0, NET_FORCE_ATOL) + checker.truth( + "omol.forces_are_not_vacuous", + float(forces.abs().max()) > 0.0, + "every OMOL force is exactly zero — the readout wake-up no longer works", + ) + + with torch.no_grad(): + cation = potential.energy_core( + POSITIONS.clone(), + Z, + EDGE_INDEX, + BATCH, + NUM_GRAPHS, + total_charge=CATION_CHARGE, + total_spin=TOTAL_SPIN, + ) + checker.close("omol.cation.energy", cation, OMOL_CATION_ENERGY, GOLDEN_ATOL) + + +def main() -> int: + """Run every seam of both variants against the embedded goldens.""" + checker = Checker() + check_matpes(checker) + check_matpes_energy_seams(checker) + check_omol(checker) + + if checker.failures: + print("\nFAILED — MACEPotential no longer reproduces the pre-restructure numbers:") + for failure in checker.failures: + print(f" {failure}") + return 1 + print("\nOK") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/regressions/mace-subpackage-restructure-05-checkpoint.py b/regressions/mace-subpackage-restructure-05-checkpoint.py new file mode 100644 index 0000000..a69efd7 --- /dev/null +++ b/regressions/mace-subpackage-restructure-05-checkpoint.py @@ -0,0 +1,675 @@ +"""Public-API checkpoint round trip for chain step mace-subpackage-restructure-05. + +Chain step 05 merges the two hand-copied official-checkpoint loaders +(``molzoo.mace_matpes.load_matpes_state_dict`` / +``molzoo.mace_omol.load_omol_state_dict``) into the one +:class:`molzoo.mace.CheckpointRemap`, and gives step 04's +:class:`molzoo.mace.MACEPotential` a ``from_checkpoint`` classmethod. Step 06 +then **deleted** the flat modules, so this file must never import them — it +uses ``molzoo.mace`` only and is expected to keep running unchanged afterwards. + +Scenario (public API only), start to finish inside one ``tempfile`` directory: + +1. build a tiny MatPES potential from a :class:`molzoo.mace.MACEMatpesSpec` + (the step-04 API) and fill **every** parameter, plus the fitted-constant + buffers, from a shape-derived deterministic rule — ``torch.linspace``, no + RNG at all, so these goldens cannot be knocked over by a change in torch's + random-number order or in module initialisation; +2. dump that ``state_dict`` under the **official cueq key names** and write it + next to a minimal official-style config json whose ``hidden_irreps`` / + ``MLP_irreps`` are the *only* statement of the channel widths — a + ``from_checkpoint`` that hard-coded ``128 / 1 / 16`` would raise on shapes + here; +3. read it back with ``MACEPotential.from_checkpoint(config, weights)`` and + assert the reconstruction is exact: every ``state_dict`` tensor + ``torch.equal`` to the source model's; +4. evaluate energy and forces on a fixed 5-atom cluster (fp64, CPU, + ``use_fallback=True``) against the hard-coded goldens below; +5. exercise the strictness doctrine on the cheap: a checkpoint key with no home + raises ``RuntimeError`` under :data:`molzoo.mace.MATPES_REMAP` and leaves the + model untouched, while the same key is merely *reported* by a + ``CheckpointRemap(MATPES_KEY_REMAP, on_unexpected="return")``. + +The synthetic checkpoint is written in the real dialect, not a convenient one: +it carries the graph constants and an ``output_mask`` (both skipped on load), +the ``r_max`` / ``num_interactions`` / ``cutoff_fn.p`` / ``pair_repulsion_fn.p`` +scalars and the ``atomic_energies_fn`` table (all *dropped* by the remap — E0 +comes from the config). Every one of those would otherwise surface as a key +"with no home" and abort the load, so the round trip passing is itself the +assertion that the drop/skip rules still hold. + +Goldens +------- +Captured by running this file at the commit below and pasting back what it +printed. There is no external oracle and none is possible: the numbers are what +*this repo* computes for a synthetic, deterministically filled model — the +point of the file is that the checkpoint round trip does not move them. + + capture command : PYTHONPATH=src python \\ + regressions/mace-subpackage-restructure-05-checkpoint.py + commit : cf50d3a (cf50d3ac6af1caa151669cc36d5baba8bfc5a988) plus + the uncommitted step-05 working tree + (``src/molzoo/mace/checkpoint.py`` and + ``MACEPotential.from_checkpoint``) + torch : 2.12.1+cpu (python 3.14.5) + date : 2026-08-08 + device / dtype : CPU, float64 (``config.set_precision("fp64")``) + oracle : this repository, self-consistent. No mace-torch, no e3nn, + no ASE, no network, no subprocess — and no RNG, seeded or + otherwise (see :func:`fill_ramp`). + observed : two consecutive runs byte-identical; the loaded model + reproduced the source model's energy and forces exactly + (deviation 0.0, not merely inside the band); Σ|F| and the + per-atom force table are pinned below. + +Tolerances follow ``tests/test_molzoo/test_mace_matpes.py:96-97``: 1e-9 eV on +energies, 1e-8 eV/Å on forces. + +Run: + PYTHONPATH=src python regressions/mace-subpackage-restructure-05-checkpoint.py +""" + +from __future__ import annotations + +import json +import sys +import tempfile +from pathlib import Path + +import torch + +from molix import config + +config.set_precision("fp64") # before any module construction: cuEq bakes in dtype + +from tensordict import TensorDict # noqa: E402 + +from molzoo.mace import ( # noqa: E402 + MATPES_KEY_REMAP, + MATPES_REMAP, + CheckpointRemap, + MACEMatpesSpec, + MACEPotential, +) + +#: Energy band in eV and force band in eV/Å — the fp64 tolerances of +#: ``tests/test_molzoo/test_mace_matpes.py:96-97``. +ENERGY_ATOL = 1e-9 +FORCE_ATOL = 1e-8 + +#: The loaded model holds the *same* tensors as the source, so it must run the +#: *same* arithmetic: its energy and forces are compared to the source model's +#: at exactly zero tolerance, not at :data:`ENERGY_ATOL`. +ROUNDTRIP_ATOL = 0.0 + +#: Newton's third law on an isolated cluster (measured ≤ 1e-16 eV/Å). +NET_FORCE_ATOL = 1e-10 + +# --------------------------------------------------------------------------- +# System: the fixed 5-atom cluster of +# regressions/mace-subpackage-restructure-04-potential.py, all 20 ordered pairs +# as directed edges. Literal coordinates in Å — reproducible anywhere. +# --------------------------------------------------------------------------- +ATOMIC_NUMBERS: list[int] = [1, 6, 8] +ATOMIC_ENERGIES: list[float] = [-13.6, -1029.0, -2041.0] # eV/atom + +POSITIONS = torch.tensor( + [ + [0.00, 0.00, 0.00], + [1.09, 0.00, 0.00], + [1.70, 1.15, 0.00], + [-0.40, 0.95, 0.30], + [2.10, -0.85, -0.50], + ], + dtype=torch.float64, +) +Z = torch.tensor([1, 6, 8, 1, 1], dtype=torch.long) +BATCH = torch.zeros(5, dtype=torch.long) +NUM_GRAPHS = 1 +EDGE_INDEX = torch.tensor( # (E, 2): [:, 0] = source, [:, 1] = target + [(i, j) for i in range(5) for j in range(5) if i != j], dtype=torch.long +) + +# --------------------------------------------------------------------------- +# The synthetic checkpoint's hyper-parameters. ``num_features`` / +# ``max_hidden_l`` / ``mlp_dim`` appear in the config json *only* as the irreps +# strings: ``from_checkpoint`` has to parse 16 / 1 / 8 out of them. +# --------------------------------------------------------------------------- +NUM_FEATURES = 16 +MAX_HIDDEN_L = 1 +MLP_DIM = 8 +HIDDEN_IRREPS = "16x0e+16x1o" +MLP_IRREPS = "8x0e" + +#: ``atomic_inter_scale`` / ``atomic_inter_shift`` as written into the config +#: json. The checkpoint's own ``scale_shift`` buffers are deliberately filled +#: *away* from these two values (see :func:`fill_potential`), so the loaded +#: model can only match the goldens if the checkpoint's numbers won. +CONFIG_SCALE = 0.7 +CONFIG_SHIFT = -0.25 + +SPEC_KWARGS = dict( + r_max=5.0, + num_bessel=4, + num_polynomial_cutoff=5, + l_max=1, + num_features=NUM_FEATURES, + max_hidden_l=MAX_HIDDEN_L, + num_interactions=2, + correlation=2, + mlp_dim=MLP_DIM, + radial_mlp=[8], + scale=CONFIG_SCALE, + shift=CONFIG_SHIFT, +) + +# --------------------------------------------------------------------------- +# Deterministic, RNG-free fill +# --------------------------------------------------------------------------- + +#: Half-width of the ``torch.linspace`` ramp every parameter is filled with. +#: Comparable to the ``N(0, 1)`` initialisation cuEquivariance linears use, and +#: chosen empirically: this tiny model's symmetric contraction (body order 2) +#: leaves the physical regime somewhere above ~1.4 — at 1.5 the forces reach +#: 4e2 eV/Å. At 1.2 the learned readouts still carry 0.30 eV of the energy and +#: 7e-2 eV/Å of the forces, so the goldens are not a ZBL-only measurement. +FILL_SPAN = 1.2 + +#: Per-name phase added to the ramp, so two tensors of the same length never +#: receive the same values (a remap that swapped two same-shaped keys would +#: otherwise round-trip unnoticed). A character sum, deliberately not +#: :func:`hash`, which is salted per process. +PHASE_MODULUS = 17 + +#: Fitted physical constants an official checkpoint carries and the remap +#: therefore has to transport: the Agnesi transform, the ZBL pair term and the +#: affine energy normalisation. They are perturbed rather than overwritten — +#: 1e-2 keeps every covalent radius and screening coefficient positive, so the +#: model stays sane, while still moving each buffer off the value construction +#: would give it. Anything outside these three families (the cuEquivariance +#: graph constants, the symmetric-contraction projection, the E0 table) is a +#: structural constant, rebuilt identically on both sides, and is left alone. +FILLED_BUFFER_PREFIXES = ("distance_transform.", "pair_repulsion.", "scale_shift.") +BUFFER_FILL_SCALE = 0.01 + +# --------------------------------------------------------------------------- +# The official cueq key dialect, written out by hand +# --------------------------------------------------------------------------- + +#: molnex prefix → official cueq prefix for the entries an official MatPES +#: checkpoint **carries into the model**. Hand-written rather than derived from +#: :data:`molzoo.mace.MATPES_KEY_REMAP`, so a wrong entry in the shipped table +#: cannot cancel itself out in this round trip; +#: :func:`check_official_dialect` asserts the two still agree. +CARRIED_OFFICIAL_NAMES: dict[str, str] = { + "node_embedding.": "node_embedding.linear.", + "bessel.freqs": "radial_embedding.bessel_fn.bessel_weights", + "distance_transform.": "radial_embedding.distance_transform.", + "pair_repulsion.": "pair_repulsion_fn.", + "z_table": "atomic_numbers", +} + +#: molnex prefix → official cueq prefix for entries the checkpoint carries and +#: the remap **drops**: the isolated-atom reference energies are rebuilt from +#: the config's ``atomic_energies``, Z-indexed, not read out of the weights. +DROPPED_OFFICIAL_NAMES: dict[str, str] = { + "atomic_energies.": "atomic_energies_fn.", +} + +#: Official-only entries with no molnex counterpart at all. Each is dropped by +#: a ``None`` row of the shipped table — exactly (``r_max``, +#: ``num_interactions``, ``pair_repulsion_fn.p``) or through the enclosing +#: prefix (``radial_embedding.cutoff_fn.``); if any stopped being dropped, the +#: load would abort with "no home" instead of reaching the goldens. +#: ``pair_repulsion_fn.p`` is also the longest-prefix probe — its exact row +#: must beat the ``pair_repulsion_fn.`` prefix that encloses it, or it would +#: land as ``pair_repulsion.p`` and be homeless. +OFFICIAL_ONLY_SCALARS: dict[str, float] = { + "r_max": 5.0, + "num_interactions": 2.0, + "radial_embedding.cutoff_fn.p": 5.0, + "pair_repulsion_fn.p": 5.0, +} + +#: An irrep mask of the kind cuEquivariance emits and rebuilds. Skipped by +#: ``rename`` on the ``output_mask`` suffix; were it not, it would be a key +#: with no home and the load would raise. +OUTPUT_MASK_KEY = "interactions.0.conv_tp.output_mask" + +#: A checkpoint key that matches no table entry and no model parameter — the +#: probe for the ``on_unexpected`` knob (the same key the unit tests use). +MYSTERY_KEY = "interactions.0.mystery_layer.weight" + +# --------------------------------------------------------------------------- +# Goldens — see the module docstring for provenance. +# --------------------------------------------------------------------------- + +#: Total energy of the cluster, eV, ``(B,)``. +ENERGY: list[float] = [-3111.798363383643] + +#: Forces, eV/Å, ``(N, 3)``. +FORCES: list[list[float]] = [ + [0.02073978897711603, -0.014283907947197812, -0.004580298404743775], + [-0.07674001496483682, -0.1308470223078055, -0.002087429078993083], + [0.10544897994819648, 0.1571999433835074, 0.0003340549276333903], + [0.022217895631466936, -0.04160730323074992, -0.012362154174687933], + [-0.07166664959194265, 0.02953829010224585, 0.0186958267307914], +] + +#: Σ|F| in eV/Å — one number that moves if any component does. +FORCE_ABS_SUM = 0.708349559401915 + +#: ``scale_shift`` as it comes back out of the checkpoint. Both differ from the +#: config's :data:`CONFIG_SCALE` / :data:`CONFIG_SHIFT`, which is what makes +#: "the checkpoint's value won" an observable claim. +SCALE_GOLDEN = 0.6886 +SHIFT_GOLDEN = -0.2609 + + +# --------------------------------------------------------------------------- +# Construction +# --------------------------------------------------------------------------- + + +def fill_ramp(name: str, count: int) -> torch.Tensor: + """Deterministic ``count`` values for the tensor called ``name``. + + A ``torch.linspace`` ramp over ``±`` :data:`FILL_SPAN` plus a per-name + phase. No RNG is involved anywhere in this file, so the goldens survive any + change to torch's random-number order or to module initialisation — the + failure mode that makes seeded goldens rot. + + ``torch.linspace(-s, s, 1)`` is ``[-s]``, so a single-element tensor gets + ``-FILL_SPAN + phase``. + + Args: + name: ``state_dict`` key, used only for the phase. + count: Number of values to produce. + + Returns: + A flat ``(count,)`` float64 tensor. + """ + phase = (sum(ord(character) for character in name) % PHASE_MODULUS) / 100.0 + return torch.linspace(-FILL_SPAN, FILL_SPAN, count, dtype=torch.float64) + phase + + +def fill_potential(potential: MACEPotential) -> None: + """Overwrite every parameter, and the fitted buffers, in place. + + Every ``nn.Parameter`` is *replaced* by :func:`fill_ramp`, so nothing the + constructor's RNG produced survives; the three fitted-constant buffer + families of :data:`FILLED_BUFFER_PREFIXES` are *perturbed* by + :data:`BUFFER_FILL_SCALE` times the same ramp, which keeps them physical + (positive radii, positive ZBL coefficients) while moving them off the value + a fresh construction would give — without that, asserting the checkpoint + restored them would be vacuous. + + Args: + potential: Model to fill in place. + + Raises: + RuntimeError: If a buffer perturbation left the tensor where it was, + which would silently hollow out the round-trip assertion. + """ + with torch.no_grad(): + for name, parameter in potential.named_parameters(): + parameter.copy_(fill_ramp(name, parameter.numel()).reshape(parameter.shape)) + for name, buffer in potential.named_buffers(): + if not name.startswith(FILLED_BUFFER_PREFIXES): + continue + ramp = fill_ramp(name, buffer.numel()).reshape(buffer.shape) + moved = buffer + BUFFER_FILL_SCALE * ramp + if torch.equal(moved, buffer): + raise RuntimeError(f"buffer {name!r} did not move — the fill asserts nothing") + buffer.copy_(moved) + + +def source_potential() -> MACEPotential: + """The tiny MatPES potential the synthetic checkpoint is dumped from.""" + spec = MACEMatpesSpec( + atomic_numbers=ATOMIC_NUMBERS, atomic_energies=ATOMIC_ENERGIES, **SPEC_KWARGS + ) + potential = MACEPotential(spec, use_fallback=True) + fill_potential(potential) + return potential.eval() + + +def official_state(potential: MACEPotential) -> dict[str, torch.Tensor]: + """Dump ``potential``'s weights under the official cueq key names. + + The inverse of the shipped remap, plus the entries a real converted + checkpoint carries and molnex has no home for: the dropped scalars of + :data:`OFFICIAL_ONLY_SCALARS` and one :data:`OUTPUT_MASK_KEY`. The + cuEquivariance graph constants keep their names — they are unlisted in the + table and skipped on the ``.graph.c`` rule. + + Args: + potential: Source model. + + Returns: + Official key → the same tensor objects. + """ + renames = sorted( + {**CARRIED_OFFICIAL_NAMES, **DROPPED_OFFICIAL_NAMES}.items(), + key=lambda item: -len(item[0]), + ) + state: dict[str, torch.Tensor] = {} + for key, value in potential.state_dict().items(): + for molnex, official in renames: + if key == molnex or key.startswith(molnex): + key = official + (key[len(molnex) :] if key != molnex else "") + break + state[key] = value + for key, value in OFFICIAL_ONLY_SCALARS.items(): + state[key] = torch.tensor(value, dtype=torch.float64) + state[OUTPUT_MASK_KEY] = torch.ones(4, dtype=torch.float64) + return state + + +def write_checkpoint(directory: Path, state: dict[str, torch.Tensor]) -> tuple[Path, Path]: + """Write an official-style MatPES config json + cueq ``state_dict``. + + The config states the channel widths **only** through ``hidden_irreps`` / + ``MLP_irreps``; ``num_features`` / ``max_hidden_l`` / ``mlp_dim`` are + nowhere in the file, so a loader that assumed the shipped ``128 / 1 / 16`` + would build a model these weights do not fit. + + Args: + directory: Destination (a ``tempfile`` directory). + state: Officially named weights. + + Returns: + ``(config_path, weights_path)``. + """ + config_path = directory / "matpes_synthetic_config.json" + config_path.write_text( + json.dumps( + { + "atomic_numbers": ATOMIC_NUMBERS, + "atomic_energies": ATOMIC_ENERGIES, + "r_max": SPEC_KWARGS["r_max"], + "num_bessel": SPEC_KWARGS["num_bessel"], + "num_polynomial_cutoff": SPEC_KWARGS["num_polynomial_cutoff"], + "max_ell": SPEC_KWARGS["l_max"], + "num_interactions": SPEC_KWARGS["num_interactions"], + "correlation": SPEC_KWARGS["correlation"], + "hidden_irreps": HIDDEN_IRREPS, + "MLP_irreps": MLP_IRREPS, + "radial_MLP": SPEC_KWARGS["radial_mlp"], + "atomic_inter_scale": CONFIG_SCALE, + "atomic_inter_shift": CONFIG_SHIFT, + }, + indent=2, + ) + ) + weights_path = directory / "matpes_synthetic_cueq_state.pt" + torch.save(state, weights_path) + return config_path, weights_path + + +def make_batch() -> TensorDict: + """Post-collate batch for :data:`POSITIONS` (atoms / edges / graphs). + + A fresh object every call: ``forward`` writes in place and swaps + ``atoms.pos`` for its own differentiation leaf. + + Returns: + A ``TensorDict`` with ``edge_diff = pos[target] - pos[source]`` and + ``edge_dist = ‖edge_diff‖``. + """ + pos = POSITIONS.clone() + edge_diff = pos[EDGE_INDEX[:, 1]] - pos[EDGE_INDEX[:, 0]] + return TensorDict( + atoms=TensorDict(Z=Z, pos=pos, batch=BATCH, batch_size=[pos.shape[0]]), + edges=TensorDict( + edge_index=EDGE_INDEX, + edge_diff=edge_diff, + edge_dist=edge_diff.norm(dim=-1), + batch_size=[EDGE_INDEX.shape[0]], + ), + graphs=TensorDict( + num_atoms=torch.tensor([pos.shape[0]], dtype=torch.long), batch_size=[NUM_GRAPHS] + ), + batch_size=[], + ) + + +def energy_forces(potential: MACEPotential) -> tuple[torch.Tensor, torch.Tensor]: + """Energy ``(B,)`` in eV and forces ``(N, 3)`` in eV/Å for the cluster.""" + out = potential(make_batch()) + return out["graphs", "energy"].detach(), out["atoms", "forces"].detach() + + +# --------------------------------------------------------------------------- +# Checking +# --------------------------------------------------------------------------- + + +class Checker: + """Collects every deviation so one run reports all failures, not the first.""" + + def __init__(self) -> None: + self.failures: list[str] = [] + + def close( + self, + name: str, + got: torch.Tensor, + want: list[float] | list[list[float]], + atol: float, + ) -> None: + """Assert ``max|got - want| <= atol`` against a literal nested list.""" + expected = torch.tensor(want, dtype=torch.float64) + if got.shape != expected.shape: + self.failures.append(f"{name}: shape {tuple(got.shape)} != {tuple(expected.shape)}") + return + deviation = float((got - expected).abs().max()) + if not deviation <= atol: + self.failures.append(f"{name}: max|Δ| = {deviation:.6e} > atol {atol:.1e}") + print(f" {name:<38} max|Δ| {deviation:.3e}") + + def scalar(self, name: str, got: float, want: float, atol: float) -> None: + """Assert ``|got - want| <= atol`` on a single number.""" + deviation = abs(got - want) + if not deviation <= atol: + self.failures.append( + f"{name}: got {got!r}, want {want!r} (|Δ| = {deviation:.6e} > atol {atol:.1e})" + ) + print(f" {name:<38} |Δ| {deviation:.3e}") + + def truth(self, name: str, holds: bool, message: str) -> None: + """Assert a boolean contract (key present / absent, error raised).""" + if not holds: + self.failures.append(f"{name}: {message}") + print(f" {name:<38} {'ok' if holds else 'FAILED'}") + + +def check_official_dialect( + checker: Checker, source: MACEPotential, state: dict[str, torch.Tensor] +) -> None: + """The synthetic checkpoint really is written in the official dialect. + + Two table-level legs against the shipped + :data:`molzoo.mace.MATPES_KEY_REMAP` (so a wrong entry there cannot cancel + itself out in the round trip), then two behavioural ones through the pure + :meth:`molzoo.mace.CheckpointRemap.rename`: everything this file wrote as + droppable leaves no trace, and everything that survives has a home. + + Args: + checker: Failure collector. + source: The model the checkpoint was dumped from. + state: The officially named checkpoint contents. + """ + print("Official key dialect (hand-written inverse vs molzoo.mace.MATPES_KEY_REMAP)") + + wrong = [ + f"{official} → {MATPES_KEY_REMAP.get(official)!r}, want {molnex!r}" + for molnex, official in CARRIED_OFFICIAL_NAMES.items() + if MATPES_KEY_REMAP.get(official) != molnex + ] + checker.truth("dialect.carried_entries_agree", not wrong, f"shipped table disagrees: {wrong}") + + carried = {value for value in MATPES_KEY_REMAP.values() if value is not None} + checker.truth( + "dialect.covers_every_carried_family", + carried == set(CARRIED_OFFICIAL_NAMES), + f"shipped table carries {sorted(carried)}, this file knows " + f"{sorted(CARRIED_OFFICIAL_NAMES)} — a new family needs a golden", + ) + + renamed = set(MATPES_REMAP.rename(state)) + droppable = { + *OFFICIAL_ONLY_SCALARS, + OUTPUT_MASK_KEY, + *(key for key in state if ".graph.c" in key), + *(key for key in state if key.startswith(tuple(DROPPED_OFFICIAL_NAMES.values()))), + } + checker.truth( + "dialect.droppable_keys_vanish", + not droppable & renamed, + f"expected no trace of {sorted(droppable & renamed)} after rename", + ) + checker.truth( + "dialect.every_survivor_has_a_home", + not renamed - set(source.state_dict()), + f"renamed keys with no home: {sorted(renamed - set(source.state_dict()))}", + ) + + +def check_roundtrip(checker: Checker, source: MACEPotential, built: MACEPotential) -> None: + """``from_checkpoint`` reconstructs the source model tensor for tensor.""" + print("\nCheckpoint round trip (from_checkpoint vs the model it was dumped from)") + reference = source.state_dict() + produced = built.state_dict() + + checker.truth( + "roundtrip.same_keys", + set(produced) == set(reference), + f"only in checkpoint model: {sorted(set(produced) - set(reference))}; " + f"only in source: {sorted(set(reference) - set(produced))}", + ) + differing = [name for name, want in reference.items() if not torch.equal(produced[name], want)] + checker.truth( + "roundtrip.every_tensor_equal", + not differing, + f"{len(differing)} tensor(s) not bit-identical: {differing[:5]}", + ) + + # The widths were parsed out of "16x0e+16x1o" / "8x0e"; had they been + # assumed, the shapes below would not have matched and the load would have + # raised on shape mismatch long before this line. + checker.truth( + "roundtrip.widths_came_from_irreps", + tuple(built.node_embedding.weight.shape) == (1, len(ATOMIC_NUMBERS) * NUM_FEATURES) + and built.readouts[-1].linear_2.weight.numel() == MLP_DIM, + "the derived channel widths do not match the config's irreps strings", + ) + + scale = float(built.scale_shift.scale) + shift = float(built.scale_shift.shift) + checker.scalar("checkpoint.scale_shift.scale", scale, SCALE_GOLDEN, 0.0) + checker.scalar("checkpoint.scale_shift.shift", shift, SHIFT_GOLDEN, 0.0) + checker.truth( + "checkpoint.beat_the_config_scale", + scale != CONFIG_SCALE and shift != CONFIG_SHIFT, + "scale/shift equal the config values — the checkpoint's own numbers were not loaded", + ) + + +def check_energy_forces(checker: Checker, source: MACEPotential, built: MACEPotential) -> None: + """Energy and forces of the loaded model, against the goldens and the source.""" + print("\nEnergy / forces of the loaded model (fp64, CPU, use_fallback=True)") + energy, forces = energy_forces(built) + + checker.close("loaded.energy", energy, ENERGY, ENERGY_ATOL) + checker.close("loaded.forces", forces, FORCES, FORCE_ATOL) + checker.scalar("loaded.force_abs_sum", float(forces.abs().sum()), FORCE_ABS_SUM, FORCE_ATOL) + checker.scalar("loaded.net_force", float(forces.sum(0).abs().max()), 0.0, NET_FORCE_ATOL) + + source_energy, source_forces = energy_forces(source) + checker.scalar( + "loaded.vs_source.energy", + float((energy - source_energy).abs().max()), + 0.0, + ROUNDTRIP_ATOL, + ) + checker.scalar( + "loaded.vs_source.forces", + float((forces - source_forces).abs().max()), + 0.0, + ROUNDTRIP_ATOL, + ) + + +def check_unhoused_key( + checker: Checker, built: MACEPotential, state: dict[str, torch.Tensor] +) -> None: + """The ``on_unexpected`` knob: MatPES refuses, the lenient policy reports. + + One leg of the strictness doctrine, cheaply: a checkpoint key with no home + means the mapping is stale, and loading the rest would leave a model that + runs, looks sane and is quietly wrong (that doctrine, and the two incidents + behind it, are quoted in the :mod:`molzoo.mace.checkpoint` docstring). + + Args: + checker: Failure collector. + built: The loaded model, re-checked afterwards — the refusal must not + have left a half-loaded model behind. + state: The officially named checkpoint contents; a tampered *copy* is + what reaches the remap. + """ + print("\nStrictness doctrine (a checkpoint key with no home)") + tampered = dict(state) + tampered[MYSTERY_KEY] = torch.zeros(3, dtype=torch.float64) + + message = "" + try: + MATPES_REMAP.load(built, tampered) + except RuntimeError as error: + message = str(error) + checker.truth( + "unhoused.raises_no_home", + "no home" in message and MYSTERY_KEY in message, + f"expected a RuntimeError naming the key, got {message!r}", + ) + + _, unexpected = CheckpointRemap(MATPES_KEY_REMAP, on_unexpected="return").load(built, tampered) + checker.truth( + "unhoused.returned_under_lenient_policy", + unexpected == [MYSTERY_KEY], + f"expected [{MYSTERY_KEY!r}], got {unexpected}", + ) + + # Neither call may have damaged the model: the refusal happens before any + # tensor is placed, and the lenient reload puts back the same weights. + energy, forces = energy_forces(built) + checker.close("unhoused.energy_unchanged", energy, ENERGY, ENERGY_ATOL) + checker.close("unhoused.forces_unchanged", forces, FORCES, FORCE_ATOL) + + +def main() -> int: + """Build a synthetic checkpoint, read it back, and check every golden.""" + checker = Checker() + source = source_potential() + state = official_state(source) + check_official_dialect(checker, source, state) + + with tempfile.TemporaryDirectory(prefix="molnex-mace-checkpoint-") as directory: + config_path, weights_path = write_checkpoint(Path(directory), state) + built = MACEPotential.from_checkpoint(config_path, weights_path, use_fallback=True).eval() + + check_roundtrip(checker, source, built) + check_energy_forces(checker, source, built) + check_unhoused_key(checker, built, state) + + if checker.failures: + print("\nFAILED — the MACE checkpoint round trip no longer reproduces the goldens:") + for failure in checker.failures: + print(f" {failure}") + return 1 + print("\nOK") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/regressions/mace-subpackage-restructure-06-wire.py b/regressions/mace-subpackage-restructure-06-wire.py new file mode 100644 index 0000000..c351883 --- /dev/null +++ b/regressions/mace-subpackage-restructure-06-wire.py @@ -0,0 +1,624 @@ +"""Public import surface and parameter inventory for chain step 06-wire. + +Chain step 06 flips `molzoo` over to the `molzoo.mace` sub-package: it deletes +the flat `molzoo/mace_matpes.py` and `molzoo/mace_omol.py`, turns +`MACEMatpes` / `MACEOMol` / `load_matpes_state_dict` / `load_omol_state_dict` +into thin aliases over `molzoo.mace.variants`, and puts the top level on **one** +PEP 562 lazy policy so that `import molzoo` no longer drags in the +cuEquivariance stack. Everything a user can see is supposed to be unchanged +except that one import cost. This file is the standalone check of that claim +(spec `.claude/specs/mace-subpackage-restructure-06-wire.md`, ac-009). + +Scenario (public API only), in one process, in this order — the order *is* part +of the test: + +1. **Lazy policy.** The script is itself the clean interpreter the policy has to + be measured in, so the probe runs before anything else pulls torch: after + `import molzoo`, neither `torch` nor `cuequivariance_torch` may be in + `sys.modules`; after `import molzoo.mace` (whose config models are imported + eagerly) both must *still* be absent; `from molix import config` then brings + torch but **not** cuEquivariance, which is what makes step 4 attributable; + touching `molzoo.MACE` finally pulls the equivariance stack in. A unit test + needs a subprocess for this (`tests/test_molzoo/test_imports.py`) because a + pytest session has imported cuEquivariance long before collection; here the + process is fresh by construction. It is also asserted to *be* fresh, so + importing this file from an already-warm interpreter fails loudly instead of + passing vacuously. +2. **Import surface.** The ten top-level names and the whole + `molzoo.mace.__all__` resolve; the top-level aliases are the *same objects* + as the sub-package's, which is what "thin alias" has to mean; `dir()` on both + modules reports exactly `sorted(__all__)`, and an unknown name still raises + `AttributeError` rather than an `ImportError` from inside the package. The + `set(__all__) <= set(dir())` subset form of the ac-001 gate is asserted + verbatim alongside the stronger equality that in fact holds. +3. **Parameter inventory.** Tiny MatPES and OMol variants are built through the + keyword surface the scripts bind and their parameter tensor count and element + total are compared to the hard-coded pre-cutover baselines below. The same + two models are then built through the `MACEPotential(spec)` form and required + to agree name-for-name and shape-for-shape with the keyword form. + +Why a parameter inventory and not an energy. Energies are the natural golden but +the wrong one here: this restructure reorders module construction, which +reorders the draws from torch's global RNG, so a freshly initialised model's +energy moves for reasons that have nothing to do with correctness. The +parameter inventory is invariant to that and is the offline-checkable proxy for +the claim that matters — an unchanged tensor inventory is exactly the condition +under which the official MatPES / OMol checkpoints still load strictly, and the +strict load is what carries the external numerical verdicts (mace-omol-port-01 +ac-006, 7e-7 eV / 4.3e-6 eV·A) across the cutover. Re-running those verdicts is +impossible in-tree by policy: mace-torch and e3nn are not importable here. + +Goldens +------- +Integers, so there is no floating-point tolerance anywhere in this file — every +comparison below is exact equality. + + capture command : git worktree add --detach 0e05959 && \\ + PYTHONPATH=/src python -c ' + import torch + from molix import config + config.set_precision("fp64") + from molzoo.mace_matpes import MACEMatpes + from molzoo.mace_omol import MACEOMol + torch.manual_seed(0) + m = MACEMatpes(atomic_numbers=[1, 6, 8], + atomic_energies=torch.tensor([-13.6, -1029.0, -2041.0]), + **TINY_MATPES_KWARGS) + print(len(list(m.parameters())), + sum(p.numel() for p in m.parameters()))' + (and the same for MACEOMol / TINY_OMOL_KWARGS) + commit : 0e05959 (0e059591b923833986344a936e60941588588de7) — the + last commit at which the flat `molzoo/mace_matpes.py` and + `molzoo/mace_omol.py` still existed. The 06-wire working + tree that deletes them is uncommitted on top of it. + torch : 2.12.1+cpu (python 3.14.5) + date : 2026-08-09 + device / dtype : CPU, float64 (`config.set_precision("fp64")`), + `use_fallback=True` for MatPES; MACEOMol takes no such + keyword and its spec defaults to the fused path. + oracle : this repository at the pre-cutover commit. No mace-torch, + no e3nn, no ASE, no network, no subprocess. The dtype and + the fallback switch do not enter the counts; they are + pinned so the capture is reproducible, not because the + numbers depend on them. + observed : the flat pre-cutover classes and the post-cutover + `molzoo.mace.variants` classes produce not merely the same + two totals but the identical `named_parameters()` list — + same names, same shapes, same order — for both variants, + and the `MACEPotential(spec)` form matches both. Only the + two totals are pinned as literals here; the keyword-vs-spec + list equality is re-checked live (`check_construction_paths`) + since it needs no golden. + +Run: + PYTHONPATH=src python regressions/mace-subpackage-restructure-06-wire.py +""" + +from __future__ import annotations + +# --------------------------------------------------------------------------- +# Lazy-policy probe. Nothing above this may import torch, so the imports below +# are deliberately interleaved with the snapshots they are being measured by -- +# moving any of them breaks the measurement rather than merely reordering it. +# --------------------------------------------------------------------------- +import sys + +#: `sys.modules` before this file touched anything. Recorded so that running +#: the file from a warm interpreter (an `import` rather than `python `) +#: is reported as a broken measurement instead of a silent pass. +FRESH_INTERPRETER = "torch" not in sys.modules and "molzoo" not in sys.modules + +import molzoo # noqa: E402 + +#: Third-party module whose presence marks "the equivariance stack is loaded". +#: The expensive import the lazy policy exists to defer; same probe as +#: `tests/test_molzoo/test_imports.py`. +CUEQ = "cuequivariance_torch" + +#: Snapshot after `import molzoo`: the top level must cost `typing` and nothing +#: else, so torch has not been imported either. +AT_MOLZOO = {"torch": "torch" in sys.modules, "cueq": CUEQ in sys.modules} + +import molzoo.mace # noqa: E402 + +#: Snapshot after `import molzoo.mace`. This one imports `molzoo.mace.spec` +#: eagerly; the sub-package docstring's claim is that the config models are +#: torch-free, so even torch must still be absent here. +AT_MOLZOO_MACE = {"torch": "torch" in sys.modules, "cueq": CUEQ in sys.modules} + +from molix import config # noqa: E402 + +config.set_precision("fp64") # before any module construction: cuEq bakes in dtype + +#: Snapshot after the `molix` import that the rest of the file needs. torch is +#: now loaded and cuEquivariance still is not -- without this leg, "cuEq +#: appeared when we touched `molzoo.MACE`" would rest on the assumption that +#: nothing in between could have brought it. +AT_MOLIX = {"torch": "torch" in sys.modules, "cueq": CUEQ in sys.modules} + +#: The lazy resolution itself. `molzoo.__getattr__` does not cache onto the +#: module, so this is an honest first touch; it is kept for the identity check +#: in `check_alias_identity`. +LAZILY_RESOLVED_MACE = molzoo.MACE + +#: Snapshot after the first model-symbol access: lazy, not absent. +AFTER_ATTRIBUTE_ACCESS = {"torch": "torch" in sys.modules, "cueq": CUEQ in sys.modules} + +import torch # noqa: E402 + +from molzoo import ( # noqa: E402 + MACE, + Allegro, + AllegroSpec, + MACEMatpes, + MACEOMol, + MACESpec, + PiNet, + PiNetSpec, + load_matpes_state_dict, + load_omol_state_dict, +) +from molzoo.mace import MACE as MACE_FROM_SUBPACKAGE # noqa: E402 +from molzoo.mace import ( # noqa: E402 + CheckpointRemap, + MACEMatpesSpec, + MACEOMolSpec, + MACEPotential, +) + +# --------------------------------------------------------------------------- +# Configurations. Literal copies of `TINY_MATPES_KWARGS` / `TINY_OMOL_KWARGS` +# in `tests/test_molzoo/test_mace/conftest.py` -- copied rather than imported, +# because a regression script must not depend on the test suite and because a +# golden whose configuration can be edited elsewhere is not a golden. +# --------------------------------------------------------------------------- + +#: Element table (z-table), strictly ascending as `torch.searchsorted` needs. +ATOMIC_NUMBERS: list[int] = [1, 6, 8] + +#: Per-element reference energies `E0` in eV/atom, in `ATOMIC_NUMBERS` order. +ATOMIC_ENERGIES: list[float] = [-13.6, -1029.0, -2041.0] + +#: Tiny MACE-MatPES hyper-parameters (l_max=1, 16 channels): structurally the +#: shipped model, small enough to build in a second on CPU. +TINY_MATPES_KWARGS: dict[str, object] = { + "r_max": 5.0, + "num_bessel": 4, + "num_polynomial_cutoff": 5, + "l_max": 1, + "num_features": 16, + "max_hidden_l": 1, + "num_interactions": 2, + "correlation": 2, + "mlp_dim": 8, + "radial_mlp": [8], + "use_fallback": True, # CPU: the fused kernels need a GPU + the ops wheel +} + +#: Tiny MACE-OMol hyper-parameters (l_max=1, 16 channels). `use_fallback` is +#: absent on purpose: `MACEOMol` takes no such keyword, so the spec's default +#: (`False`, the fused path) is what both construction paths must agree on. +TINY_OMOL_KWARGS: dict[str, object] = { + "r_max": 5.0, + "num_bessel": 4, + "num_polynomial_cutoff": 5, + "l_max": 1, + "num_features": 16, + "num_interactions": 2, + "correlation": 2, + "mlp_dim": 8, + "edge_channels": 8, +} + +#: Construction seed. The counts below are seed-independent -- initialisation +#: fills tensors, it does not decide how many there are -- so this is hygiene, +#: not a load-bearing golden input. +SEED = 0 + +# --------------------------------------------------------------------------- +# Goldens -- see the module docstring for provenance. +# --------------------------------------------------------------------------- + +#: `len(list(m.parameters()))` and `sum(p.numel() for p in m.parameters())` for +#: `MACEMatpes(**TINY_MATPES_KWARGS)`, captured from the flat pre-cutover +#: `molzoo.mace_matpes` at commit 0e05959. +MATPES_PARAM_TENSORS = 21 +MATPES_PARAM_ELEMENTS = 6852 + +#: The same two totals for `MACEOMol(**TINY_OMOL_KWARGS)`, captured from the +#: flat pre-cutover `molzoo.mace_omol` at the same commit. Larger than MatPES +#: because of the charge/spin conditioning and the per-l edge channels. +OMOL_PARAM_TENSORS = 73 +OMOL_PARAM_ELEMENTS = 16795 + +#: The names `scripts/matpes_port/run_nve.py:45`, +#: `benchmarks/bench_mace_matpes.py:38` and `benchmarks/bench_trainer_throughput.py:26` +#: bind, plus the Allegro / PiNet pair that joined the same lazy table in 06. +#: Compared against `molzoo.__all__` as a set, so a name silently dropped from +#: the export list fails here even though the `from molzoo import ...` above +#: would still have succeeded through `__getattr__`. +EXPECTED_TOP_LEVEL: frozenset[str] = frozenset( + { + "Allegro", + "AllegroSpec", + "MACE", + "MACEMatpes", + "MACEOMol", + "MACESpec", + "PiNet", + "PiNetSpec", + "load_matpes_state_dict", + "load_omol_state_dict", + } +) + +#: The six names ac-001 requires `molzoo.mace.__all__` to contain *at least*. +#: The rest of that list is checked for resolvability, not for membership, so +#: adding a symbol to the sub-package does not fail this file. +REQUIRED_SUBPACKAGE_NAMES: frozenset[str] = frozenset( + { + "MACE", + "MACESpec", + "MACEMatpes", + "MACEOMol", + "load_matpes_state_dict", + "load_omol_state_dict", + } +) + +#: A name in neither lazy table: the probe for "the export tables are closed". +UNKNOWN_NAME = "NotAThing" + + +# --------------------------------------------------------------------------- +# Checking +# --------------------------------------------------------------------------- + + +class Checker: + """Collects every deviation so one run reports all failures, not the first.""" + + def __init__(self) -> None: + self.failures: list[str] = [] + + def count(self, name: str, got: int, want: int) -> None: + """Assert an integer golden exactly -- no tolerance applies to a count.""" + if got != want: + self.failures.append(f"{name}: got {got}, want {want}") + print(f" {name:<46} {got:>6} {'ok' if got == want else 'FAILED'}") + + def truth(self, name: str, holds: bool, message: str) -> None: + """Assert a boolean contract (a name resolves, a module is absent, ...).""" + if not holds: + self.failures.append(f"{name}: {message}") + print(f" {name:<46} {'ok' if holds else 'FAILED':>6}") + + +def check_lazy_policy(checker: Checker) -> None: + """The four `sys.modules` snapshots taken while this file was importing. + + The whole point of the ladder is attribution: each snapshot narrows what + could have loaded the equivariance stack, so the final `cueq is present` + can only be blamed on the `molzoo.MACE` access. + + Args: + checker: Failure collector. + """ + print("Lazy policy (snapshots taken during this file's own import)") + checker.truth( + "lazy.interpreter_was_fresh", + FRESH_INTERPRETER, + "torch or molzoo was already imported -- run this file as the entry " + "point (`python regressions/...py`), not as an imported module; the " + "snapshots below measure nothing otherwise", + ) + checker.truth( + "lazy.import_molzoo_is_torch_free", + AT_MOLZOO["torch"] is False, + "`import molzoo` pulled torch -- the top level is supposed to import nothing but `typing`", + ) + checker.truth( + "lazy.import_molzoo_is_cueq_free", + AT_MOLZOO["cueq"] is False, + f"`import molzoo` pulled {CUEQ} -- the lazy policy is not in force", + ) + checker.truth( + "lazy.import_molzoo_mace_is_torch_free", + AT_MOLZOO_MACE["torch"] is False, + "`import molzoo.mace` pulled torch -- its eagerly imported config " + "models are documented as torch-free", + ) + checker.truth( + "lazy.import_molzoo_mace_is_cueq_free", + AT_MOLZOO_MACE["cueq"] is False, + f"`import molzoo.mace` pulled {CUEQ} -- the sub-package re-export " + "surface must stay lazy too", + ) + checker.truth( + "lazy.molix_brings_torch_not_cueq", + AT_MOLIX["torch"] is True and AT_MOLIX["cueq"] is False, + f"expected torch and no {CUEQ} after `from molix import config`, got " + f"{AT_MOLIX} -- the next leg cannot attribute the stack to the " + "attribute access", + ) + checker.truth( + "lazy.attribute_access_brings_cueq", + AFTER_ATTRIBUTE_ACCESS["cueq"] is True, + f"`molzoo.MACE` did not pull {CUEQ} -- lazy has become absent, and the " + "symbol is resolving to something other than the real encoder", + ) + checker.truth( + "lazy.resolved_symbol_is_a_class", + isinstance(LAZILY_RESOLVED_MACE, type), + f"`molzoo.MACE` resolved to {LAZILY_RESOLVED_MACE!r}, not a class", + ) + + +def check_import_surface(checker: Checker) -> None: + """Both export tables: complete, closed, and reported by `dir()`. + + Args: + checker: Failure collector. + """ + print("\nImport surface (molzoo and molzoo.mace export tables)") + checker.truth( + "surface.top_level_all_is_the_expected_set", + set(molzoo.__all__) == EXPECTED_TOP_LEVEL, + f"molzoo.__all__ is {sorted(molzoo.__all__)}, expected {sorted(EXPECTED_TOP_LEVEL)}", + ) + unresolved_top = [name for name in molzoo.__all__ if getattr(molzoo, name, None) is None] + checker.truth( + "surface.every_top_level_name_resolves", + not unresolved_top, + f"lazy table advertises names that do not resolve: {unresolved_top}", + ) + checker.truth( + "surface.top_level_symbols_are_bound", + all( + symbol is not None + for symbol in ( + MACE, + MACESpec, + MACEMatpes, + MACEOMol, + load_matpes_state_dict, + load_omol_state_dict, + Allegro, + AllegroSpec, + PiNet, + PiNetSpec, + ) + ), + "one of the ten `from molzoo import ...` bindings is None", + ) + checker.truth( + "surface.subpackage_all_covers_required", + REQUIRED_SUBPACKAGE_NAMES <= set(molzoo.mace.__all__), + f"molzoo.mace.__all__ is missing " + f"{sorted(REQUIRED_SUBPACKAGE_NAMES - set(molzoo.mace.__all__))}", + ) + unresolved_sub = [ + name for name in molzoo.mace.__all__ if getattr(molzoo.mace, name, None) is None + ] + checker.truth( + "surface.every_subpackage_name_resolves", + not unresolved_sub, + f"molzoo.mace advertises names that do not resolve: {unresolved_sub}", + ) + checker.truth( + "surface.subpackage_classes_are_bound", + isinstance(MACE_FROM_SUBPACKAGE, type) + and isinstance(MACEPotential, type) + and isinstance(CheckpointRemap, type) + and isinstance(MACEMatpesSpec, type) + and isinstance(MACEOMolSpec, type), + "`from molzoo.mace import MACE, MACEPotential, CheckpointRemap, " + "MACEMatpesSpec, MACEOMolSpec` did not all bind classes", + ) + + # ac-001's gate is the subset form; the equality is what actually holds, + # because both `__dir__`s return `sorted(__all__)`. Pinning the stronger + # one keeps `dir()` and `import *` from drifting apart in either direction. + checker.truth( + "surface.ac001_subpackage_all_within_dir", + set(molzoo.mace.__all__) <= set(dir(molzoo.mace)), + f"dir(molzoo.mace) is missing {sorted(set(molzoo.mace.__all__) - set(dir(molzoo.mace)))}", + ) + checker.truth( + "surface.top_level_dir_is_sorted_all", + dir(molzoo) == sorted(molzoo.__all__), + f"dir(molzoo) is {dir(molzoo)}, expected {sorted(molzoo.__all__)}", + ) + checker.truth( + "surface.subpackage_dir_is_sorted_all", + dir(molzoo.mace) == sorted(molzoo.mace.__all__), + f"dir(molzoo.mace) is {dir(molzoo.mace)}, expected {sorted(molzoo.mace.__all__)}", + ) + + for module in (molzoo, molzoo.mace): + raised = False + try: + getattr(module, UNKNOWN_NAME) + except AttributeError: + raised = True + checker.truth( + f"surface.{module.__name__}_table_is_closed", + raised, + f"{module.__name__}.{UNKNOWN_NAME} did not raise AttributeError -- a " + "typo must fail as a missing attribute, not as an ImportError from " + "inside the package", + ) + + +def check_alias_identity(checker: Checker) -> None: + """The top-level names are the sub-package's objects, not copies of them. + + "Thin alias" is only meaningful as an identity claim: a second definition + that happened to behave the same would satisfy every other check in this + file and still be the duplication 06-wire exists to remove. + + Args: + checker: Failure collector. + """ + print("\nAlias identity (top level is molzoo.mace, not a second definition)") + for name, top_level in ( + ("MACE", MACE), + ("MACESpec", MACESpec), + ("MACEMatpes", MACEMatpes), + ("MACEOMol", MACEOMol), + ("load_matpes_state_dict", load_matpes_state_dict), + ("load_omol_state_dict", load_omol_state_dict), + ): + checker.truth( + f"alias.{name}", + top_level is getattr(molzoo.mace, name), + f"molzoo.{name} is not molzoo.mace.{name}", + ) + checker.truth( + "alias.MACE_binding_matches_attribute", + MACE is MACE_FROM_SUBPACKAGE is LAZILY_RESOLVED_MACE, + "`from molzoo import MACE`, `from molzoo.mace import MACE` and " + "`molzoo.MACE` gave different objects", + ) + checker.truth( + "alias.variants_are_potentials", + issubclass(MACEMatpes, MACEPotential) and issubclass(MACEOMol, MACEPotential), + "the named variants no longer specialise MACEPotential", + ) + checker.truth( + "alias.variant_specs_share_the_base", + issubclass(MACEMatpesSpec, MACESpec) and issubclass(MACEOMolSpec, MACESpec), + "MACESpec is no longer the shared foundation-variant configuration base", + ) + + +def build_variants() -> tuple[MACEMatpes, MACEOMol]: + """The two tiny foundation models, built through the keyword surface. + + `atomic_energies` is passed as a `torch.Tensor`, the way + `scripts/matpes_port/run_nve.py:152` passes it, even though the spec holds a + plain list -- the keyword adapters have to accept both. + + Returns: + `(matpes, omol)`, fp64 on CPU. + """ + energies = torch.tensor(ATOMIC_ENERGIES) + torch.manual_seed(SEED) + matpes = MACEMatpes( + atomic_numbers=list(ATOMIC_NUMBERS), atomic_energies=energies, **TINY_MATPES_KWARGS + ) + torch.manual_seed(SEED) + omol = MACEOMol( + atomic_numbers=list(ATOMIC_NUMBERS), atomic_energies=energies, **TINY_OMOL_KWARGS + ) + return matpes, omol + + +def build_from_specs() -> tuple[MACEPotential, MACEPotential]: + """The same two models, built the way new code is told to build them. + + Returns: + `(matpes, omol)` as plain `MACEPotential`s over the two variant specs. + """ + torch.manual_seed(SEED) + matpes = MACEPotential( + MACEMatpesSpec( + atomic_numbers=ATOMIC_NUMBERS, atomic_energies=ATOMIC_ENERGIES, **TINY_MATPES_KWARGS + ) + ) + torch.manual_seed(SEED) + omol = MACEPotential( + MACEOMolSpec( + atomic_numbers=ATOMIC_NUMBERS, + atomic_energies=ATOMIC_ENERGIES, + use_fallback=False, # MACEOMol takes no such keyword; the spec defaults + **TINY_OMOL_KWARGS, + ) + ) + return matpes, omol + + +def check_parameter_inventory( + checker: Checker, name: str, model: torch.nn.Module, tensors: int, elements: int +) -> None: + """One model's parameter tensor count and element total, against the goldens. + + Args: + checker: Failure collector. + name: Label for the printed rows. + model: The constructed variant. + tensors: Golden `len(list(model.parameters()))`. + elements: Golden `sum(p.numel() for p in model.parameters())`. + """ + parameters = list(model.parameters()) + checker.count(f"{name}.param_tensors", len(parameters), tensors) + checker.count(f"{name}.param_elements", sum(p.numel() for p in parameters), elements) + + +def check_construction_paths( + checker: Checker, name: str, keyword: torch.nn.Module, spec: torch.nn.Module +) -> None: + """The keyword adapter and the spec form build the same module graph. + + Not a golden -- a relation, so it needs no captured literal and cannot rot. + It is the assertion that `variants.py` really is an adapter: if it re-derived + a hyper-parameter, an irreps string or an `E0` table of its own (which its + docstring forbids), the two inventories would part company here. + + Args: + checker: Failure collector. + name: Label for the printed rows. + keyword: Model from the keyword constructor. + spec: Model from `MACEPotential(spec)`. + """ + from_keyword = [(key, tuple(value.shape)) for key, value in keyword.named_parameters()] + from_spec = [(key, tuple(value.shape)) for key, value in spec.named_parameters()] + differing = [ + f"{left} != {right}" for left, right in zip(from_keyword, from_spec) if left != right + ] + checker.truth( + f"{name}.keyword_and_spec_agree", + from_keyword == from_spec, + f"{len(differing)} entr(y/ies) differ: {differing[:5]}; " + f"keyword has {len(from_keyword)} tensors, spec {len(from_spec)}", + ) + + +def main() -> int: + """Check the import surface, the lazy policy and the parameter goldens.""" + checker = Checker() + check_lazy_policy(checker) + check_import_surface(checker) + check_alias_identity(checker) + + print("\nParameter inventory (tiny variants, fp64 CPU, vs pre-cutover goldens)") + keyword_matpes, keyword_omol = build_variants() + check_parameter_inventory( + checker, "matpes", keyword_matpes, MATPES_PARAM_TENSORS, MATPES_PARAM_ELEMENTS + ) + check_parameter_inventory( + checker, "omol", keyword_omol, OMOL_PARAM_TENSORS, OMOL_PARAM_ELEMENTS + ) + + print("\nConstruction paths (keyword adapter vs MACEPotential(spec))") + spec_matpes, spec_omol = build_from_specs() + check_construction_paths(checker, "matpes", keyword_matpes, spec_matpes) + check_construction_paths(checker, "omol", keyword_omol, spec_omol) + check_parameter_inventory( + checker, "matpes.spec", spec_matpes, MATPES_PARAM_TENSORS, MATPES_PARAM_ELEMENTS + ) + check_parameter_inventory( + checker, "omol.spec", spec_omol, OMOL_PARAM_TENSORS, OMOL_PARAM_ELEMENTS + ) + + if checker.failures: + print("\nFAILED — the molzoo public wiring no longer matches the 06-wire contract:") + for failure in checker.failures: + print(f" {failure}") + return 1 + print("\nOK") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/regressions/mace-subpackage-restructure-07-cleanup.py b/regressions/mace-subpackage-restructure-07-cleanup.py new file mode 100644 index 0000000..9c97f27 --- /dev/null +++ b/regressions/mace-subpackage-restructure-07-cleanup.py @@ -0,0 +1,369 @@ +"""Refactor invariance of the public MACE energy core (chain step 07-cleanup). + +Chain step 07 closes three tails: the two in-tree consumers +(`scripts/matpes_port/run_nve.py`, `benchmarks/bench_mace_matpes.py`) stop +reaching into the private core `model._compute_energy` and call the public +`energy_core` landed in step 04; `molpot.composition.EnergyForceModel` is +deleted; and `molpot.derivation.protocol.ensure_graphs` learns a `num_graphs` +argument so the `graphs` namespace it builds carries the schema-conforming +`batch_size=[B]`. The physics is supposed to be untouched — this file is the +standalone check of that claim (spec +`.claude/specs/mace-subpackage-restructure-07-cleanup.md`). + +Scenario (public API only), two sections: + +1. **Energy and forces through the public core.** A hard-coded 6-atom water + dimer is pushed through `MACEMatpes.energy_core` inside a literal copy of + `run_nve._matpes_energy_forces`' closure shape — position leaf with + `requires_grad_(True)`, energy under `torch.enable_grad()`, and + `autograd_forces_from_energy` **outside** any compile region (no `Compiler` + here: `Compiler(cuda_graphs=True)` needs a GPU, and what is being guarded is + the numbers, not the capture). The positional argument order + `(pos, Z, edge_index, batch, num_graphs, shifts)` is the part of the + re-point that could silently rot, so it is reproduced verbatim. Total + energy, `max|F|` and `F[0]` are compared against goldens captured through + the *private* core the re-point replaced. +2. **`ensure_graphs` schema.** The one-line public-surface change of this step: + `ensure_graphs(TensorDict(batch_size=[]), num_graphs=3)["graphs"].batch_size` + must be `torch.Size([3])`, the `"graphs": batch_size=[B]` shape CLAUDE.md + specifies and `src/molzoo/pinet/potential.py` already reads `[0]` from. + +Random weights under a fixed seed, and an `E0` table (`[-1.0, -8.0]` eV/atom) +that is literal rather than physical: what is pinned is *refactor invariance*, +not accuracy against a reference implementation. The chain's parity against +upstream MACE is a separate, already-recorded verdict (`src/molzoo/specs/ +mace_matpes.md` §7.1) and cannot be re-run in-tree — mace-torch and e3nn are +not importable here by policy. + +Goldens +------- + oracle : molnex itself. No third-party oracle (no mace-torch, no e3nn, + no ASE), no network, no subprocess. + path : captured on the PRIVATE core `model._compute_energy` — the path + task 3 of this spec re-points away from — so the literals below + predate the change they are guarding. In the same capture the + public `energy_core` was run on the identical inputs and came + out bit-equal (`|dE| = 0.0`, `max|dF| = 0.0`, + `torch.equal` True for both), which is expected rather than + lucky: `MACEMatpes._compute_energy.__func__ is + MACEPotential.energy_core` was also asserted True at capture + time. The private name is deliberately *not* touched by this + script — public API only. + commit : e825a51 (e825a515342a5f5f41b70d6443bab5f5ee957299), chain tip 06 + "refactor(molzoo): cut over to the mace subpackage; retire flat + modules", `src/` clean. + command : PYTHONPATH=src python capture_goldens_07.py + torch : 2.12.1+cpu (tensordict 0.13.0) + date : 2026-08-09 + device : CPU, float64 (`molix.config.set_precision("fp64")` before any + construction), `use_fallback=True` — no + `cuequivariance-ops-torch` wheel on this host, so the pure-torch + cuEquivariance path is the only one available. + observed : repeat runs at a fixed thread count are bit-identical, and the + energy is bit-identical at every thread count tried. The + *forces* are not: re-measured on this host at + `OMP_NUM_THREADS` 1 / 2 / 4 / 8, `max|F|` moves by up to + 3.2e-12 eV·A^-1 against the literals below (worst case + `OMP_NUM_THREADS=1`; 2, 4 and 8 agree with each other to + 4.5e-13). The goldens are the default-thread values and + reproduce exactly there. So: do **not** tighten this file to + `torch.equal`. The spec's 1e-9 band is ~300x above that + reduction-order jitter and still ~2000 ULP below anything a + real behaviour change would produce, which is exactly the + separation it was chosen for. + +Run: + PYTHONPATH=src python regressions/mace-subpackage-restructure-07-cleanup.py +""" + +from __future__ import annotations + +import sys +import textwrap +import traceback +from collections.abc import Callable + +import torch + +from molix import config + +# Must precede every construction: the layers bake `config.ftype` in at +# __init__ time, so switching precision afterwards silently leaves an fp32 +# model behind and the goldens below stop meaning anything. +config.set_precision("fp64") + +from tensordict import TensorDict # noqa: E402 + +from molpot.derivation.force import autograd_forces_from_energy # noqa: E402 +from molpot.derivation.protocol import ensure_graphs # noqa: E402 +from molzoo.mace import MACEMatpes # noqa: E402 + +# --------------------------------------------------------------------------- +# The system. A 6-atom non-periodic water dimer (A), O-O 3.2 A, so that both +# intra- and inter-molecular edges fall inside the 4.0 A cutoff. Literal +# coordinates: a golden whose geometry is generated is not a golden. +# --------------------------------------------------------------------------- + +#: Cartesian positions in Angstrom, in `ATOMIC_NUMBERS` order. +POSITIONS: list[list[float]] = [ + [0.00000, 0.00000, 0.00000], # O + [0.95720, 0.00000, 0.00000], # H + [-0.23999, 0.92663, 0.00000], # H + [3.20000, 0.10000, 0.20000], # O + [3.80000, 0.85000, -0.10000], # H + [3.60000, -0.70000, 0.60000], # H +] + +#: Per-atom atomic numbers. +ATOMIC_NUMBERS: list[int] = [8, 1, 1, 8, 1, 1] + +#: Number of graphs `B` in the batch: one molecule-pair, so `batch = zeros(N)`. +NUM_GRAPHS = 1 + +# --------------------------------------------------------------------------- +# The model. Small MACE-MatPES built through the public keyword constructor. +# --------------------------------------------------------------------------- + +#: Element table (z-table), strictly ascending as `torch.searchsorted` needs. +Z_TABLE: list[int] = [1, 8] + +#: Reference energies `E0` in eV/atom, in `Z_TABLE` order. Arbitrary but +#: literal, and non-zero so the isolated-atom reference path contributes to the +#: total instead of being a silent no-op. +ATOMIC_ENERGIES: list[float] = [-1.0, -8.0] + +#: Cutoff radius in Angstrom; also the neighbour-list radius below. +R_MAX = 4.0 + +# --------------------------------------------------------------------------- +# Goldens. See the module docstring for provenance. +# --------------------------------------------------------------------------- + +#: Edge count of the `r < R_MAX` ordered-pair graph on `POSITIONS` (structural, +#: so exact): every pair of the 6 atoms except the four O-H/H-H pairs that +#: straddle the two molecules at more than 4 A. +GOLDEN_NUM_EDGES = 26 + +#: Total energy in eV (sum over the `(B,)` per-graph energies). +GOLDEN_ENERGY = -4014.3426025127233 + +#: Largest force component magnitude in eV/A. +GOLDEN_MAX_FORCE = 2833.787986314126 + +#: Force on atom 0 (the first O) in eV/A. +GOLDEN_FORCE_ATOM0: list[float] = [ + -2833.787986314126, + -2349.1750081421415, + -548.5925878141118, +] + +#: Refactor-invariance tolerance on energies, eV (spec "Domain basis"). +ENERGY_TOL = 1e-9 + +#: Refactor-invariance tolerance on forces, eV/A (spec "Domain basis"). +FORCE_TOL = 1e-9 + + +class Checker: + """Collects every deviation so one run reports all failures, not the first.""" + + def __init__(self) -> None: + self.failures: list[str] = [] + + def count(self, name: str, got: int, want: int) -> None: + """Assert an integer golden exactly -- no tolerance applies to a count.""" + ok = got == want + if not ok: + self.failures.append(f"{name}: got {got}, want {want}") + print(f" {name:<30} {got:>24} {'ok' if ok else 'FAILED'}") + + def close(self, name: str, got: float, want: float, tol: float, unit: str) -> None: + """Assert a float golden within the spec's refactor-invariance band.""" + delta = abs(got - want) + ok = delta <= tol + if not ok: + self.failures.append( + f"{name}: got {got!r}, want {want!r} ({unit}); |delta|={delta:.3e} > {tol:g}" + ) + print(f" {name:<30} {got!r:>24} |d|={delta:.3e} {unit:<9} {'ok' if ok else 'FAILED'}") + + def truth(self, name: str, holds: bool, message: str) -> None: + """Assert a boolean contract (a shape, a `batch_size`, ...).""" + if not holds: + self.failures.append(f"{name}: {message}") + print(f" {name:<30} {'ok' if holds else 'FAILED':>24}") + + +def build_edges(pos: torch.Tensor, r_max: float) -> torch.Tensor: + """Full bidirectional neighbour graph as `(E, 2)`, `[:,0]`=source, `[:,1]`=target. + + Non-periodic, so no `shifts` accompany it. `torch.cdist` over all ordered + pairs is O(N^2) and fine at N=6 -- the production path uses + `molix.md.NeighborList`, which this file deliberately does not pull + in: it would add a moving part between the goldens and the core under test. + + Args: + pos: Positions `(N, 3)` in Angstrom. + r_max: Cutoff radius in Angstrom. + + Returns: + Edge index `(E, 2)`, self-pairs excluded. + """ + dist = torch.cdist(pos, pos) + mask = (dist < r_max) & ~torch.eye(pos.shape[0], dtype=torch.bool) + return torch.nonzero(mask, as_tuple=False) + + +def build_model() -> MACEMatpes: + """Small MACE-MatPES with seeded random weights, through the public kwargs. + + The unpassed spec defaults are part of the golden as much as the passed + keywords are; at capture time they resolved to `num_polynomial_cutoff=5`, + `scale=1.0`, `shift=0.0`, `max_hidden_l=1`, `radial_mlp=[64, 64, 64]`, + `interaction='density'`, `readout='per_layer'`, + `distance_transform='agnesi'`, `pair_repulsion='zbl'`, + `conditioning='none'`, for a total of 35792 parameters. + + Returns: + The model in `eval()` mode. + """ + torch.manual_seed(0) + return MACEMatpes( + atomic_numbers=Z_TABLE, + atomic_energies=ATOMIC_ENERGIES, + r_max=R_MAX, + num_bessel=8, + l_max=2, + num_features=16, + num_interactions=2, + correlation=3, + mlp_dim=8, + use_fallback=True, + ).eval() + + +def energy_forces( + core: Callable[ + [torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, int, torch.Tensor | None], + torch.Tensor, + ], + pos: torch.Tensor, + Z: torch.Tensor, + edge_index: torch.Tensor, + batch: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """`run_nve._matpes_energy_forces`' closure shape, minus `Compiler`. + + Reproduced rather than imported: the script under `scripts/` is not a + public API, and the point is to pin the *shape* of the call the re-point + left behind -- leaf `requires_grad_`, energy under `torch.enable_grad()`, + `autograd.grad` outside any compile region, and the positional order + `(pos, Z, edge_index, batch, num_graphs, shifts)` with `shifts=None` for a + non-periodic cell. + + Args: + core: The public `energy_core`, per-graph energy `(B,)` out. + pos: Positions `(N, 3)`; detached and re-leafed here. + Z: Atomic numbers `(N,)`. + edge_index: Edge index `(E, 2)`. + batch: Graph membership `(N,)`. + + Returns: + Per-graph energies `(B,)` in eV and forces `(N, 3)` in eV/A, detached. + """ + leaf = pos.detach().requires_grad_(True) + with torch.enable_grad(): + energy = core(leaf, Z, edge_index, batch, NUM_GRAPHS, None) + forces = autograd_forces_from_energy(energy, leaf) + return energy.detach(), forces.detach() + + +def check_energy_core(checker: Checker) -> None: + """Section 1 -- the public core reproduces the pre-re-point goldens. + + Args: + checker: Failure collector. + """ + print("Section 1 - public energy_core vs pre-re-point goldens (fp64, CPU, fallback)") + model = build_model() + pos = torch.tensor(POSITIONS, dtype=config.ftype) + Z = torch.tensor(ATOMIC_NUMBERS, dtype=torch.long) + edge_index = build_edges(pos, R_MAX) + batch = torch.zeros(Z.shape[0], dtype=torch.long) + + checker.count("n_atoms", int(pos.shape[0]), len(ATOMIC_NUMBERS)) + checker.count("n_edges", int(edge_index.shape[0]), GOLDEN_NUM_EDGES) + + energy, forces = energy_forces(model.energy_core, pos, Z, edge_index, batch) + + checker.truth( + "energy.shape", + tuple(energy.shape) == (NUM_GRAPHS,), + f"energy_core returned {tuple(energy.shape)}, want ({NUM_GRAPHS},) per-graph energies", + ) + checker.truth( + "forces.shape", + tuple(forces.shape) == (len(ATOMIC_NUMBERS), 3), + f"got {tuple(forces.shape)}, want ({len(ATOMIC_NUMBERS)}, 3)", + ) + checker.close("E_total", float(energy.sum()), GOLDEN_ENERGY, ENERGY_TOL, "eV") + checker.close("max|F|", float(forces.abs().max()), GOLDEN_MAX_FORCE, FORCE_TOL, "eV/A") + for index, (axis, want) in enumerate(zip("xyz", GOLDEN_FORCE_ATOM0)): + checker.close(f"F[0].{axis}", float(forces[0][index]), want, FORCE_TOL, "eV/A") + + +def check_ensure_graphs(checker: Checker) -> None: + """Section 2 -- `ensure_graphs` builds a schema-conforming `graphs` namespace. + + Args: + checker: Failure collector. + """ + print("\nSection 2 - ensure_graphs(num_graphs=3) graphs schema") + batch = ensure_graphs(TensorDict(batch_size=[]), num_graphs=3) + got = batch["graphs"].batch_size + checker.truth( + "graphs.batch_size", + got == torch.Size([3]), + f"got {got}, want torch.Size([3]) -- the CLAUDE.md graphs schema is batch_size=[B]; " + "a consumer reading batch['graphs'].batch_size[0] raises IndexError on torch.Size([])", + ) + + +def main() -> int: + """Run both sections and report a single PASS/FAIL verdict. + + A section that *raises* is a section that failed: the signatures this file + pins (`energy_core`'s six positional parameters, `ensure_graphs`' + `num_graphs` keyword) drift by `TypeError` rather than by a wrong number, + so an escaping exception would otherwise kill the run before the + `RESULT:` line the caller reads. The traceback is printed, then folded + into the verdict. + + Returns: + `0` on PASS, `1` on FAIL. + """ + print("molnex regression - mace-subpackage-restructure-07-cleanup") + print(f"torch={torch.__version__} device=cpu dtype={config.ftype} use_fallback=True\n") + + checker = Checker() + for section in (check_energy_core, check_ensure_graphs): + try: + section(checker) + except Exception: + checker.failures.append( + f"{section.__name__} raised:\n" + f"{textwrap.indent(traceback.format_exc().rstrip(), ' ')}" + ) + + if checker.failures: + print("\nThe 07-cleanup contract is broken:") + for failure in checker.failures: + print(f" {failure}") + print("RESULT: FAIL") + return 1 + print("\nRESULT: PASS") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/regressions/md-neighborlist-skin-01-guard.py b/regressions/md-neighborlist-skin-01-guard.py new file mode 100644 index 0000000..4a4e95a --- /dev/null +++ b/regressions/md-neighborlist-skin-01-guard.py @@ -0,0 +1,455 @@ +"""Public-API scenario for the `NeighborList` cutoff bound. + +Spec: `md-neighborlist-skin-01-guard`. + +The claim this file pins, in one screenful, for the golden triclinic cell + + cell = [[10, 0, 0], + [ 6, 8, 0], + [ 0, 0, 10]] V = |det| = 800 A^3 + +is that the admissible cutoff is bounded by **half the smallest perpendicular +cell width**, not by half the shortest row norm: + + ||a_2 x a_3|| = 100, ||a_3 x a_1|| = 100, ||a_1 x a_2|| = 80 (A^2) + w = (V/100, V/100, V/80) = (8.000, 8.000, 10.000) A + bound = min_i w_i / 2 = 4.000 A <- what the constructor must enforce + min_i ||a_i|| / 2 = 5.000 A <- what the old row-norm guard allowed + +Everything between 4.000 A and 5.000 A is the bug: the old guard admitted those +cutoffs, and the kernel's *sequential* minimum-image reduction (subtract +round(dz/c_zz)*a_3, then round(dy/b_yy)*a_2, then round(dx/a_xx)*a_1) then +returns a displacement longer than the true minimum image, so pairs that are +inside the cutoff are silently dropped — wrong energy, wrong forces, no error. + +Part 1 — the bound. `cutoff=5.0` (the exact value the old guard admitted) and +`cutoff=4.5` both raise `ValueError`, and the message names the measured bound +`4.000 A` and the minimum width `8.000 A`, so the *number* is pinned and not +merely the failure. `cutoff=4.0` — exactly the bound — constructs, and +`cutoff=4.0001` does not; that brackets the bound to a ten-thousandth of an +Angstrom through the public error message alone. `cutoff=3.9` constructs and +actually builds a non-empty edge set, so "accepted" means "built", not "did not +raise". + +Part 2 — completeness at the admitted cutoff. Twelve atoms, written below as +literal fractional coordinates and mapped to Cartesian by `frac @ cell`, are +handed to `NeighborList(cutoff=3.9)`. The live half-pair set it +produces is compared against a brute-force reference computed **in this file**: +for every i < j, minimise ||r_j - r_i + n . cell|| over all 27 shifts +n in {-1, 0, 1}^3. Same pairs, same distances, nothing missed. Since +3.9 A < w_min/2 = 4.0 A, at most one periodic image of a pair can lie inside +the cutoff, so that minimum *is* the complete answer — the brute force is an +analytic oracle, not a third-party one. Eight of the fourteen reference pairs +are cross-boundary (non-zero shift n), which is what makes this a periodic test +rather than an open-boundary one. + +The neighbour list is symmetry-expanded (`E = 2 * n_pairs`), so `edge_index` +rows are deduplicated to `(min, max)` half pairs here; that every pair appears +exactly twice is asserted, since the deduplication would otherwise hide a +missing reverse edge. Distances are recovered from the public buffers as +`|| pos[target] - pos[source] + shifts ||`, the identity `shifts` is defined by. + +That this layout is a real counterexample and not just a passing one was +checked out of band while writing the file: with the guard forced back to the +old row-norm bound, the same twelve atoms at `cutoff=5.0` give 41 reference +half pairs and only 39 reported ones — pairs `(1, 8)` at 4.769 A and `(3, 10)` +at 4.827 A are silently dropped. That measurement is **not** asserted here: it +needs a private helper monkeypatched, and this file stays on the public API. +Part 1 refusing 5.0 A is the public-API form of the same claim. + +What is deliberately **not** pinned: `capacity` and the dead-edge tail padding +(fixed-capacity buffering is a different concern from the min-image bound), and +anything timed. + +Goldens +------- + capture command : PYTHONPATH=src python regressions/md-neighborlist-skin-01-guard.py + The one measured literal is N_HALF_PAIRS = 14, captured + from the in-script brute-force reference (not from the + neighbour list) while writing this file. Every other + golden is arithmetic over the cell above: V = 800, the + three face areas 100 / 100 / 80, w = (8, 8, 10), the + bound 4.000. + commit : 786d2b7 (786d2b74370371b039bae0fb6412a8f48466e28c), with + the `md-neighborlist-skin-01-guard` working tree on top + (the perpendicular-width guard in + `src/molix/md/neighbors.py` is new there). + torch : 2.12.1+cpu (python 3.14.5) + date : 2026-08-09 + device / dtype : CPU, float64 throughout — the cell and the positions are + float64 literals, so the neighbour list's buffers are + float64 too. float64 is what makes the 1e-9 A distance + tolerance meaningful; the observed agreement is exact. + oracle : none. No third-party package (torch is the repo's own + core dependency), no network, no subprocess, no RNG, no + wall-clock value, no filesystem access. + tolerance : 1e-9 A on distances (float64 "position" band, 1e-8, with a + decade of slack unused: the observed maximum deviation is + 0.0, bit-for-bit). Exact equality on pair sets and counts. + +Run: + PYTHONPATH=src python regressions/md-neighborlist-skin-01-guard.py +""" + +from __future__ import annotations + +import itertools +import sys + +import torch + +from molix.md import NeighborList + +# --------------------------------------------------------------------------- +# The golden triclinic cell. Rows are the cell vectors a_1, a_2, a_3 (A). +# +# Lower triangular, which is the form the kernel's reduction assumes. a_2 is +# sheared by 6 A along x, and that shear is the whole story: it costs the cell +# 2 A of perpendicular width along x (w_1 = 8, not 10) while leaving ||a_1|| = +# ||a_3|| = 10 and ||a_2|| = 10 — every row norm is 10, so the row-norm guard +# sees a 10 A cube that is not there. +# --------------------------------------------------------------------------- + +CELL = torch.tensor( + [ + [10.0, 0.0, 0.0], + [6.0, 8.0, 0.0], + [0.0, 0.0, 10.0], + ], + dtype=torch.float64, +) + +#: `V / max_i ||a_j x a_k||` = 800 / 100. Not read from the code under test — +#: the private helper is never imported here; this is the paper value. +MIN_WIDTH = 8.0 + +#: The contract: `cutoff <= min_i w_i / 2`. Formatted as `4.000 A` by the +#: constructor's error message, which is the only public window onto it. +BOUND = 4.0 + +#: What the superseded row-norm guard admitted: `min_i ||a_i|| / 2` = 10 / 2. +#: Every cutoff in (4.000, 5.000] used to be accepted and is now refused. +OLD_ROW_NORM_BOUND = 5.0 + +#: Substrings the `ValueError` must carry, so the numbers survive a reword. +BOUND_TEXT = "4.000 A" +MIN_WIDTH_TEXT = "8.000 A" + +# --------------------------------------------------------------------------- +# Twelve atoms as literal fractional coordinates (no RNG, no fixture). +# +# Cartesian is `frac @ cell`, i.e. r = (10*f1 + 6*f2, 8*f2, 10*f3). The layout +# is three loose layers along b with a deliberate spread along c, chosen so that +# +# * fourteen half pairs sit inside the 3.9 A cutoff and eight of them only +# through a periodic image, and +# * no pair separation lands near the cutoff — the closest miss is 4.139 A and +# the closest hit 3.582 A, a 0.239 A margin either side of 3.9, so the pair +# *set* cannot flip on an arithmetic reordering (the tolerance that matters +# for set membership is ~1e-1, not ~1e-9). +# +# Closest approach overall is 3.036 A, so nothing is unphysically overlapped. +# --------------------------------------------------------------------------- + +FRACTIONAL = torch.tensor( + [ + [0.00, 0.06, 0.10], + [0.30, 0.12, 0.15], + [0.46, 0.07, 0.85], + [0.76, 0.15, 0.55], + [0.06, 0.40, 0.35], + [0.36, 0.45, 0.60], + [0.61, 0.38, 0.05], + [0.85, 0.42, 0.81], + [0.05, 0.70, 0.20], + [0.33, 0.75, 0.91], + [0.65, 0.68, 0.45], + [0.82, 0.72, 0.70], + ], + dtype=torch.float64, +) + +#: The admitted cutoff for part 2: below the 4.000 A bound, above the old +#: guard's threshold for nothing — it is simply a legal cutoff at which the +#: minimum-image reduction is guaranteed complete. +CUTOFF = 3.9 + +#: **Captured once** from the in-script brute-force reference below (not from +#: the neighbour list), at the commit in the header. Fourteen of the 66 half +#: pairs are within 3.9 A; the neighbour list is symmetry-expanded, so this is +#: 28 live edges. +N_HALF_PAIRS = 14 + +#: Of those fourteen, the number reachable only across a periodic boundary +#: (minimising shift n != 0). Same capture. Pinned separately because a +#: neighbour list that silently lost periodicity would still find the other six. +N_CROSS_BOUNDARY_PAIRS = 8 + +#: Position-band tolerance in Angstrom for float64 (CLAUDE tester contract: +#: 1e-8 numerical, 1e-12 exact). Observed deviation is 0.0. +DIST_TOL = 1e-9 + +#: All 27 periodic images n in {-1, 0, 1}^3. Enough because the search radius +#: 3.9 A is below every half-width (4.0, 4.0, 5.0 A), so no second-shell image +#: can be the minimiser. +IMAGE_SHIFTS: tuple[tuple[int, ...], ...] = tuple(itertools.product((-1, 0, 1), repeat=3)) + + +# --------------------------------------------------------------------------- +# Checking +# --------------------------------------------------------------------------- + + +class Checker: + """Collects every deviation so one run reports all failures, not the first.""" + + def __init__(self) -> None: + self.failures: list[str] = [] + + def _row(self, name: str, got: object, ok: bool) -> None: + print(f" {name:<40} {got!s:<32} {'ok' if ok else 'FAILED'}") + + def exact(self, name: str, got: object, want: object) -> None: + """Assert a count, a pair set or a string — no tolerance applies.""" + ok = got == want + if not ok: + self.failures.append(f"{name}: got {got!r}, want {want!r}") + self._row(name, got, ok) + + def within(self, name: str, got: float, tol: float) -> None: + """Assert a measured deviation is at most *tol* (Angstrom).""" + ok = got <= tol + if not ok: + self.failures.append(f"{name}: deviation {got!r} A exceeds {tol!r} A") + self._row(name, f"{got:.3e} A", ok) + + def truth(self, name: str, holds: bool, message: str) -> None: + """Assert a boolean contract (a cutoff was refused, a list was built, ...).""" + if not holds: + self.failures.append(f"{name}: {message}") + self._row(name, holds, holds) + + +def positions() -> torch.Tensor: + """Cartesian coordinates ``(12, 3)`` in Angstrom for the literal layout. + + Returns: + ``FRACTIONAL @ CELL`` in float64 — fractional coordinates are the + readable form, Cartesian is what the public API takes. + """ + return FRACTIONAL @ CELL + + +def refuses(cutoff: float) -> str | None: + """Construct at *cutoff* and report the refusal message, if any. + + Args: + cutoff: Model cutoff ``r_cut`` in Angstrom. + + Returns: + The ``ValueError`` message if the constructor refused, else ``None``. + """ + try: + NeighborList(cell=CELL, cutoff=cutoff, positions=positions()) + except ValueError as error: + return str(error) + return None + + +def brute_force_pairs(pos: torch.Tensor) -> dict[tuple[int, int], float]: + """Minimum-image half pairs within :data:`CUTOFF`, by exhaustive search. + + For every ``i < j`` the true minimum-image separation is + ``min_n ||r_j - r_i + n . cell||`` over the 27 shifts ``n in {-1,0,1}^3``. + No shortcut, no sequential reduction — this is the reference the kernel's + reduction is supposed to reproduce. + + Args: + pos: Cartesian positions ``(N, 3)`` in Angstrom, float64. + + Returns: + ``{(i, j): distance}`` in Angstrom for every half pair whose minimum + image lies within :data:`CUTOFF`, with ``i < j``. + """ + images = torch.tensor(IMAGE_SHIFTS, dtype=pos.dtype) @ CELL # (27, 3) + inside: dict[tuple[int, int], float] = {} + n_atoms = int(pos.shape[0]) + for i in range(n_atoms): + for j in range(i + 1, n_atoms): + distance = float(torch.linalg.norm(pos[j] - pos[i] + images, dim=-1).min()) + if distance <= CUTOFF: + inside[(i, j)] = distance + return inside + + +def brute_force_cross_boundary(pos: torch.Tensor) -> int: + """How many reference pairs are reachable only through a periodic image. + + Args: + pos: Cartesian positions ``(N, 3)`` in Angstrom, float64. + + Returns: + The count of in-cutoff half pairs whose minimising shift is non-zero. + """ + shifts = torch.tensor(IMAGE_SHIFTS, dtype=pos.dtype) + images = shifts @ CELL + count = 0 + n_atoms = int(pos.shape[0]) + for i in range(n_atoms): + for j in range(i + 1, n_atoms): + norms = torch.linalg.norm(pos[j] - pos[i] + images, dim=-1) + best = int(norms.argmin()) + if float(norms[best]) <= CUTOFF and bool(shifts[best].abs().sum() > 0): + count += 1 + return count + + +def live_half_pairs( + neighbors: NeighborList, pos: torch.Tensor +) -> dict[tuple[int, int], list[float]]: + """Deduplicate the live, symmetry-expanded edge buffer into half pairs. + + The list is built with ``symmetry=True``, so each pair occupies two rows + (``i -> j`` and ``j -> i``); both are collected under the ``(min, max)`` key + so the caller can assert each pair appears exactly twice. Distances come + from the public buffers via the identity ``shifts`` is defined by: + ``edge_diff = pos[target] - pos[source] + shifts``. + + Args: + neighbors: A built list; only its public ``edge_index`` / ``shifts`` / + ``num_edges`` members are read. + pos: The Cartesian positions the list was built at, ``(N, 3)``. + + Returns: + ``{(i, j): [distance, ...]}`` in Angstrom over the live rows + ``[0, num_edges)``. + """ + edge_index = neighbors.edge_index[: neighbors.num_edges] + source, target = edge_index[:, 0], edge_index[:, 1] + edge_diff = pos[target] - pos[source] + neighbors.shifts[: neighbors.num_edges] + edge_dist = torch.linalg.norm(edge_diff, dim=-1) + pairs: dict[tuple[int, int], list[float]] = {} + for row in range(int(edge_index.shape[0])): + i, j = int(source[row]), int(target[row]) + pairs.setdefault((min(i, j), max(i, j)), []).append(float(edge_dist[row])) + return pairs + + +def check_bound(checker: Checker) -> None: + """Part 1 — the constructor bounds the cutoff at 4.000 A, not 5.000 A. + + Args: + checker: Failure collector. + """ + print("Bound (cell rows 10 / 6,8 / 10 A; V = 800 A^3; w = (8.000, 8.000, 10.000) A)") + + old_guard_message = refuses(OLD_ROW_NORM_BOUND) + checker.truth( + "bound.rejects_5.0", + old_guard_message is not None, + "cutoff=5.0 A was accepted — that is half the shortest row norm, not " + "half the smallest perpendicular width; the min-image reduction drops " + "pairs inside the cutoff there", + ) + checker.truth( + "bound.rejects_4.5", + refuses(4.5) is not None, + "cutoff=4.5 A was accepted — above the 4.000 A bound, below the old " + "row-norm guard's 5.000 A, i.e. squarely in the silently-wrong window", + ) + checker.truth( + "bound.message_names_4.000_A", + old_guard_message is not None and BOUND_TEXT in old_guard_message, + f"the refusal does not contain {BOUND_TEXT!r}, so the measured bound is " + f"not observable through the public API: {old_guard_message!r}", + ) + checker.truth( + "bound.message_names_8.000_A", + old_guard_message is not None and MIN_WIDTH_TEXT in old_guard_message, + f"the refusal does not contain {MIN_WIDTH_TEXT!r} (min_i w_i = " + f"{MIN_WIDTH} A): {old_guard_message!r}", + ) + # Brackets the bound: accepted at exactly min_i w_i / 2, refused a + # ten-thousandth of an Angstrom above it. + checker.truth( + "bound.accepts_exactly_4.0", + refuses(BOUND) is None, + "cutoff=4.000 A was refused; the contract is `cutoff <= min_i w_i / 2`, " + "so the bound itself is admissible", + ) + checker.truth( + "bound.rejects_4.0001", + refuses(4.0001) is not None, + "cutoff=4.0001 A was accepted; the bound is 4.000 A exactly", + ) + + accepted = NeighborList(cell=CELL, cutoff=CUTOFF, positions=positions()) + checker.truth( + "bound.accepts_3.9_and_builds", + accepted.num_edges > 0, + f"cutoff={CUTOFF} A constructed but produced {accepted.num_edges} edges — " + "acceptance must mean 'actually built', not merely 'did not raise'", + ) + + +def check_completeness(checker: Checker) -> None: + """Part 2 — at the admitted cutoff the list equals the 27-image reference. + + Args: + checker: Failure collector. + """ + print(f"\nCompleteness at cutoff = {CUTOFF} A (12 atoms, brute force over 27 images)") + pos = positions() + reference = brute_force_pairs(pos) + checker.exact("reference.n_half_pairs", len(reference), N_HALF_PAIRS) + checker.exact( + "reference.n_cross_boundary", + brute_force_cross_boundary(pos), + N_CROSS_BOUNDARY_PAIRS, + ) + + neighbors = NeighborList(cell=CELL, cutoff=CUTOFF, positions=pos) + observed = live_half_pairs(neighbors, pos) + checker.exact("list.num_edges", neighbors.num_edges, 2 * N_HALF_PAIRS) + checker.exact("list.n_half_pairs", len(observed), N_HALF_PAIRS) + checker.truth( + "list.each_pair_bidirectional", + all(len(distances) == 2 for distances in observed.values()), + "a half pair does not occupy exactly two rows: " + f"{sorted(pair for pair, d in observed.items() if len(d) != 2)}", + ) + + missed = sorted(set(reference) - set(observed)) + spurious = sorted(set(observed) - set(reference)) + checker.exact("completeness.missed_pairs", tuple(missed), ()) + checker.exact("completeness.spurious_pairs", tuple(spurious), ()) + if missed: + print(" pairs inside the cutoff that the list never reported:") + for pair in missed: + print(f" {pair}: true minimum image {reference[pair]:.9f} A") + + shared = set(reference) & set(observed) + deviation = max( + (abs(distance - reference[pair]) for pair in shared for distance in observed[pair]), + default=0.0, + ) + checker.within("completeness.max_distance_deviation", deviation, DIST_TOL) + + +def main() -> int: + """Pin the cutoff bound and the minimum-image completeness beneath it.""" + checker = Checker() + check_bound(checker) + check_completeness(checker) + + if checker.failures: + print("\nFAILED — NeighborList no longer matches the golden bound / reference:") + for failure in checker.failures: + print(f" {failure}") + return 1 + print("\nOK") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/regressions/md-neighborlist-skin-02-prune.py b/regressions/md-neighborlist-skin-02-prune.py new file mode 100644 index 0000000..d2b613d --- /dev/null +++ b/regressions/md-neighborlist-skin-02-prune.py @@ -0,0 +1,305 @@ +"""Public-API scenario for the neighbour-graph homonym prune. + +Spec: `md-neighborlist-skin-02-prune`. + +Two dead neighbour-graph implementations are removed from the tree so that the +name `NeighborList` is unambiguous: + + molix.nn.locality.NeighborList an nn.Module wrapper over + molix.F.locality.get_neighbor_pairs; + zero call sites + molpot.graph.radius_graph the same kernel plus a cross-molecule + mask, with the pre-Edge-Convention sign + (pos_j - pos_i, unnegated); zero call + sites, no __init__.py, never a declared + surface + +Stage is `experimental`, so this is a **hard removal** — no alias, no +DeprecationWarning, no shim. This file is the public-API form of the two +claims that makes such a removal safe. + +Section 1 — the deleted surface stays deleted. `import molix.nn` and +`import molpot` still succeed (the prune is import-time-only: nothing in the +training loop, the MD driver or any encoder resolved either symbol), while +`molix.nn.NeighborList`, the `molix.nn.locality` module, the `molpot.graph` +package, the `molpot.graph.radius` module and a `molpot.radius_graph` +attribute are all absent. `find_spec("molpot.graph")` is checked rather than +an import: `src/molpot/graph/` has no `__init__.py`, so a *surviving directory* +— a leftover `__pycache__`, say — is still a live namespace package and makes +the spec non-None even though nothing imports. That leftover is precisely the +drift being caught. + +Section 2 — the capability was never lost, only the duplicate wrapper. The +surviving `molix.data.tasks.NeighborList` (module home +`molix.data.tasks.neighbor`) is run on a hard-coded 3-atom chain and must still +produce the documented bidirectional edge list. + +The goldens are arithmetic on the page, not a captured measurement. For atoms +at x = 0.0, 1.0, 2.0 A on a line and `cutoff=1.5 A`: + + d(0,1) = 1.0 A <= 1.5 in + d(1,2) = 1.0 A <= 1.5 in + d(0,2) = 2.0 A > 1.5 out + +so two half pairs survive, and `symmetry=True` (the default, the full +bidirectional list every aggregating model assumes) doubles them: + + E = 2 x 2 = 4 -> edge_index (4, 2), edge_dist == 1.0 in all four entries + +Only the shape and the distances are pinned. Row order and the source/target +orientation within a row are **not** asserted: both are kernel-internal and the +Edge Convention (`edge_index[:,0]` = source, `edge_diff = pos[target] - +pos[source]`) is already pinned by the unit suite; re-asserting an ordering +here would make this file fail on a legal kernel change. + +Goldens +------- + capture command : PYTHONPATH=src python regressions/md-neighborlist-skin-02-prune.py + Nothing was captured from a run. Every literal below is + derived on the page from the chain geometry above: the + four names of `molix.nn.__all__`, E = 4, d = 1.0 A. + commit : 0111076 (0111076aba240e3e3b7b326f33341bf601fcd6ed), with + the `md-neighborlist-skin-02-prune` working tree on top + (at 0111076 itself Section 1 fails by construction — the + two modules are still present; that is the RED this file + was written against). + torch : 2.12.1+cpu (python 3.14.5) + date : 2026-08-09 + device / dtype : CPU, float64 — the positions are float64 literals, so the + neighbour list's distances come back float64 too. + oracle : none. No third-party package (torch is the repo's own + core dependency), no network, no subprocess, no RNG, no + wall-clock value, no filesystem access. + tolerance : 1e-6 A on the distances (spec-mandated; the float64 + "position" band is 1e-8 and the observed deviation is + 0.0, so six decades of slack are unused). Exact equality + on every import-surface assertion and on the edge shape. + +Run: + PYTHONPATH=src python regressions/md-neighborlist-skin-02-prune.py +""" + +from __future__ import annotations + +import importlib +import importlib.util +import sys + +import torch + +import molix.nn +import molpot +from molix.data.tasks import NeighborList + +# --------------------------------------------------------------------------- +# Section 1 goldens — the surviving public surface of `molix.nn`. +# +# Hard-coded as a list, not a set: the order pins the alphabetization rule +# (`.claude/notes/notes.md:243`) that the deletion's rewrite has to honour. +# --------------------------------------------------------------------------- + +SURVIVING_NN_EXPORTS = ["BatchAggregation", "KeyedMLP", "KeyedMLPSpec", "ScatterSum"] + +#: Fully-qualified names that must no longer resolve to anything importable. +DELETED_MODULES = ("molix.nn.locality", "molpot.graph") + +# --------------------------------------------------------------------------- +# Section 2 goldens — the 3-atom chain. +# +# Three atoms on the x axis one Angstrom apart; y = z = 0 throughout, so every +# separation is a difference of the x column and can be read off by eye. +# --------------------------------------------------------------------------- + +POSITIONS = torch.tensor( + [ + [0.0, 0.0, 0.0], + [1.0, 0.0, 0.0], + [2.0, 0.0, 0.0], + ], + dtype=torch.float64, +) + +#: Between the nearest-neighbour separation (1.0 A, in) and the end-to-end one +#: (2.0 A, out), with a 0.5 A margin either side — the pair *set* cannot flip +#: on arithmetic noise. +CUTOFF = 1.5 + +#: Two half pairs — (0,1) and (1,2) — each expanded to a forward and a reverse +#: edge by `symmetry=True`. +EXPECTED_EDGE_SHAPE = (4, 2) + +#: Every surviving pair is a nearest-neighbour pair, so all four rows carry the +#: same distance. +EXPECTED_EDGE_DIST = 1.0 + +#: Spec-mandated distance tolerance in Angstrom. +DIST_TOL = 1e-6 + + +# --------------------------------------------------------------------------- +# Checking +# --------------------------------------------------------------------------- + + +class Checker: + """Collects every deviation so one run reports all failures, not the first.""" + + def __init__(self) -> None: + self.failures: list[str] = [] + + def _row(self, name: str, got: object, ok: bool) -> None: + print(f" {name:<42} {got!s:<34} {'ok' if ok else 'FAILED'}") + + def exact(self, name: str, got: object, want: object) -> None: + """Assert a name list, a shape or a flag — no tolerance applies.""" + ok = got == want + if not ok: + self.failures.append(f"{name}: got {got!r}, want {want!r}") + self._row(name, got, ok) + + def within(self, name: str, got: float, tol: float) -> None: + """Assert a measured deviation is at most *tol* (Angstrom).""" + ok = got <= tol + if not ok: + self.failures.append(f"{name}: deviation {got!r} A exceeds {tol!r} A") + self._row(name, f"{got:.3e} A", ok) + + def truth(self, name: str, holds: bool, message: str) -> None: + """Assert a boolean contract (a module is gone, an import raised, ...).""" + if not holds: + self.failures.append(f"{name}: {message}") + self._row(name, holds, holds) + + +def import_error_of(module: str) -> str | None: + """Import *module* and report the ``ModuleNotFoundError``, if any. + + Args: + module: Fully-qualified module name, e.g. ``"molpot.graph.radius"``. + + Returns: + The exception message if the import raised :class:`ModuleNotFoundError`, + else ``None`` — the module is still importable, which is the drift. + """ + try: + importlib.import_module(module) + except ModuleNotFoundError as error: + return str(error) + return None + + +def check_deleted_surface(checker: Checker) -> None: + """Section 1 — neither dead neighbour-graph symbol resolves any more. + + Args: + checker: Failure collector. + """ + print("Deleted surface (molix.nn.locality.NeighborList, molpot.graph.radius_graph)") + + # The prune is import-time-only; the two packages must still come up. + checker.truth( + "surface.molix_nn_imports", + molix.nn is not None, + "`import molix.nn` did not yield a module — the deletion was supposed " + "to drop one re-export line, not break the package", + ) + checker.truth( + "surface.molpot_imports", + molpot is not None, + "`import molpot` did not yield a module — `molpot.graph` was never a " + "declared surface, so removing it must be invisible here", + ) + + checker.exact("surface.molix_nn_all", list(molix.nn.__all__), SURVIVING_NN_EXPORTS) + checker.truth( + "surface.no_neighborlist_attribute", + not hasattr(molix.nn, "NeighborList"), + "`molix.nn.NeighborList` still resolves; the name must stay free for " + "molix.data.tasks.neighbor.NeighborList / molix.md.NeighborList", + ) + checker.truth( + "surface.no_neighborlist_export", + "NeighborList" not in molix.nn.__all__, + "`NeighborList` is back in `molix.nn.__all__`", + ) + + for module in DELETED_MODULES: + # `find_spec`, not `import`: `src/molpot/graph/` has no `__init__.py`, + # so a surviving directory (e.g. a stale `__pycache__`) is a live + # namespace package with a non-None spec and no import error. + spec = importlib.util.find_spec(module) + checker.truth( + f"surface.{module}_has_no_spec", + spec is None, + f"`{module}` still resolves to {spec!r} — the module file or its " + "directory (stale `__pycache__` included) survived the deletion", + ) + + checker.truth( + "surface.molpot_graph_radius_unimportable", + import_error_of("molpot.graph.radius") is not None, + "`import molpot.graph.radius` succeeded; the pre-Edge-Convention " + "radius_graph builder is still in the tree", + ) + checker.truth( + "surface.molpot_has_no_radius_graph", + not hasattr(molpot, "radius_graph"), + "`molpot.radius_graph` resolves; the deleted free function was never " + "exported at package level and must not appear now", + ) + + +def check_surviving_neighbor_list(checker: Checker) -> None: + """Section 2 — the live pipeline neighbour list still builds the chain graph. + + Args: + checker: Failure collector. + """ + print(f"\nSurviving neighbour list (3-atom chain at x = 0, 1, 2 A; cutoff = {CUTOFF} A)") + + # `molix.data.tasks.NeighborList` is the package re-export of + # `molix.data.tasks.neighbor.NeighborList` — the same class object. + task = NeighborList(cutoff=CUTOFF, symmetry=True) + sample = task.execute({"pos": POSITIONS}) + + edge_index = sample["edge_index"] + edge_dist = sample["edge_dist"] + + checker.exact("chain.edge_index_shape", tuple(edge_index.shape), EXPECTED_EDGE_SHAPE) + checker.truth( + "chain.edge_dist_all_one", + bool( + torch.allclose( + edge_dist, + torch.full_like(edge_dist, EXPECTED_EDGE_DIST), + atol=DIST_TOL, + rtol=0.0, + ) + ), + f"edge distances {edge_dist.tolist()} are not all {EXPECTED_EDGE_DIST} A — " + "within a 1.5 A cutoff only the two nearest-neighbour pairs survive, and " + "both are exactly 1.0 A", + ) + deviation = ( + float((edge_dist - EXPECTED_EDGE_DIST).abs().max()) if edge_dist.numel() else float("inf") + ) + checker.within("chain.max_distance_deviation", deviation, DIST_TOL) + + +def main() -> int: + """Pin the two deletions and the surviving neighbour list's behaviour.""" + checker = Checker() + check_deleted_surface(checker) + check_surviving_neighbor_list(checker) + + if checker.failures: + print("\nFAILED — the pruned neighbour-graph surface is back, or the survivor drifted:") + for failure in checker.failures: + print(f" {failure}") + return 1 + print("\nOK") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/regressions/md-neighborlist-skin-03-rename.py b/regressions/md-neighborlist-skin-03-rename.py new file mode 100644 index 0000000..4d1dc18 --- /dev/null +++ b/regressions/md-neighborlist-skin-03-rename.py @@ -0,0 +1,391 @@ +"""Public-API scenario for the `PeriodicNeighborList` -> `NeighborList` rename. + +Spec: `md-neighborlist-skin-03-rename`. + +The rename is declared **behaviour-neutral**: identifiers, docstrings, `__all__` +entries and prose move; no logic, signature, default or numerical result does. +A rename that quietly changed a number would be indistinguishable from a +successful one if the only evidence were "the suite is green after the suite was +edited" — every unit-test call site had its identifier substituted in the same +commit. This file is the independent witness: it was written against the +*renamed* symbol only, it touches nothing private, and every literal it asserts +is arithmetic on the page rather than a value read off a run. + +Section 1 — the surface. `from molix.md import NeighborList` resolves, and +`molix.md.PeriodicNeighborList` does not exist. Stage is `experimental` and the +spec forbids a back-compat alias, so the old name must be *gone*, not shimmed: +a `PeriodicNeighborList = NeighborList` line, a module `__getattr__` fallback or +a stale re-export would all be caught here. `"PeriodicNeighborList" not in +molix.md.__all__` is checked separately from the attribute, because a name can +survive in `__all__` (breaking `from molix.md import *`) while `hasattr` is +already False, and vice versa. + +Section 2 — the numbers the rename must not have moved. A 3x3x3 simple-cubic +lattice, spacing 3.0 A, in a cubic 9.0 A cell, at `cutoff=3.5 A`: + + nearest-neighbour separation 3.000 A <= 3.5 in (6 per site) + face-diagonal separation sqrt2*3 = 4.243 A > 3.5 out + body-diagonal separation sqrt3*3 = 5.196 A > 3.5 out + +so each of the 27 sites has exactly its 6 axis neighbours (+-x, +-y, +-z) inside +the cutoff — distinct sites, since 3 lattice points per axis means the +3.0 A and +-3.0 A neighbours are different atoms, not the same image twice. Under the repo +convention (`symmetry=True`, full bidirectional list): + + num_edges = 27 * 6 = 162 + capacity = ceil(1.35 * 162) = ceil(218.7) = 219 (default capacity_factor) + edge_index (219, 2) shifts (219, 3) + +The 0.5 A margin on either side of the cutoff (3.0 in, 4.243 out) means no pair +can flip class on floating-point noise, which is what makes `162` a hard integer +golden and not a tolerance question. `cutoff = 3.5 A <= 4.5 A = 9.0/2` also +clears the constructor's half-perpendicular-width bound, so the minimum-image +reduction is complete here. + +Section 3 — the dead-edge tail. Rows `[num_edges, capacity)` (57 of them) are +padding and must be *inert*: source and target both atom 0, displaced by +`DEAD_EDGE_CUTOFF_FACTOR (10.0) * cutoff (3.5) = 35.0 A`, far outside every +cutoff envelope, so the padding contributes exactly zero energy and — since +`pos[0] - pos[0]` cancels — exactly zero force. 35.0 A is asserted directly on +the norm of the padding rows. + +Section 4 — rebuild under a rigid translation. Shifting every atom by +`(1.234, 0, 0)` A preserves every pairwise displacement exactly, so the periodic +neighbour set is unchanged: `num_edges` back to 162, buffer shapes untouched +(the whole point of fixed capacity — the force path stays CUDA-graph +capturable), and `rebuild_count == 1` after exactly one `rebuild` call. + +Drift policy +------------ +A disagreement between this file and the runtime is a **DEFECT REPORT, never a +golden edit**. Every literal here is derived analytically from the lattice above +and from the two documented constants (`capacity_factor=1.35`, +`DEAD_EDGE_CUTOFF_FACTOR=10.0`); none was captured from a run. If `num_edges` +comes back as anything but 162, the neighbour path lost or gained pairs, and the +correct response is to open a defect against `src/molix/md/neighbors.py` — not to +adjust the number below to whatever the run printed. + +Goldens +------- + capture command : PYTHONPATH=src python regressions/md-neighborlist-skin-03-rename.py + Nothing was captured from a run. Every literal is + arithmetic on the page: 27*6 = 162, ceil(1.35*162) = 219, + 10.0*3.5 = 35.0, sqrt(2)*3 = 4.243 > 3.5. + commit : d7b0ea2 (d7b0ea2bceb2740fc14b21dd5901a72bbb2c0228), with + the `md-neighborlist-skin-03-rename` working tree on top + (at d7b0ea2 itself Section 1 fails by construction — the + class is still `PeriodicNeighborList` there; that is the + RED this file was written against). + torch : 2.12.1+cpu (python 3.14.5) + date : 2026-08-09 + device / dtype : CPU, float64 — the cell and the positions are float64 + literals, so the neighbour list's `shifts` buffer is + float64 too. `edge_index` is int64 regardless. + oracle : none. No third-party package (torch is the repo's own + core dependency), no network, no subprocess, no RNG, no + wall-clock value, no filesystem access. + tolerance : 1e-12 A on the dead-edge shift norm (float64 "position" + exact band; 35.0 = 10.0 * 3.5 is exact in binary and the + observed deviation is 0.0). Exact integer equality on + edge counts, capacity, shapes, rebuild count and every + surface assertion. + +Run: + PYTHONPATH=src python regressions/md-neighborlist-skin-03-rename.py +""" + +from __future__ import annotations + +import itertools +import sys + +import torch + +import molix.md +from molix.md import NeighborList + +# --------------------------------------------------------------------------- +# Section 1 goldens — the renamed public surface. +# --------------------------------------------------------------------------- + +#: The name that must have disappeared entirely: no attribute, no `__all__` +#: entry, no alias. `stage: experimental` buys the hard rename. +RETIRED_NAME = "PeriodicNeighborList" + +#: The name that must resolve on `molix.md` after the rename. +CURRENT_NAME = "NeighborList" + +# --------------------------------------------------------------------------- +# Section 2 goldens — the 3x3x3 simple-cubic lattice. +# +# Three lattice points per axis at 0 / 3 / 6 A in a 9 A cube: the spacing is +# uniform under periodicity (6 -> 9 == 0), so every site is equivalent and the +# edge count is exactly 27 * (number of in-cutoff neighbours per site). +# --------------------------------------------------------------------------- + +#: Lattice spacing in Angstrom. +SPACING = 3.0 + +#: Cubic cell edge in Angstrom — three spacings, so the lattice tiles exactly. +BOX = 9.0 + +#: Between the nearest-neighbour separation (3.0 A, in) and the face diagonal +#: (sqrt2 * 3.0 = 4.243 A, out), and at most half the 9.0 A perpendicular width +#: (4.5 A), so minimum image is complete. +CUTOFF = 3.5 + +CELL = torch.tensor( + [ + [BOX, 0.0, 0.0], + [0.0, BOX, 0.0], + [0.0, 0.0, BOX], + ], + dtype=torch.float64, +) + +#: 27 sites x 6 axis neighbours each, doubled-counted as directed edges by the +#: full-bidirectional convention (`symmetry=True`, the MD list's fixed setting). +EXPECTED_NUM_EDGES = 162 + +#: ceil(1.35 * 162) = ceil(218.7), with the constructor's default +#: `capacity_factor=1.35`. +EXPECTED_CAPACITY = 219 + +EXPECTED_EDGE_INDEX_SHAPE = (EXPECTED_CAPACITY, 2) +EXPECTED_SHIFTS_SHAPE = (EXPECTED_CAPACITY, 3) + +# --------------------------------------------------------------------------- +# Section 3 goldens — the inert padding tail. +# --------------------------------------------------------------------------- + +#: `DEAD_EDGE_CUTOFF_FACTOR (10.0) * CUTOFF (3.5)`, written out rather than +#: imported: this file pins the *number* a dead edge carries, so importing the +#: constant would make the assertion tautological. +EXPECTED_DEAD_SHIFT_NORM = 35.0 + +#: Both endpoints of a dead edge are atom 0, so `pos[0] - pos[0]` cancels and no +#: spurious force reaches it; the *set* of indices in the padding rows is +#: therefore the single value 0. +EXPECTED_DEAD_ENDPOINTS = (0,) + +#: float64 "position" exact band; 35.0 is representable exactly. +NORM_TOL = 1e-12 + +# --------------------------------------------------------------------------- +# Section 4 goldens — rigid translation. +# --------------------------------------------------------------------------- + +#: An arbitrary non-lattice offset: rigid, so every pairwise displacement (and +#: hence every minimum image) is preserved exactly, but not a symmetry of the +#: lattice, so a rebuild that silently reused stale state would still be caught. +TRANSLATION = torch.tensor([1.234, 0.0, 0.0], dtype=torch.float64) + +#: One `rebuild` call; construction itself does not count as a rebuild. +EXPECTED_REBUILD_COUNT = 1 + + +def simple_cubic_lattice() -> torch.Tensor: + """Build the 3x3x3 simple-cubic lattice used by every section. + + Returns: + Positions ``(27, 3)`` in Angstrom, float64: the Cartesian product of + ``{0.0, 3.0, 6.0}`` with itself three times, in lexicographic order. + """ + coordinates = (0.0 * SPACING, 1.0 * SPACING, 2.0 * SPACING) + return torch.tensor( + list(itertools.product(coordinates, repeat=3)), + dtype=torch.float64, + ) + + +# --------------------------------------------------------------------------- +# Checking +# --------------------------------------------------------------------------- + + +class Checker: + """Collects every deviation so one run reports all failures, not the first.""" + + def __init__(self) -> None: + self.failures: list[str] = [] + + def _row(self, name: str, got: object, ok: bool) -> None: + print(f" {name:<42} {got!s:<24} {'ok' if ok else 'FAILED'}") + + def exact(self, name: str, got: object, want: object) -> None: + """Assert an integer, a shape or a name — no tolerance applies.""" + ok = got == want + if not ok: + self.failures.append(f"{name}: got {got!r}, want {want!r}") + self._row(name, got, ok) + + def within(self, name: str, got: float, tol: float) -> None: + """Assert a measured deviation is at most *tol* (Angstrom).""" + ok = got <= tol + if not ok: + self.failures.append(f"{name}: deviation {got!r} A exceeds {tol!r} A") + self._row(name, f"{got:.3e} A", ok) + + def truth(self, name: str, holds: bool, message: str) -> None: + """Assert a boolean contract (a name is gone, a shape held, ...).""" + if not holds: + self.failures.append(f"{name}: {message}") + self._row(name, holds, holds) + + +def check_renamed_surface(checker: Checker) -> None: + """Section 1 — the new name resolves and the old one is gone without a shim. + + Args: + checker: Failure collector. + """ + print("Renamed surface (molix.md.NeighborList)") + + checker.truth( + "surface.neighborlist_attribute", + hasattr(molix.md, CURRENT_NAME), + "`molix.md.NeighborList` does not resolve; the MD buffer owner must be " + "reachable under the bare name freed by link 02", + ) + checker.truth( + "surface.neighborlist_export", + CURRENT_NAME in molix.md.__all__, + "`NeighborList` is missing from `molix.md.__all__`; the attribute alone " + "is not the declared surface", + ) + checker.truth( + "surface.no_periodic_attribute", + not hasattr(molix.md, RETIRED_NAME), + "`molix.md.PeriodicNeighborList` still resolves — a back-compat alias, " + "a module `__getattr__` fallback or a stale re-export survived; the " + "spec forbids all three at `stage: experimental`", + ) + checker.truth( + "surface.no_periodic_export", + RETIRED_NAME not in molix.md.__all__, + "`PeriodicNeighborList` is still in `molix.md.__all__`; " + "`from molix.md import *` would resurrect the retired name", + ) + checker.exact("surface.class_name", NeighborList.__name__, CURRENT_NAME) + + +def check_lattice_buffers(checker: Checker, neighbor_list: NeighborList) -> None: + """Section 2 — edge count, capacity and buffer shapes on the lattice. + + Args: + checker: Failure collector. + neighbor_list: List built on the 3x3x3 lattice at ``cutoff=3.5 A``. + """ + print(f"\nLattice buffers (3x3x3 sc, {SPACING} A spacing, {BOX} A cell, cutoff {CUTOFF} A)") + + checker.exact("lattice.num_edges", neighbor_list.num_edges, EXPECTED_NUM_EDGES) + checker.exact("lattice.capacity", neighbor_list.capacity, EXPECTED_CAPACITY) + checker.exact( + "lattice.edge_index_shape", + tuple(neighbor_list.edge_index.shape), + EXPECTED_EDGE_INDEX_SHAPE, + ) + checker.exact( + "lattice.shifts_shape", + tuple(neighbor_list.shifts.shape), + EXPECTED_SHIFTS_SHAPE, + ) + + +def check_dead_edge_tail(checker: Checker, neighbor_list: NeighborList) -> None: + """Section 3 — the padding rows carry an inert 35.0 A self-loop. + + Args: + checker: Failure collector. + neighbor_list: List built on the 3x3x3 lattice at ``cutoff=3.5 A``. + """ + print(f"\nDead-edge tail (rows [{EXPECTED_NUM_EDGES}, {EXPECTED_CAPACITY}))") + + dead_shifts = neighbor_list.shifts[neighbor_list.num_edges :] + dead_edges = neighbor_list.edge_index[neighbor_list.num_edges :] + + checker.exact( + "dead.row_count", + int(dead_shifts.shape[0]), + EXPECTED_CAPACITY - EXPECTED_NUM_EDGES, + ) + + norms = torch.linalg.norm(dead_shifts, dim=-1) + deviation = ( + float((norms - EXPECTED_DEAD_SHIFT_NORM).abs().max()) if norms.numel() else float("inf") + ) + checker.within("dead.shift_norm_deviation", deviation, NORM_TOL) + checker.truth( + "dead.shift_norm_is_35A", + deviation <= NORM_TOL, + f"dead-edge shift norms {norms.unique().tolist()} A are not " + f"{EXPECTED_DEAD_SHIFT_NORM} A (= 10.0 x {CUTOFF} A); padding closer than " + "the cutoff would leak a spurious pair contribution into the energy", + ) + + endpoints = tuple(int(value) for value in dead_edges.unique().tolist()) + checker.truth( + "dead.endpoints_are_atom_zero", + endpoints == EXPECTED_DEAD_ENDPOINTS, + f"dead-edge endpoints {endpoints} are not all atom 0; a dead edge must " + "be a self-loop so `pos[0] - pos[0]` cancels and no force reaches it", + ) + + +def check_rigid_translation(checker: Checker, neighbor_list: NeighborList) -> None: + """Section 4 — a rigid shift leaves the neighbour set and the shapes alone. + + Args: + checker: Failure collector. + neighbor_list: List built on the 3x3x3 lattice; rebuilt in place here. + """ + print(f"\nRigid translation ({TRANSLATION.tolist()} A)") + + neighbor_list.rebuild(simple_cubic_lattice() + TRANSLATION) + + checker.exact("translated.num_edges", neighbor_list.num_edges, EXPECTED_NUM_EDGES) + checker.exact("translated.capacity", neighbor_list.capacity, EXPECTED_CAPACITY) + checker.exact( + "translated.edge_index_shape", + tuple(neighbor_list.edge_index.shape), + EXPECTED_EDGE_INDEX_SHAPE, + ) + checker.exact( + "translated.shifts_shape", + tuple(neighbor_list.shifts.shape), + EXPECTED_SHIFTS_SHAPE, + ) + checker.exact("translated.rebuild_count", neighbor_list.rebuild_count, EXPECTED_REBUILD_COUNT) + + +def main() -> int: + """Pin the renamed surface and the buffer numbers the rename must not move.""" + checker = Checker() + + check_renamed_surface(checker) + + neighbor_list = NeighborList( + cell=CELL, + cutoff=CUTOFF, + positions=simple_cubic_lattice(), + ) + check_lattice_buffers(checker, neighbor_list) + check_dead_edge_tail(checker, neighbor_list) + check_rigid_translation(checker, neighbor_list) + + if checker.failures: + print("\nFAILED — the rename was not behaviour-neutral, or the old name is back:") + for failure in checker.failures: + print(f" {failure}") + print( + "\nEvery golden above is analytic (27*6 = 162, ceil(1.35*162) = 219, " + "10.0*3.5 = 35.0). A disagreement is a DEFECT REPORT against the " + "neighbour path, never a reason to edit the literal." + ) + return 1 + print("\nOK") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/regressions/md-neighborlist-skin-04-policy.py b/regressions/md-neighborlist-skin-04-policy.py new file mode 100644 index 0000000..87a606d --- /dev/null +++ b/regressions/md-neighborlist-skin-04-policy.py @@ -0,0 +1,483 @@ +"""Public-API scenario for the Verlet skin + `neigh_modify` rebuild policy. + +Spec: `md-neighborlist-skin-04-policy`. + +`molix.md.NeighborList` now builds at `r_build = cutoff + skin` and decides once +per force evaluation, in `update(positions)`, whether the frozen list is still +valid — under the LAMMPS `neigh_modify every/delay/check` gate. The whole value +of that machinery is *when it rebuilds*: too eager and the skin bought nothing, +too lazy and pairs are silently missed (an `O(1)`, one-signed energy injection +per missed event, not an `O(dt^2)` discretisation artefact). A unit test can +pin one gate arm at a time; this file pins the **decision sequence** a user +actually observes, end to end, over three scripted schedules whose every literal +is arithmetic on this page. + +The shared configuration is a 64-atom simple-cubic lattice (4x4x4 sites, 3.0 A +spacing) in a 12.0 A cube, CPU float64. Nothing moves by itself: displacements +are applied by hand, so no integrator, no thermostat and no RNG stand between +the schedule and the counters. + +Section 1 — the crystallographic goldens (what the skin costs) +-------------------------------------------------------------- +Simple cubic at `a = 3.0 A` has exact neighbour shells: + + 6 neighbours at a = 3.000 A + 12 neighbours at sqrt2*a = 4.243 A + 8 neighbours at sqrt3*a = 5.196 A + 6 neighbours at 2*a = 6.000 A + +With `cutoff = 3.5 A, skin = 1.5 A`: + + r_build = 3.5 + 1.5 = 5.0 A + in-build shells: 6 + 12 = 18 (5.196 A is out; 4.243 A is in) + num_edges = 64 * 18 = 1152 (full bidirectional list, `symmetry=True`) + capacity = ceil(1.35 * 1152) = ceil(1555.2) = 1556 + +versus the bare cutoff `skin = 0.0`: + + r_build = 3.5 A; only the 6-shell is in + num_edges = 64 * 6 = 384 + capacity = ceil(1.35 * 384) = ceil(518.4) = 519 + +1152/384 = 3.0 is the `(1 + skin/r_cut)^3 = (1 + 1.5/3.5)^3 = 2.91` growth law +measured directly (the lattice is discrete, so the ratio lands on the next shell +rather than exactly on the continuum estimate). Four sites per axis means the +`+a` and `-a` neighbours are distinct atoms, so the shell multiplicities above +are not double-counted images. The `2*a = 6.0 A` shell is exactly the +minimum-image half-width and would be ambiguous — it sits well outside +`r_build = 5.0 A`, so it never enters. `r_build = 5.0 A <= 6.0 A = 12.0/2` +also clears the constructor's half-perpendicular-width guard. + +Section 2 — gating arithmetic (`every=2, delay=4, check=True`) +-------------------------------------------------------------- +Atom 0 is translated by `+0.1 A` along x per update; every other atom is frozen, +so the largest displacement since the last build is exactly `0.1 * ago` A. +Half-skin is `skin/2 = 0.75 A`. The gate is conjunctive — a rebuild is +*permitted* only at `ago >= delay` **and** `ago % every == 0`: + + ago: 1 2 3 4 5 6 7 8 + >= delay 4: . . . x x x x x + % every 2: . x . x . x . x + permitted: . . . x . x . x + displacement: 0.4 0.6 0.8 A + vs 0.75 A: no no YES + +so the first rebuild lands at the 8th update since the build, not the 4th +(displacement too small) and not the 7th (`ago` odd). `rebuild` resets `ago` to +0, so the pattern repeats verbatim: over 40 updates `update()` returns True at +exactly `{8, 16, 24, 32, 40}` — `rebuild_count == 5`. Every one of those fired +at `ago == 8`, never at the first permitted opportunity `ago == max(every, +delay) == 4`, so `ndanger == 0`: no rebuild was ever overdue. + +This one schedule falsifies all three ways the gate can be wrong. A disjunctive +gate (`or` instead of `and`) would fire at update 6. A gate that ignored the +displacement check would fire at update 4. A `>=` instead of the strict `>` on +the half-skin would not change these numbers (0.4/0.6/0.8 never equal 0.75) — +which is why the strict comparison is pinned in the unit tests instead, and why +this file pins the *schedule*. + +Section 3 — the `ndanger` alarm (`every=1, delay=0, check=True`) +---------------------------------------------------------------- +At the default gate `max(every, delay) == 1`, so *every* rebuild lands on the +first permitted opportunity and every rebuild is counted dangerous. Atom 0 is +jumped `1.0 A` along x and back, alternately, so the displacement since the last +build is `1.0 A > 0.75 A` on every single update: 5 updates give +`rebuild_count == 5` and `ndanger == 5`. + +That is the alarm behaving as designed, not a bug: `ndanger` says "a rebuild +fired at the earliest moment the gate allowed, so it may already have been +overdue on a step the gate skipped". With `every=1` no step is skipped, so the +counter carries no information there — pinned here precisely so the documented +degenerate reading stays visible and nobody "fixes" it into silence. + +Section 4 — cadence only (`every=5, delay=0, check=False`) +---------------------------------------------------------- +Nothing moves at all. With `check=False` the displacement criterion — and with +it the unwrapped-positions guard — is skipped entirely, so the list rebuilds on +raw cadence: updates `{5, 10, 15, 20}` out of 20, `rebuild_count == 4`, and +`ndanger == 0` because `ndanger` is only ever incremented inside the +displacement branch. Rebuilding four times on a configuration that never +changed is the visible price of `check=False`; under `check=True` the same run +would rebuild zero times. + +Drift policy +------------ +A disagreement between this file and the runtime is a **DEFECT REPORT, never a +golden edit**. Every literal below is derived on this page from the lattice +geometry, the two documented constants (`capacity_factor = 1.35`, shell +multiplicities of the simple-cubic lattice) and the LAMMPS gate as specified; +none was read off a run and then written down. If the fired-update sequence +comes back as anything but `{8, 16, 24, 32, 40}`, the gate changed semantics — +open a defect against `src/molix/md/neighbors.py` and fix the gate, do not +retune the tuple below to whatever the run printed. + +Goldens +------- + capture command : PYTHONPATH=src python regressions/md-neighborlist-skin-04-policy.py + Nothing was captured from a run. Every literal is + arithmetic on the page: 64*18 = 1152, 64*6 = 384, + ceil(1.35*1152) = 1556, ceil(1.35*384) = 519, + 3.5 + 1.5 = 5.0, and the three gate tables above. + commit : 58fb180 (58fb180faab256bd03083782fcda8e3af9c4a883), with + the `md-neighborlist-skin-04-policy` working tree on top + (at 58fb180 itself `NeighborList` has no `skin` / `update` + at all, so construction raises `TypeError` — that is the + RED this file was written against). + torch : 2.12.1+cpu (python 3.14.5) + date : 2026-08-09 + device / dtype : CPU, float64 — the cell and the positions are float64 + literals, so `shifts` and `_x_hold` are float64 too. + `edge_index` is int64 regardless. + oracle : none. No third-party package (torch is the repo's own + core dependency), no network, no subprocess, no RNG, no + wall-clock value, no filesystem access. LAMMPS is the + *specified* behaviour, not a runtime dependency: its gate + is transcribed into the tables above, never executed. + tolerance : 1e-12 A on `r_build` (float64 "position" exact band; both + 3.5 and 1.5 are exactly representable, so 5.0 is exact and + the observed deviation is 0.0). Exact integer equality on + every edge count, capacity, rebuild count, `ago`, + `ndanger` and fired-update sequence — they are counters, + and a counter has no tolerance. + +Run: + PYTHONPATH=src python regressions/md-neighborlist-skin-04-policy.py +""" + +from __future__ import annotations + +import itertools +import sys + +import torch + +from molix.md import NeighborList + +# --------------------------------------------------------------------------- +# The shared lattice. +# --------------------------------------------------------------------------- + +#: Lattice spacing `a` in Angstrom. +SPACING = 3.0 + +#: Sites per axis; 4 keeps the `+a` and `-a` neighbours distinct atoms. +N_SIDE = 4 + +#: Cubic cell edge in Angstrom — four spacings, so the lattice tiles exactly. +BOX = SPACING * N_SIDE + +#: Interaction cutoff in Angstrom: between the 3.0 A first shell (in) and the +#: 4.243 A second shell (out), so no pair can flip class on rounding. +CUTOFF = 3.5 + +#: Verlet skin in Angstrom. `r_build = 5.0 A` sits between the 4.243 A second +#: shell (in) and the 5.196 A third shell (out) — again a hard integer boundary, +#: not a tolerance question — and half-skin is a round 0.75 A. +SKIN = 1.5 + +CELL = torch.tensor( + [ + [BOX, 0.0, 0.0], + [0.0, BOX, 0.0], + [0.0, 0.0, BOX], + ], + dtype=torch.float64, +) + +# --------------------------------------------------------------------------- +# Section 1 goldens — what the skin costs. +# --------------------------------------------------------------------------- + +#: `CUTOFF + SKIN`, written out rather than recomputed from the two constants: +#: this file pins the derived radius, so deriving it here would be tautological. +EXPECTED_R_BUILD = 5.0 + +#: 64 sites x (6 at 3.0 A + 12 at 4.243 A) directed edges, both directions kept +#: (`symmetry=True`, the MD list's fixed setting). +EXPECTED_EDGES_AT_R_BUILD = 1152 + +#: 64 sites x 6 at 3.0 A — the same lattice with `skin = 0.0`. +EXPECTED_EDGES_AT_CUTOFF = 384 + +#: ceil(1.35 * 1152) = ceil(1555.2), with the default `capacity_factor = 1.35`. +EXPECTED_CAPACITY_AT_R_BUILD = 1556 + +#: ceil(1.35 * 384) = ceil(518.4). +EXPECTED_CAPACITY_AT_CUTOFF = 519 + +#: float64 "position" exact band; 3.5 + 1.5 = 5.0 is exact in binary. +LENGTH_TOL = 1e-12 + +# --------------------------------------------------------------------------- +# Section 2 goldens — `every=2, delay=4, check=True`. +# --------------------------------------------------------------------------- + +#: Per-update translation of atom 0 along x, in Angstrom. +STEP_DISPLACEMENT = 0.1 + +#: Updates driven in scenario 1. +SCENARIO_1_UPDATES = 40 + +SCENARIO_1_EVERY = 2 +SCENARIO_1_DELAY = 4 + +#: Permitted `ago` values are {4, 6, 8, ...}; the displacement `0.1 * ago` first +#: exceeds half-skin 0.75 A at `ago == 8`, and `rebuild` re-phases `ago` to 0. +SCENARIO_1_FIRED_UPDATES = (8, 16, 24, 32, 40) + +SCENARIO_1_REBUILD_COUNT = 5 + +#: Every rebuild fired at `ago == 8`, never at the first permitted opportunity +#: `max(every, delay) == 4`, so no rebuild was ever overdue. +SCENARIO_1_NDANGER = 0 + +#: The 40th update *is* a rebuild, so the clock is back at 0 when the loop ends. +SCENARIO_1_FINAL_AGO = 0 + +# --------------------------------------------------------------------------- +# Section 3 goldens — `every=1, delay=0, check=True`. +# --------------------------------------------------------------------------- + +#: Jump of atom 0 along x, in Angstrom; alternated with 0.0 so the displacement +#: *since the last build* is 1.0 A on every update, not a growing drift. +JUMP_DISPLACEMENT = 1.0 + +SCENARIO_2_UPDATES = 5 + +#: 1.0 A > half-skin 0.75 A every time, and `every=1` permits every step. +SCENARIO_2_FIRED_UPDATES = (1, 2, 3, 4, 5) + +SCENARIO_2_REBUILD_COUNT = 5 + +#: `max(every, delay) == max(1, 0) == 1 == ago` at every rebuild: the documented +#: degenerate reading of the alarm, where it merely counts rebuilds. +SCENARIO_2_NDANGER = 5 + +# --------------------------------------------------------------------------- +# Section 4 goldens — `every=5, delay=0, check=False`, nothing moves. +# --------------------------------------------------------------------------- + +SCENARIO_3_UPDATES = 20 +SCENARIO_3_EVERY = 5 + +#: Pure cadence: `ago % 5 == 0` and `check=False` short-circuits before the +#: displacement branch, so a motionless system still rebuilds four times. +SCENARIO_3_FIRED_UPDATES = (5, 10, 15, 20) + +SCENARIO_3_REBUILD_COUNT = 4 + +#: `ndanger` is only ever incremented inside the displacement branch, which +#: `check=False` skips. +SCENARIO_3_NDANGER = 0 + + +def simple_cubic_lattice() -> torch.Tensor: + """Build the 4x4x4 simple-cubic lattice every section runs on. + + Returns: + Positions ``(64, 3)`` in Angstrom, float64: the Cartesian product of + ``{0.0, 3.0, 6.0, 9.0}`` with itself three times, lexicographic order. + """ + coordinates = tuple(index * SPACING for index in range(N_SIDE)) + return torch.tensor( + list(itertools.product(coordinates, repeat=3)), + dtype=torch.float64, + ) + + +def displaced(base: torch.Tensor, offset: float) -> torch.Tensor: + """Return *base* with atom 0 translated by *offset* Angstrom along x. + + Args: + base: Reference positions ``(N, 3)`` in Angstrom. + offset: Translation of atom 0 along x, in Angstrom. + + Returns: + A fresh ``(N, 3)`` tensor; *base* is never mutated, so each schedule is + an absolute displacement from the lattice rather than an accumulation + of rounding. + """ + positions = base.clone() + positions[0, 0] += offset + return positions + + +# --------------------------------------------------------------------------- +# Checking +# --------------------------------------------------------------------------- + + +class Checker: + """Collects every deviation so one run reports all failures, not the first.""" + + def __init__(self) -> None: + self.failures: list[str] = [] + + def _row(self, name: str, got: object, ok: bool) -> None: + print(f" {name:<44} {got!s:<26} {'ok' if ok else 'FAILED'}") + + def exact(self, name: str, got: object, want: object) -> None: + """Assert a counter, a shape or an update sequence — no tolerance applies.""" + ok = got == want + if not ok: + self.failures.append(f"{name}: got {got!r}, want {want!r}") + self._row(name, got, ok) + + def within(self, name: str, got: float, tol: float) -> None: + """Assert a measured deviation is at most *tol* (Angstrom).""" + ok = got <= tol + if not ok: + self.failures.append(f"{name}: deviation {got!r} A exceeds {tol!r} A") + self._row(name, f"{got:.3e} A", ok) + + +def check_skin_geometry(checker: Checker) -> None: + """Section 1 — `r_build`, the edge counts it implies, and the fresh counters. + + Args: + checker: Failure collector. + """ + print(f"Skin geometry ({N_SIDE}x{N_SIDE}x{N_SIDE} sc, {SPACING} A spacing, {BOX} A cell)") + + lattice = simple_cubic_lattice() + skinned = NeighborList(cell=CELL, cutoff=CUTOFF, positions=lattice, skin=SKIN) + bare = NeighborList(cell=CELL, cutoff=CUTOFF, positions=lattice, skin=0.0) + + checker.within("skin.r_build_deviation", abs(skinned.r_build - EXPECTED_R_BUILD), LENGTH_TOL) + checker.exact("skin.cutoff_unchanged", skinned.cutoff, CUTOFF) + checker.exact("skin.skin", skinned.skin, SKIN) + checker.exact("skin.num_edges_at_r_build", skinned.num_edges, EXPECTED_EDGES_AT_R_BUILD) + checker.exact("skin.capacity_at_r_build", skinned.capacity, EXPECTED_CAPACITY_AT_R_BUILD) + + checker.exact("bare.num_edges_at_cutoff", bare.num_edges, EXPECTED_EDGES_AT_CUTOFF) + checker.exact("bare.capacity_at_cutoff", bare.capacity, EXPECTED_CAPACITY_AT_CUTOFF) + + checker.exact("fresh.ago", skinned.ago, 0) + checker.exact("fresh.rebuild_count", skinned.rebuild_count, 0) + checker.exact("fresh.ndanger", skinned.ndanger, 0) + + +def check_gate_schedule(checker: Checker) -> None: + """Section 2 — the `every=2, delay=4, check=True` decision sequence. + + Args: + checker: Failure collector. + """ + print( + f"\nGate schedule (every={SCENARIO_1_EVERY}, delay={SCENARIO_1_DELAY}, check=True, " + f"{STEP_DISPLACEMENT} A/update on atom 0)" + ) + + lattice = simple_cubic_lattice() + neighbor_list = NeighborList( + cell=CELL, + cutoff=CUTOFF, + positions=lattice, + skin=SKIN, + every=SCENARIO_1_EVERY, + delay=SCENARIO_1_DELAY, + check=True, + ) + + fired: list[int] = [] + for update in range(1, SCENARIO_1_UPDATES + 1): + if neighbor_list.update(displaced(lattice, STEP_DISPLACEMENT * update)): + fired.append(update) + + checker.exact("gate.fired_updates", tuple(fired), SCENARIO_1_FIRED_UPDATES) + checker.exact("gate.rebuild_count", neighbor_list.rebuild_count, SCENARIO_1_REBUILD_COUNT) + checker.exact("gate.ndanger", neighbor_list.ndanger, SCENARIO_1_NDANGER) + checker.exact("gate.final_ago", neighbor_list.ago, SCENARIO_1_FINAL_AGO) + + +def check_danger_alarm(checker: Checker) -> None: + """Section 3 — `ndanger` at the default gate, where every rebuild is dangerous. + + Args: + checker: Failure collector. + """ + print(f"\nDanger alarm (every=1, delay=0, check=True, {JUMP_DISPLACEMENT} A jump per update)") + + lattice = simple_cubic_lattice() + neighbor_list = NeighborList( + cell=CELL, + cutoff=CUTOFF, + positions=lattice, + skin=SKIN, + every=1, + delay=0, + check=True, + ) + + fired: list[int] = [] + for update in range(1, SCENARIO_2_UPDATES + 1): + # Alternate there-and-back so the displacement *since the last build* is + # 1.0 A every time, rather than a drift that would also pass a broken + # gate reading absolute position. + offset = JUMP_DISPLACEMENT if update % 2 else 0.0 + if neighbor_list.update(displaced(lattice, offset)): + fired.append(update) + + checker.exact("danger.fired_updates", tuple(fired), SCENARIO_2_FIRED_UPDATES) + checker.exact("danger.rebuild_count", neighbor_list.rebuild_count, SCENARIO_2_REBUILD_COUNT) + checker.exact("danger.ndanger", neighbor_list.ndanger, SCENARIO_2_NDANGER) + + +def check_cadence_only(checker: Checker) -> None: + """Section 4 — `check=False` rebuilds a motionless system on cadence alone. + + Args: + checker: Failure collector. + """ + print(f"\nCadence only (every={SCENARIO_3_EVERY}, delay=0, check=False, nothing moves)") + + lattice = simple_cubic_lattice() + neighbor_list = NeighborList( + cell=CELL, + cutoff=CUTOFF, + positions=lattice, + skin=SKIN, + every=SCENARIO_3_EVERY, + delay=0, + check=False, + ) + + fired: list[int] = [] + for update in range(1, SCENARIO_3_UPDATES + 1): + if neighbor_list.update(lattice): + fired.append(update) + + checker.exact("cadence.fired_updates", tuple(fired), SCENARIO_3_FIRED_UPDATES) + checker.exact("cadence.rebuild_count", neighbor_list.rebuild_count, SCENARIO_3_REBUILD_COUNT) + checker.exact("cadence.ndanger", neighbor_list.ndanger, SCENARIO_3_NDANGER) + checker.exact("cadence.num_edges_unchanged", neighbor_list.num_edges, EXPECTED_EDGES_AT_R_BUILD) + + +def main() -> int: + """Pin the rebuild-decision sequence the skin policy produces.""" + checker = Checker() + + check_skin_geometry(checker) + check_gate_schedule(checker) + check_danger_alarm(checker) + check_cadence_only(checker) + + if checker.failures: + print("\nFAILED — the rebuild policy no longer matches the LAMMPS gate:") + for failure in checker.failures: + print(f" {failure}") + print( + "\nEvery golden above is analytic (64*18 = 1152, 64*6 = 384, " + "ceil(1.35*1152) = 1556, and the three gate tables in the module " + "docstring). A disagreement is a DEFECT REPORT against " + "src/molix/md/neighbors.py, never a reason to edit the literal: a " + "gate that rebuilds late misses pairs and leaks NVE energy, and a " + "golden retuned to match it would hide exactly that." + ) + return 1 + print("\nOK") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/regressions/md-neighborlist-skin-05-bind.py b/regressions/md-neighborlist-skin-05-bind.py new file mode 100644 index 0000000..088f28d --- /dev/null +++ b/regressions/md-neighborlist-skin-05-bind.py @@ -0,0 +1,643 @@ +"""Public-API scenario for the TensorDict bind surface of `molix.md.NeighborList`. + +Spec: `md-neighborlist-skin-05-bind`. + +`nl.build(batch)` rebuilds the list at `batch["atoms", "pos"]` and writes +`batch["edges"]` as a `TensorDict` holding the live `edge_index` / `shifts` +buffers **by reference**; `nl.update(batch)` then drives the LAMMPS +`every/delay/check` policy from the same batch, dispatching on +`TensorDictBase` so the raw `(N, 3)` MD hot path and the batch path share one +door. The load-bearing property is the *tie*: because every rebuild is in +place, a consumer holding that batch sees the current neighbour set with no +re-binding and no shape change — which is exactly what keeps a compiled or +CUDA-graph-captured force path valid across a rebuild. + +A unit test can pin the identity `batch["edges", "edge_index"] is nl.edge_index` +at bind time. What it cannot show is that the tie *stays alive across a driven +run* and that what a consumer reads through the batch is still physically +right afterwards. That is this file: bind once, drive twenty steps through the +batch, then reconstruct the interatomic distances **out of the batch** and check +them against crystallography. A binding that ever copied on assignment would +sail through the first three sections and fail the fourth with a frozen +neighbour set — which is the failure mode worth a regression file. + +The configuration is the same 64-atom simple-cubic lattice link 04 uses (4x4x4 +sites, 3.0 A spacing) in a 12.0 A cube, CPU float64, +`NeighborList(cutoff=3.5, skin=1.5, every=1, delay=0, check=True)`. Nothing +moves by itself: positions are written by hand, so no integrator, no thermostat +and no RNG stand between the schedule and the counters. + +Section 1 — the lattice goldens (inherited, re-pinned at the bind) +------------------------------------------------------------------ +Simple cubic at `a = 3.0 A` has exact neighbour shells: + + 6 neighbours at a = 3.000 A + 12 neighbours at sqrt2*a = 4.243 A + 8 neighbours at sqrt3*a = 5.196 A + +With `cutoff = 3.5 A, skin = 1.5 A`: + + r_build = 3.5 + 1.5 = 5.0 A (the 5.196 A shell is out, 4.243 A is in) + num_edges = 64 * (6 + 12) = 1152 (full bidirectional list, `symmetry=True`) + capacity = ceil(1.35 * 1152) = ceil(1555.2) = 1556 + +`r_build = 5.0 A <= 6.0 A = 12.0/2` clears the constructor's +half-perpendicular-width guard. Four sites per axis keeps the `+a` and `-a` +neighbours distinct atoms, so the multiplicities are not double-counted images. + +Section 2 — what `build` binds +------------------------------- +`build` returns the *same* batch object (so `potential(nl.build(batch))` +composes with the repo's `forward(td) -> td` convention), and the two leaves it +writes are the list's own buffers, not copies: + + nl.build(batch) is batch + batch["edges", "edge_index"] is nl.edge_index + batch["edges", "shifts"] is nl.shifts + set(batch["edges"].keys()) == {"edge_index", "shifts"} + batch["edges"].batch_size == [1556] (capacity, not num_edges) + +`build` is a *binding* operation, not physics: `rebuild_count` stays 0 across it +(it also runs on every `.to()` re-sync, and counting it would make a dtype cast +look like a rebuild), while `ago` restarts at 0 because the buffers are fresh. + +Section 3 — driving the policy through the batch +------------------------------------------------- +Every update translates **all 64 atoms** rigidly by `+0.2 A` along x, written in +place into the batch's `("atoms", "pos")` leaf — the shape an integrator +produces. Half-skin is `skin/2 = 0.75 A`, and the criterion is strict `>`: + + ago: 1 2 3 4 + displacement: 0.2 0.4 0.6 0.8 A + vs 0.75 A: no no no YES + +`every=1, delay=0` permits every step, so the first rebuild lands at `ago == 4`; +`rebuild` re-phases `ago` to 0, so the pattern repeats verbatim. Over 20 +updates `update(batch)` returns True at exactly `{4, 8, 12, 16, 20}` — +`rebuild_count == 5`, and the 20th update *is* a rebuild so the run ends at +`ago == 0`. + +`ndanger == 0`: at this gate `max(every, delay) == 1`, but every rebuild fired at +`ago == 4`, never at that first permitted opportunity. This is the *informative* +corner of the alarm, and the complement of link 04's Section 3 — there +`every=1` with a 1.0 A jump made every rebuild fire at `ago == 1` and +`ndanger` merely counted rebuilds. Here the skin does its job and the counter +correctly reports "no rebuild was ever overdue". + +Section 4 — reading the physics back **out of the batch** +---------------------------------------------------------- +A rigid translation preserves every minimum-image displacement exactly, so the +neighbour set is invariant: `num_edges == 1152` still, after five rebuilds and a +total 4.0 A drift (the lattice is deliberately left *unwrapped* — x runs to +13.0 A in a 12.0 A cell — which the raw-displacement invariant requires). + +The check reconstructs distances the way a potential does, through the batch and +not through `nl`: + + r = || pos[target] - pos[source] + shift || over edge_index[:num_edges] + +with `pos`, `edge_index` and `shifts` all read from `batch`. Against the shell +table of Section 1 that must give, exactly: + + min r = 3.000 A (first shell) + max r = sqrt2*a = 4.2426406871 A (second shell) + count(r <= 3.5 A) = 64 * 6 = 384 (interaction cutoff) + count(3.5 < r <= 5.0 A) = 64 * 12 = 768 (the skin band) + count(r > 5.0 A) = 0 (nothing beyond r_build) + +Section 5 — dispatch equivalence at the public-API level +--------------------------------------------------------- +A second, identically constructed list is driven over the *same* schedule with +raw `update(pos)` tensors instead of `update(batch)`. Both end at +`rebuild_count == 5`, `ndanger == 0`, `num_edges == 1152`, and `torch.equal` on +the full `edge_index` / `shifts` buffers — dead padding tail included. One +policy, one state machine, two front doors. + +Drift policy +------------ +A disagreement between this file and the runtime is a **DEFECT REPORT, never a +golden edit**. Every literal below is derived on this page from the lattice +geometry, the documented `capacity_factor = 1.35`, and the LAMMPS gate as +specified; none was read off a run and then written down. If `max r` comes back +as anything but `sqrt2 * 3.0 A`, or `num_edges` drifts off 1152 after the run, +the bind went stale — open a defect against `src/molix/md/neighbors.py`, do not +retune the literal. A batch that silently *copied* on assignment would show up +here as exactly that kind of drift, and a golden edited to match it would hide a +frozen potential-energy surface. + +Known spec-text correction (recorded, not tuned) +------------------------------------------------- +The spec's "Regression example" paragraph asks for `max r == 5.0 A`. That is a +slip in the prose, contradicted by the spec's own shell table three lines +earlier: `5.0 A` is `r_build`, an upper **bound** on realised pair distances, not +a realised one. Simple-cubic distances are `3 * sqrt(i^2+j^2+k^2)` A, i.e. +`{3.000, 4.243, 5.196, 6.000, ...}` — `5.0` is not in the set, so no pair can +ever sit there. This file therefore pins the crystallographic value +`max r = sqrt2 * a = 4.2426406871 A` **and** the bound the prose was reaching +for (`count(r > r_build) == 0`, `max r < r_build`). The golden comes from the +lattice, not from a run; the spec sentence should be corrected to match. + +Goldens +------- + capture command : PYTHONPATH=src python regressions/md-neighborlist-skin-05-bind.py + Nothing was captured from a run. Every literal is + arithmetic on the page: 64*18 = 1152, 64*6 = 384, + 64*12 = 768, ceil(1.35*1152) = 1556, 3.5 + 1.5 = 5.0, + 3.0*sqrt(2) = 4.2426406871192851, and the gate table of + Section 3. + commit : d59a38f (d59a38fc77b700daec57c5a90676d42d32cf01be), with + the `md-neighborlist-skin-05-bind` working tree on top + (at d59a38f itself `NeighborList` has no `build` at all, + so `nl.build(batch)` raises `AttributeError` — that is the + RED this file was written against). + torch : 2.12.1+cpu (python 3.14.5) + date : 2026-08-09 + device / dtype : CPU, float64 — the cell and the positions are float64 + literals, so `shifts` and `_x_hold` are float64 too. + `edge_index` is int64 regardless. + oracle : none. No third-party package (torch and tensordict are + the repo's own core dependencies), no network, no + subprocess, no RNG, no wall-clock value, no filesystem + access. LAMMPS is the *specified* behaviour, not a + runtime dependency: its gate is transcribed into the table + in Section 3, never executed. + tolerance : 1e-12 A on `r_build` (float64 "position" exact band; 3.5 + and 1.5 are exactly representable, so 5.0 is exact). + 1e-9 A on the reconstructed pair distances, which pass + through a subtraction, an addition and a norm in float64 + (observed deviation ~1e-15 A). Exact integer equality on + every edge count, capacity, rebuild count, `ago`, + `ndanger`, pair tally and fired-update sequence — they are + counters, and a counter has no tolerance. Exact identity + (`is`) on the bound leaves: the whole point. + +Run: + PYTHONPATH=src python regressions/md-neighborlist-skin-05-bind.py +""" + +from __future__ import annotations + +import itertools +import math +import sys + +import torch +from tensordict import TensorDict + +from molix.md import NeighborList + +# --------------------------------------------------------------------------- +# The lattice every section runs on. +# --------------------------------------------------------------------------- + +#: Lattice spacing `a` in Angstrom. +SPACING = 3.0 + +#: Sites per axis; 4 keeps the `+a` and `-a` neighbours distinct atoms. +N_SIDE = 4 + +#: Cubic cell edge in Angstrom — four spacings, so the lattice tiles exactly. +BOX = SPACING * N_SIDE + +#: Atomic number written into the batch. Argon: chemically inert here, since no +#: potential is evaluated — the key exists because a batch carries it. +ATOMIC_NUMBER = 18 + +#: Interaction cutoff in Angstrom: between the 3.0 A first shell (in) and the +#: 4.243 A second shell (out), so no pair can flip class on rounding. +CUTOFF = 3.5 + +#: Verlet skin in Angstrom. `r_build = 5.0 A` sits between the 4.243 A second +#: shell (in) and the 5.196 A third shell (out) — a hard integer boundary, not a +#: tolerance question — and half-skin is a round 0.75 A. +SKIN = 1.5 + +CELL = torch.tensor( + [ + [BOX, 0.0, 0.0], + [0.0, BOX, 0.0], + [0.0, 0.0, BOX], + ], + dtype=torch.float64, +) + +# --------------------------------------------------------------------------- +# Section 1 goldens — the lattice at `r_build`. +# --------------------------------------------------------------------------- + +#: `CUTOFF + SKIN`, written out rather than recomputed from the two constants: +#: this file pins the derived radius, so deriving it here would be tautological. +EXPECTED_R_BUILD = 5.0 + +#: 64 sites x (6 at 3.0 A + 12 at 4.243 A) directed edges, both directions kept +#: (`symmetry=True`, the MD list's fixed setting). +EXPECTED_NUM_EDGES = 1152 + +#: ceil(1.35 * 1152) = ceil(1555.2), with the default `capacity_factor = 1.35`. +EXPECTED_CAPACITY = 1556 + +#: float64 "position" exact band; 3.5 + 1.5 = 5.0 is exact in binary. +LENGTH_TOL = 1e-12 + +# --------------------------------------------------------------------------- +# Section 2 goldens — the bind. +# --------------------------------------------------------------------------- + +#: `build` binds exactly the two live buffers, replacing the namespace wholesale. +EXPECTED_EDGE_KEYS = ("edge_index", "shifts") + +# --------------------------------------------------------------------------- +# Section 3 goldens — `every=1, delay=0, check=True` under a rigid drift. +# --------------------------------------------------------------------------- + +#: Rigid translation of *all* atoms along x per update, in Angstrom. +STEP_DISPLACEMENT = 0.2 + +#: Updates driven through the batch. +N_UPDATES = 20 + +#: Displacement since the last build is `0.2 * ago` A; it first exceeds +#: half-skin 0.75 A at `ago == 4` (0.8 A), and `rebuild` re-phases `ago` to 0. +EXPECTED_FIRED_UPDATES = (4, 8, 12, 16, 20) + +EXPECTED_REBUILD_COUNT = 5 + +#: Every rebuild fired at `ago == 4`, never at the first permitted opportunity +#: `max(every, delay) == 1`, so no rebuild was ever overdue. +EXPECTED_NDANGER = 0 + +#: The 20th update *is* a rebuild, so the clock is back at 0 when the loop ends. +EXPECTED_FINAL_AGO = 0 + +# --------------------------------------------------------------------------- +# Section 4 goldens — the geometry read back out of the batch. +# --------------------------------------------------------------------------- + +#: First shell: the lattice spacing itself, exactly. +EXPECTED_MIN_PAIR_DISTANCE = SPACING + +#: Second shell — the face diagonal `sqrt2 * a = 4.2426406871192851 A`, the +#: largest distance that fits inside `r_build = 5.0 A`. Crystallography, not a +#: captured run value: the next shell is `sqrt3 * a = 5.196 A`, outside. +EXPECTED_MAX_PAIR_DISTANCE = SPACING * math.sqrt(2.0) + +#: 64 sites x 6 first-shell neighbours, inside the *interaction* cutoff 3.5 A. +EXPECTED_PAIRS_WITHIN_CUTOFF = 384 + +#: 64 sites x 12 second-shell neighbours — the skin band `(3.5, 5.0] A`, edges +#: the list carries and the model's own envelope masks to zero. +EXPECTED_PAIRS_IN_SKIN_BAND = 768 + +#: `r_build` is a hard horizon: the kernel builds at it, so nothing is beyond. +EXPECTED_PAIRS_BEYOND_R_BUILD = 0 + +#: float64 through a subtract, an add and a 3-vector norm; observed ~1e-15 A. +DISTANCE_TOL = 1e-9 + + +def simple_cubic_lattice() -> torch.Tensor: + """Build the 4x4x4 simple-cubic lattice every section runs on. + + Returns: + Positions ``(64, 3)`` in Angstrom, float64: the Cartesian product of + ``{0.0, 3.0, 6.0, 9.0}`` with itself three times, lexicographic order. + """ + coordinates = tuple(index * SPACING for index in range(N_SIDE)) + return torch.tensor( + list(itertools.product(coordinates, repeat=3)), + dtype=torch.float64, + ) + + +def md_batch(positions: torch.Tensor) -> TensorDict: + """Wrap *positions* in the batch shape an MD driver hands the list. + + Args: + positions: Positions ``(N, 3)`` in Angstrom, float64. + + Returns: + A ``TensorDict`` with root ``batch_size=[]`` carrying ``("atoms", + "pos")`` / ``"Z"`` / ``"batch"`` at ``batch_size=[N]`` — the two-tier + contract's post-collate tier, single system, no ``"edges"`` yet. The + positions leaf is stored by reference, so writing into it in place is + what an integrator does. + """ + n_atoms = int(positions.shape[0]) + return TensorDict( + { + "atoms": TensorDict( + { + "pos": positions, + "Z": torch.full((n_atoms,), ATOMIC_NUMBER, dtype=torch.long), + "batch": torch.zeros(n_atoms, dtype=torch.long), + }, + batch_size=[n_atoms], + ), + }, + batch_size=[], + ) + + +def rigidly_translated(base: torch.Tensor, offset: float) -> torch.Tensor: + """Return *base* with **every** atom translated by *offset* along x. + + A rigid translation leaves every interatomic displacement — and so every + minimum image — exactly unchanged, which is what lets Section 4 assert the + neighbour set is bit-identical after a 4.0 A drift. + + Args: + base: Reference positions ``(N, 3)`` in Angstrom. + offset: Translation along x, in Angstrom, applied to all atoms. + + Returns: + A fresh ``(N, 3)`` tensor; *base* is never mutated, so each step is an + absolute displacement from the lattice rather than an accumulation of + rounding. + """ + shift = torch.tensor([offset, 0.0, 0.0], dtype=base.dtype, device=base.device) + return base + shift + + +def pair_distances(batch: TensorDict, num_edges: int) -> torch.Tensor: + """Reconstruct live pair distances **through the batch**, as a potential does. + + Reads positions, edge indices and periodic shifts from *batch* only — never + from the list — so a binding that copied instead of aliasing shows up as a + stale neighbour set here. + + Args: + batch: Batch carrying ``("atoms", "pos")`` and the bound + ``("edges", "edge_index")`` / ``("edges", "shifts")``. + num_edges: Live edge count; rows ``[num_edges, capacity)`` are dead + padding and are excluded. + + Returns: + Minimum-image distances ``(num_edges,)`` in Angstrom, computed as + ``|| pos[target] - pos[source] + shift ||`` per the repo edge + convention (``edge_index[:, 0]`` source, ``[:, 1]`` target). + """ + edge_index = batch["edges", "edge_index"][:num_edges] + shifts = batch["edges", "shifts"][:num_edges] + positions = batch["atoms", "pos"] + source, target = edge_index[:, 0], edge_index[:, 1] + return torch.linalg.norm(positions[target] - positions[source] + shifts, dim=-1) + + +def fresh_list(positions: torch.Tensor) -> NeighborList: + """Construct the one configuration every section shares. + + Args: + positions: Initial positions ``(N, 3)`` in Angstrom the list builds at. + + Returns: + A ``NeighborList`` at ``cutoff=3.5 A``, ``skin=1.5 A``, default gate + ``every=1, delay=0, check=True``. + """ + return NeighborList( + cell=CELL, + cutoff=CUTOFF, + positions=positions, + skin=SKIN, + every=1, + delay=0, + check=True, + ) + + +# --------------------------------------------------------------------------- +# Checking +# --------------------------------------------------------------------------- + + +class Checker: + """Collects every deviation so one run reports all failures, not the first.""" + + def __init__(self) -> None: + self.failures: list[str] = [] + + def _row(self, name: str, got: object, ok: bool) -> None: + print(f" {name:<44} {got!s:<26} {'ok' if ok else 'FAILED'}") + + def exact(self, name: str, got: object, want: object) -> None: + """Assert a counter, a shape, an identity or a sequence — no tolerance.""" + ok = got == want + if not ok: + self.failures.append(f"{name}: got {got!r}, want {want!r}") + self._row(name, got, ok) + + def within(self, name: str, got: float, tol: float) -> None: + """Assert a measured deviation is at most *tol* (Angstrom).""" + ok = got <= tol + if not ok: + self.failures.append(f"{name}: deviation {got!r} A exceeds {tol!r} A") + self._row(name, f"{got:.3e} A", ok) + + +def check_bind(checker: Checker) -> tuple[NeighborList, TensorDict]: + """Sections 1-2 — the lattice at `r_build`, and what `build` writes. + + Args: + checker: Failure collector. + + Returns: + The bound list and its batch, for the driven schedule to reuse. + """ + print(f"Bind ({N_SIDE}x{N_SIDE}x{N_SIDE} sc, {SPACING} A spacing, {BOX} A cell, float64 CPU)") + + lattice = simple_cubic_lattice() + neighbor_list = fresh_list(lattice.clone()) + batch = md_batch(lattice.clone()) + + checker.within( + "bind.r_build_deviation", + abs(neighbor_list.r_build - EXPECTED_R_BUILD), + LENGTH_TOL, + ) + checker.exact("bind.num_edges", neighbor_list.num_edges, EXPECTED_NUM_EDGES) + checker.exact("bind.capacity", neighbor_list.capacity, EXPECTED_CAPACITY) + + bound = neighbor_list.build(batch) + + # The returned object *is* the argument, so `potential(nl.build(batch))` + # composes with the repo's forward(td) -> td convention. + checker.exact("bind.returns_same_batch", bound is batch, True) + # Identity, not equality: a TensorDict that copied on assignment would give + # a frozen neighbour set that still compares equal at bind time. + checker.exact( + "bind.edge_index_is_buffer", + batch["edges", "edge_index"] is neighbor_list.edge_index, + True, + ) + checker.exact( + "bind.shifts_is_buffer", + batch["edges", "shifts"] is neighbor_list.shifts, + True, + ) + checker.exact("bind.edge_keys", tuple(sorted(batch["edges"].keys())), EXPECTED_EDGE_KEYS) + # Capacity, not num_edges: the buffers are fixed-shape, tail padded dead. + checker.exact( + "bind.edges_batch_size", + tuple(batch["edges"].batch_size), + (EXPECTED_CAPACITY,), + ) + # A bind is not physics: it also runs on every .to() re-sync. + checker.exact("bind.rebuild_count", neighbor_list.rebuild_count, 0) + checker.exact("bind.ago", neighbor_list.ago, 0) + checker.exact("bind.ndanger", neighbor_list.ndanger, 0) + + return neighbor_list, batch + + +def check_driven_schedule(checker: Checker, neighbor_list: NeighborList, batch: TensorDict) -> None: + """Section 3 — twenty `update(batch)` calls under a rigid +0.2 A/step drift. + + Args: + checker: Failure collector. + neighbor_list: The list bound in Section 2. + batch: Its bound batch; positions are written in place, as an integrator + would, so the bound leaves are never rebound. + """ + print( + f"\nDriven schedule (every=1, delay=0, check=True, rigid {STEP_DISPLACEMENT} A/update " + f"on all {N_SIDE**3} atoms)" + ) + + lattice = simple_cubic_lattice() + positions = batch["atoms", "pos"] + + fired: list[int] = [] + for update in range(1, N_UPDATES + 1): + positions.copy_(rigidly_translated(lattice, STEP_DISPLACEMENT * update)) + if neighbor_list.update(batch): + fired.append(update) + + checker.exact("driven.fired_updates", tuple(fired), EXPECTED_FIRED_UPDATES) + checker.exact("driven.rebuild_count", neighbor_list.rebuild_count, EXPECTED_REBUILD_COUNT) + checker.exact("driven.ndanger", neighbor_list.ndanger, EXPECTED_NDANGER) + checker.exact("driven.final_ago", neighbor_list.ago, EXPECTED_FINAL_AGO) + # The tie survived five in-place rebuilds: still the same tensors. + checker.exact("driven.pos_leaf_not_rebound", batch["atoms", "pos"] is positions, True) + checker.exact( + "driven.edge_index_still_bound", + batch["edges", "edge_index"] is neighbor_list.edge_index, + True, + ) + checker.exact( + "driven.shifts_still_bound", + batch["edges", "shifts"] is neighbor_list.shifts, + True, + ) + + +def check_geometry_through_batch( + checker: Checker, neighbor_list: NeighborList, batch: TensorDict +) -> None: + """Section 4 — reconstruct the shells from the batch after the driven run. + + Args: + checker: Failure collector. + neighbor_list: The driven list (read only for ``num_edges``). + batch: Its bound batch — the sole source of positions, indices, shifts. + """ + print("\nGeometry read back through the batch (after 4.0 A of rigid, unwrapped drift)") + + # A rigid translation preserves every minimum image, so the neighbour set is + # invariant under the whole schedule. + checker.exact("geometry.num_edges_unchanged", neighbor_list.num_edges, EXPECTED_NUM_EDGES) + + distances = pair_distances(batch, neighbor_list.num_edges) + + checker.exact("geometry.distances_counted", int(distances.numel()), EXPECTED_NUM_EDGES) + checker.within( + "geometry.min_distance_deviation", + abs(float(distances.min()) - EXPECTED_MIN_PAIR_DISTANCE), + DISTANCE_TOL, + ) + checker.within( + "geometry.max_distance_deviation", + abs(float(distances.max()) - EXPECTED_MAX_PAIR_DISTANCE), + DISTANCE_TOL, + ) + checker.exact( + "geometry.pairs_within_cutoff", + int((distances <= CUTOFF).sum()), + EXPECTED_PAIRS_WITHIN_CUTOFF, + ) + checker.exact( + "geometry.pairs_in_skin_band", + int(((distances > CUTOFF) & (distances <= EXPECTED_R_BUILD)).sum()), + EXPECTED_PAIRS_IN_SKIN_BAND, + ) + checker.exact( + "geometry.pairs_beyond_r_build", + int((distances > EXPECTED_R_BUILD).sum()), + EXPECTED_PAIRS_BEYOND_R_BUILD, + ) + + +def check_dispatch_equivalence(checker: Checker, driven: NeighborList) -> None: + """Section 5 — the raw-tensor front door reaches the same state. + + Args: + checker: Failure collector. + driven: The batch-driven list to compare buffers against. + """ + print("\nDispatch equivalence (same schedule through raw update(pos))") + + lattice = simple_cubic_lattice() + neighbor_list = fresh_list(lattice.clone()) + + fired: list[int] = [] + for update in range(1, N_UPDATES + 1): + positions = rigidly_translated(lattice, STEP_DISPLACEMENT * update) + if neighbor_list.update(positions): + fired.append(update) + + checker.exact("raw.fired_updates", tuple(fired), EXPECTED_FIRED_UPDATES) + checker.exact("raw.rebuild_count", neighbor_list.rebuild_count, EXPECTED_REBUILD_COUNT) + checker.exact("raw.ndanger", neighbor_list.ndanger, EXPECTED_NDANGER) + checker.exact("raw.num_edges", neighbor_list.num_edges, driven.num_edges) + # Whole buffers, dead padding tail included — not just the live prefix. + checker.exact( + "raw.edge_index_equals_batch_path", + bool(torch.equal(neighbor_list.edge_index, driven.edge_index)), + True, + ) + checker.exact( + "raw.shifts_equals_batch_path", + bool(torch.equal(neighbor_list.shifts, driven.shifts)), + True, + ) + + +def main() -> int: + """Pin the bind surface: build, drive through the batch, read the physics back.""" + checker = Checker() + + neighbor_list, batch = check_bind(checker) + check_driven_schedule(checker, neighbor_list, batch) + check_geometry_through_batch(checker, neighbor_list, batch) + check_dispatch_equivalence(checker, neighbor_list) + + if checker.failures: + print("\nFAILED — the TensorDict bind surface no longer holds:") + for failure in checker.failures: + print(f" {failure}") + print( + "\nEvery golden above is analytic (64*18 = 1152, 64*6 = 384, " + "64*12 = 768, ceil(1.35*1152) = 1556, 3.0*sqrt(2) = 4.2426406871, " + "and the gate table in the module docstring). A disagreement is a " + "DEFECT REPORT against src/molix/md/neighbors.py, never a reason to " + "edit the literal. In particular a broken identity assertion means " + "the batch holds a *copy* of the buffers: the neighbour set a " + "potential reads would freeze at bind time while the list keeps " + "rebuilding, and a golden retuned to match would hide a frozen " + "potential-energy surface." + ) + return 1 + print("\nOK") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/regressions/md-neighborlist-skin-06-bins.py b/regressions/md-neighborlist-skin-06-bins.py new file mode 100644 index 0000000..468ea92 --- /dev/null +++ b/regressions/md-neighborlist-skin-06-bins.py @@ -0,0 +1,768 @@ +"""Public-API scenario for the binned (cell-list) build backend of the MD neighbour list. + +Spec: `md-neighborlist-skin-06-bins`. + +`molix.md.NeighborList` gained a keyword-only `bin` argument selecting a second +build **backend**: `bin=None` (default) hands the system to the compiled O(N^2) +pair kernel, a float switches to the pure-torch periodic cell list. `bin` is a +**cost** knob, never a physics knob — which is a claim about the *edge set*, and +that is what this file pins. + +The oracle is **in repo**: the untouched `bin=None` kernel path. Nothing here +imports, subprocesses or downloads a third party (torch is the repo's own core +dependency); no network, no filesystem, no RNG, no wall-clock *assertion*. Every +system is built from integer arithmetic or written-out literals on this page. + +What set equality alone cannot see +---------------------------------- +A binned build searches a stencil of neighbouring bins. When `2*k_i + 1 > n_i` +the raw stencil wraps onto the same bin twice, so a correct-looking *set* of +edges can hide every pair being emitted **twice**. Comparison here is therefore +three-fold, on both paths: + +1. `set(canonical_keys_binned) == set(canonical_keys_kernel)` — same pairs, same + periodic images; +2. `len(set(directed_keys)) == num_edges` — no duplicate directed edge; +3. `2 * len(set(canonical_keys)) == num_edges` — the bidirectional list collapses + exactly 2:1 onto orientation-free keys. + +A directed key is `(source, target, round(shift, 6))`. A canonical key is the +same triple re-oriented low index first, with the shift negated when `s > t`, so +it is orientation-free while still separating periodic images of the same pair. +Six decimals is far coarser than float64 noise on shifts that are integer +combinations of cell vectors (~1e-15 A here) and far finer than the smallest +distinct shift component (6.0 A), so the rounding can neither merge two images +nor split one. + +Section 1 — cubic lattice, automatic bin, the aliasing regime +------------------------------------------------------------- +64 atoms on a 4x4x4 simple-cubic lattice, `a = 3.0 A`, in a 12.0 A cube (built +in-script from `arange`/`meshgrid`; no RNG). `cutoff = 3.5 A`, `skin = 1.5 A`: + + r_build = 3.5 + 1.5 = 5.0 A (<= 12.0/2 = 6.0 A, guard clear) + bin=0.0 => requested b = r_build/2 = 2.5 A (LAMMPS `binsize_optimal`) + n_i = floor(12.0 / 2.5) = floor(4.8) = 4 => n_bins == (4, 4, 4) + b_i = 12.0 / 4 = 3.0 A + k_i = ceil(5.0 / 3.0) = 2 => 2*k_i + 1 = 5 > 4 = n_i + +so the raw stencil wraps onto itself: this is exactly the regime where the +unique-residue construction is load-bearing, and where a naive stencil would +double every edge while keeping the set right. + +Simple cubic at `a = 3.0 A` has exact shells — 6 at 3.000 A, 12 at 4.243 A, 8 at +5.196 A — so `r_build = 5.0 A` admits 6 + 12 = 18 neighbours per atom: + + num_edges = 64 * 18 = 1152 directed edges (full bidirectional list) + unordered pairs = 1152 / 2 = 576 + +Four sites per axis keeps the `+a` and `-a` neighbours distinct atoms, so the +multiplicities are not double-counted images. Both shells sit far from the +5.0 A radius (4.243 and 5.196), so no pair can flip class on rounding. + +Section 2 — the grid is derived from `r_build`, not `cutoff` +------------------------------------------------------------- +The same lattice at `skin = 0.0` (`r_build = 3.5 A`): + + requested b = 1.75 A; n_i = floor(12.0/1.75) = floor(6.857) = 6 + n_bins == (6, 6, 6); b_i = 2.0 A; k_i = ceil(3.5/2.0) = 2 + 2*k_i + 1 = 5 <= 6 => 125 of 216 bins searched — real pruning + num_edges = 64 * 6 = 384 (only the 3.0 A shell is inside 3.5 A) + +`(6,6,6)` versus section 1's `(4,4,4)` on the *same cell* is what pins the grid +to `r_build`: a backend that binned on `cutoff` would report `(6,6,6)` in section +1 too, and one that binned on `r_build` while *filtering* on `cutoff` would +return 384 edges there instead of 1152. + +Section 3 — the bin size is a cost knob +---------------------------------------- +Section 1's system, built with explicit bin thicknesses: + + bin=5.0 A => n_i = floor(12.0/5.0) = 2 => n_bins == (2, 2, 2) + bin=12.0 A => n_i = floor(12.0/12.0) = 1 => n_bins == (1, 1, 1) + +Both must reproduce section 1's canonical edge set **exactly** — same 576 +unordered pairs, same 1152 directed edges. `(1,1,1)` is the documented graceful +degeneration: one bin holding all 64 atoms is an all-pairs search, correct and +merely not faster. Three different grids over one configuration is the strongest +available statement that `bin` never moves an atom's neighbours. + +Section 4 — triclinic, where the perpendicular width matters +------------------------------------------------------------- +Cell rows `a_1 = (10, 0, 0)`, `a_2 = (6, 8, 0)`, `a_3 = (0, 0, 10)` (Angstrom), +`V = 800 A^3`, perpendicular widths `w_i = V / ||a_j x a_k||`: + + w = (800/100, 800/100, 800/80) = (8, 8, 10) A (guard: min w / 2 = 4.0 A) + +12 atoms are written out as **literal fractional coordinates** and mapped by +`frac @ cell`; `cutoff = 3.0 A`, `skin = 0.5 A` => `r_build = 3.5 A`, strictly +inside the 4.0 A guard. With `bin=0.0` (requested 1.75 A): + + n_bins == (floor(8/1.75), floor(8/1.75), floor(10/1.75)) = (4, 4, 5) + +which is **discriminating**: sizing on the row norm `||a_2|| = 10` instead of the +perpendicular width `w_2 = 8` would give `(4, 5, 5)`. This is also the fixture +that exercises the `|f_i| <= 1/2` lemma — the binned path's minimum image comes +from fractional rounding, the kernel's from a sequential diagonal reduction, and +on a sheared cell those two agree only because `r_build <= min_i w_i / 2`. + +The edge count `24` here is **not** analytic: it was captured once from the +in-repo kernel oracle at implementation time (see Goldens). The script also +recomputes, from the literals alone, the closest minimum-image pair distance and +the margin from `r_build` over all 66 pairs and all 27 images, and refuses to run +on a fixture whose margin has collapsed — a float tie at the radius would make +the comparison flaky rather than wrong. + +Section 5 — unwrapped positions +-------------------------------- +Section 1's lattice with 8 atoms translated by whole cell vectors (`+-12 A`, +including two diagonal images). The binned path wraps fractional coordinates +into `[0, 1)` **for indexing only**: displacements are taken from the unwrapped +coordinates, and the shifts refer to the stored, unwrapped positions. So both +paths must still agree edge for edge, and the multiset of +`(min(s,t), max(s,t), round(||pos[t] - pos[s] + shift||, 6))` must be *identical* +to section 1's — the physics is translation invariant even though the individual +shift labels are not. A backend that differenced the wrapped copy would return +the same distances but wrong shifts; a backend that stored wrapped positions +would break the link-04 raw-displacement invariant. Comparing distances *and* +edge sets catches both. + +Advisory timing (asserted nowhere) +----------------------------------- +The two build wall clocks for a 512-atom / 24 A cube are printed for information. +**No timing threshold is asserted in this file, and none belongs here**: the +numbers are machine-, thread- and build-bound (the binned path loses to the +kernel at high `OMP_NUM_THREADS`, where per-op OpenMP region overhead dominates +its 125-offset loop). They are printed, never checked. + +Drift policy +------------ +A disagreement between this file and the runtime is a **DEFECT REPORT, never a +golden edit**. Sections 1-3 and 5 are analytic — simple-cubic shell +multiplicities, `floor(w_i / b)` arithmetic and translation invariance, all +derived on this page — and section 4's single captured literal came from the +kernel oracle, which this file's whole point is to hold the binned path against. +If a count or a set comparison comes back wrong, open a defect against +`src/molix/md/neighbors.py`; do not retune a literal to whatever the run printed. +A duplicated edge doubles a pair's contribution to the energy and a missing one +deletes it, and both of those are silent at runtime — this file is where they are +supposed to become loud. + +Goldens +------- + capture command : PYTHONPATH=src python regressions/md-neighborlist-skin-06-bins.py + Sections 1-3 and 5 captured nothing: 64*18 = 1152, + 64*6 = 384, floor(12/2.5) = 4, floor(12/1.75) = 6, + floor(12/5) = 2, floor(12/12) = 1, floor(8/1.75) = 4 and + floor(10/1.75) = 5 are arithmetic on this page. + Section 4's `TRICLINIC_EDGES = 24` was captured ONCE, from + the in-repo `bin=None` kernel path on the 12 literal + fractional coordinates below (which also reported closest + pair 2.012461 A and margin 0.145898 A from r_build). + commit : 639c9df (639c9df8818102a602b96fa326dbc4e7eab5b013), with + the `md-neighborlist-skin-06-bins` working tree on top (at + 639c9df itself `NeighborList.__init__` takes no `bin`, so + every scenario here raises `TypeError` — that is the RED + this file was written against). + torch : 2.12.1+cpu (python 3.14.5) + date : 2026-08-09 + device / dtype : CPU, float64 — cells and positions are float64 literals, so + `shifts` is float64 too; `edge_index` is int64 regardless. + oracle : in repo. The compiled O(N^2) kernel path reached through + the public constructor as `bin=None` — the backend this + link does not touch. No third-party package, no network, no + subprocess, no downloaded reference, no RNG. + tolerance : exact integer equality on every edge count, bin-count tuple + and set cardinality — they are counters, and a counter has + no tolerance. Shifts and distances are compared as keys + rounded to 6 decimals (float64 "position" band is 1e-12; + the values here are integer combinations of cell vectors + carrying ~1e-15 A of noise, and the smallest distinct shift + component is 6.0 A, so the rounding is unambiguous). + +Run: + PYTHONPATH=src python regressions/md-neighborlist-skin-06-bins.py +""" + +from __future__ import annotations + +import itertools +import sys +import time + +import torch + +from molix.md import NeighborList + +# --------------------------------------------------------------------------- +# The cubic system of sections 1, 2, 3 and 5. +# --------------------------------------------------------------------------- + +#: Lattice spacing `a` in Angstrom. +SPACING = 3.0 + +#: Sites per axis; 4 keeps the `+a` and `-a` neighbours distinct atoms. +N_SIDE = 4 + +#: Cubic cell edge in Angstrom — four spacings, so the lattice tiles exactly. +BOX = SPACING * N_SIDE + +#: Interaction cutoff in Angstrom: between the 3.0 A first shell (in) and the +#: 4.243 A second shell (out), so no pair can flip class on rounding. +CUTOFF = 3.5 + +#: Verlet skin in Angstrom. `r_build = 5.0 A` sits between the 4.243 A second +#: shell (in) and the 5.196 A third shell (out). +SKIN = 1.5 + +CUBIC_CELL = torch.tensor( + [ + [BOX, 0.0, 0.0], + [0.0, BOX, 0.0], + [0.0, 0.0, BOX], + ], + dtype=torch.float64, +) + +#: `bin=0.0` asks for the automatic `r_build / 2` thickness (LAMMPS +#: `nbin_standard`); every other float is an explicit requested thickness in A. +AUTO_BIN = 0.0 + +# --------------------------------------------------------------------------- +# Section 1 goldens — automatic bin at `r_build = 5.0 A`. +# --------------------------------------------------------------------------- + +#: floor(12.0 / (5.0/2)) = floor(4.8) = 4 bins per axis; b_i = 3.0 A and +#: k_i = ceil(5.0/3.0) = 2, so 2*k_i + 1 = 5 > 4 — the wrapping stencil. +CUBIC_N_BINS_AT_R_BUILD = (4, 4, 4) + +#: 64 sites x (6 at 3.0 A + 12 at 4.243 A), both directions kept. +CUBIC_EDGES_AT_R_BUILD = 1152 + +#: Directed edges collapse 2:1 onto orientation-free keys. +CUBIC_PAIRS_AT_R_BUILD = CUBIC_EDGES_AT_R_BUILD // 2 + +# --------------------------------------------------------------------------- +# Section 2 goldens — same cell, bare cutoff (`skin = 0.0`). +# --------------------------------------------------------------------------- + +#: floor(12.0 / (3.5/2)) = floor(6.857) = 6; b_i = 2.0 A, k_i = 2, so the +#: stencil is 125 of 216 bins — pruning, not aliasing. +CUBIC_N_BINS_AT_CUTOFF = (6, 6, 6) + +#: 64 sites x 6 at 3.0 A — only the first shell is inside 3.5 A. +CUBIC_EDGES_AT_CUTOFF = 384 + +# --------------------------------------------------------------------------- +# Section 3 goldens — explicit bin thicknesses on section 1's system. +# --------------------------------------------------------------------------- + +#: (requested thickness in A, expected `n_bins`). 5.0 A gives two bins per axis +#: (b_i = 6.0 A, k_i = 1, residues {0,1} = the whole grid); 12.0 A is the whole +#: cell in one bin — the documented graceful degeneration to an all-pairs search. +EXPLICIT_BINS: tuple[tuple[float, tuple[int, int, int]], ...] = ( + (5.0, (2, 2, 2)), + (12.0, (1, 1, 1)), +) + +# --------------------------------------------------------------------------- +# Section 4 — the triclinic system. +# --------------------------------------------------------------------------- + +TRICLINIC_CELL = torch.tensor( + [ + [10.0, 0.0, 0.0], + [6.0, 8.0, 0.0], + [0.0, 0.0, 10.0], + ], + dtype=torch.float64, +) + +#: 12 literal fractional coordinates, mapped to Angstrom by `frac @ cell`. +#: Chosen so the closest minimum-image pair is 2.012 A (well clear of the `r > 0` +#: filter) and no pair sits within 0.145 A of `r_build = 3.5 A` (well clear of a +#: float tie at the radius) — both re-verified at run time below. +TRICLINIC_FRACTIONAL = torch.tensor( + [ + [0.00, 0.15, 0.20], + [0.00, 0.45, 0.50], + [0.30, 0.20, 0.90], + [0.35, 0.25, 0.15], + [0.35, 0.50, 0.25], + [0.35, 0.70, 0.40], + [0.50, 0.15, 0.80], + [0.60, 0.40, 0.70], + [0.75, 0.05, 0.75], + [0.75, 0.60, 0.00], + [0.90, 0.20, 0.80], + [0.90, 0.25, 0.00], + ], + dtype=torch.float64, +) + +#: Interaction cutoff and skin in Angstrom: `r_build = 3.5 A`, strictly inside +#: the guard `min_i w_i / 2 = 4.0 A` for this cell. +TRICLINIC_CUTOFF = 3.0 +TRICLINIC_SKIN = 0.5 + +#: floor(w / 1.75) on the perpendicular widths w = (8, 8, 10) A. Sizing on the +#: row norm ||a_2|| = 10 A instead would give (4, 5, 5). +TRICLINIC_N_BINS = (4, 4, 5) + +#: Captured ONCE from the in-repo `bin=None` kernel oracle on the literals above +#: (2026-08-09, torch 2.12.1+cpu, CPU float64): 12 unordered pairs, both +#: directions kept. A disagreement is a defect report, not a new capture. +TRICLINIC_EDGES = 24 + +#: Smallest allowed `| r - r_build |` over all pairs and images, in Angstrom. +#: Not a golden — a precondition floor, ~5 orders above float64 noise on these +#: distances and ~5 orders below the fixture's actual 0.145898 A margin. It fails +#: only if the literals above are edited into a tie at the build radius. +MARGIN_FLOOR = 1e-6 + +# --------------------------------------------------------------------------- +# Section 5 — whole-cell translations of section 1's lattice. +# --------------------------------------------------------------------------- + +#: `(atom index, cell-vector image)` for the 8 translated atoms: three axes in +#: each direction plus two diagonals, so the fixture is not a single-axis +#: special case. Each atom moves by an exact lattice vector, so the periodic +#: system is unchanged and only the *labels* (shifts) may move. +TRANSLATED_ATOMS: tuple[tuple[int, tuple[int, int, int]], ...] = ( + (0, (1, 0, 0)), + (5, (0, 1, 0)), + (9, (0, 0, 1)), + (17, (-1, 0, 0)), + (23, (0, -1, 0)), + (38, (0, 0, -1)), + (47, (1, -1, 0)), + (60, (-1, 1, 1)), +) + +# --------------------------------------------------------------------------- +# Advisory timing system (printed, never asserted). +# --------------------------------------------------------------------------- + +#: 8x8x8 sites at 3.0 A — 512 atoms in a 24 A cube, where `n_bins == (9,9,9)` +#: and the stencil prunes to 125 of 729 bins. +TIMING_N_SIDE = 8 + +# --------------------------------------------------------------------------- +# Key construction — the comparison currency. +# --------------------------------------------------------------------------- + +#: Decimals kept when rounding a shift component or a distance into a key. +#: Coarser than float64 noise on integer combinations of cell vectors (~1e-15 A), +#: finer than the smallest distinct shift component (6.0 A). +KEY_DECIMALS = 6 + +Shift = tuple[float, ...] +EdgeKey = tuple[int, int, Shift] +PairDistance = tuple[int, int, float] + + +def cubic_lattice(n_side: int) -> torch.Tensor: + """Build a simple-cubic lattice of `n_side**3` sites at `SPACING` Angstrom. + + Args: + n_side: Sites per axis. + + Returns: + Positions ``(n_side**3, 3)`` in Angstrom, float64, in lexicographic + order. Built from ``arange``/``meshgrid``: no RNG, no data file. + """ + axis = torch.arange(n_side, dtype=torch.float64) * SPACING + grid_x, grid_y, grid_z = torch.meshgrid(axis, axis, axis, indexing="ij") + return torch.stack((grid_x.reshape(-1), grid_y.reshape(-1), grid_z.reshape(-1)), dim=-1) + + +def translated_lattice() -> torch.Tensor: + """Section 1's lattice with `TRANSLATED_ATOMS` moved by whole cell vectors. + + Returns: + Positions ``(64, 3)`` in Angstrom, float64 — the same periodic system as + :func:`cubic_lattice`, deliberately *not* wrapped back into the box. + """ + positions = cubic_lattice(N_SIDE) + for atom, image in TRANSLATED_ATOMS: + positions[atom] += torch.tensor(image, dtype=torch.float64) @ CUBIC_CELL + return positions + + +def directed_keys(neighbor_list: NeighborList) -> list[EdgeKey]: + """Read the live edges as `(source, target, rounded shift)` triples. + + Args: + neighbor_list: A built list; only its public buffers are read. + + Returns: + One key per live directed edge, in buffer order (which is *not* part of + the contract — every comparison below is a set or a sorted multiset). + """ + edge_index = neighbor_list.edge_index[: neighbor_list.num_edges].tolist() + shifts = neighbor_list.shifts[: neighbor_list.num_edges].tolist() + return [ + (int(source), int(target), tuple(round(component, KEY_DECIMALS) for component in shift)) + for (source, target), shift in zip(edge_index, shifts, strict=True) + ] + + +def canonical_keys(neighbor_list: NeighborList) -> list[EdgeKey]: + """Re-orient :func:`directed_keys` low index first, negating the shift. + + The shift is the periodic remainder of ``pos[target] - pos[source]``, so it + flips sign wholesale with the edge. Re-orienting makes the key blind to which + way round the bidirectional list stored a pair while still separating + periodic images of it. + + Args: + neighbor_list: A built list. + + Returns: + One orientation-free key per live directed edge; a duplicate-free + bidirectional list yields each unordered pair exactly twice. + """ + canonical: list[EdgeKey] = [] + for source, target, shift in directed_keys(neighbor_list): + if source > target: + source, target, shift = target, source, tuple(-component for component in shift) + canonical.append((source, target, shift)) + return canonical + + +def pair_distances(neighbor_list: NeighborList, positions: torch.Tensor) -> list[PairDistance]: + """Reconstruct each live edge's length from the stored positions and shifts. + + ``edge_diff = pos[target] - pos[source] + shift`` is the documented way a + consumer recovers the displacement, and it is the step that fails if the + shifts refer to a wrapped copy of the positions rather than the stored ones. + + Args: + neighbor_list: A built list. + positions: The positions it was built at ``(N, 3)`` in Angstrom. + + Returns: + Sorted ``(min index, max index, rounded distance in Angstrom)`` triples — + a multiset, so a duplicated edge is visible as a repeated entry. + """ + edge_index = neighbor_list.edge_index[: neighbor_list.num_edges] + shifts = neighbor_list.shifts[: neighbor_list.num_edges] + displacement = positions[edge_index[:, 1]] - positions[edge_index[:, 0]] + shifts + distance = torch.linalg.norm(displacement, dim=-1) + return sorted( + (min(int(source), int(target)), max(int(source), int(target)), round(length, KEY_DECIMALS)) + for (source, target), length in zip(edge_index.tolist(), distance.tolist(), strict=True) + ) + + +def minimum_image_extremes(positions: torch.Tensor, cell: torch.Tensor) -> tuple[float, float]: + """Closest minimum-image pair distance, and the distance closest to `r_build`. + + Brute force over all unordered pairs and all 27 images — independent of the + neighbour list under test, so it can serve as its precondition. + + Args: + positions: Positions ``(N, 3)`` in Angstrom. + cell: Cell vectors ``(3, 3)`` in Angstrom, one per row. + + Returns: + ``(closest, nearest_to_boundary)`` in Angstrom: the smallest pair + distance, and the pair distance minimising ``|r - r_build|`` (returned as + the distance itself, so the caller reports the margin it cares about). + """ + images = torch.tensor(list(itertools.product((-1, 0, 1), repeat=3)), dtype=torch.float64) @ cell + n_atoms = int(positions.shape[0]) + closest = float("inf") + boundary = float("inf") + r_build = TRICLINIC_CUTOFF + TRICLINIC_SKIN + for first, second in itertools.combinations(range(n_atoms), 2): + separation = positions[second] - positions[first] + images + distance = float(torch.linalg.norm(separation, dim=-1).min()) + closest = min(closest, distance) + if abs(distance - r_build) < abs(boundary - r_build): + boundary = distance + return closest, boundary + + +# --------------------------------------------------------------------------- +# Checking +# --------------------------------------------------------------------------- + + +class Checker: + """Collects every deviation so one run reports all failures, not the first.""" + + def __init__(self) -> None: + self.failures: list[str] = [] + + def _row(self, name: str, got: object, ok: bool) -> None: + print(f" {name:<46} {got!s:<24} {'ok' if ok else 'FAILED'}") + + def exact(self, name: str, got: object, want: object) -> None: + """Assert a counter, a bin-count tuple or a set cardinality.""" + ok = got == want + if not ok: + self.failures.append(f"{name}: got {got!r}, want {want!r}") + self._row(name, got, ok) + + def same_edges(self, name: str, binned: list[EdgeKey], kernel: list[EdgeKey]) -> None: + """Assert the two backends produced the same canonical edge set.""" + binned_set, kernel_set = set(binned), set(kernel) + ok = binned_set == kernel_set + if not ok: + only_binned = sorted(binned_set - kernel_set)[:3] + only_kernel = sorted(kernel_set - binned_set)[:3] + self.failures.append( + f"{name}: {len(binned_set - kernel_set)} keys only in the binned path " + f"(e.g. {only_binned}), {len(kernel_set - binned_set)} only in the kernel " + f"path (e.g. {only_kernel})" + ) + self._row(name, f"{len(binned_set)} keys", ok) + + def same_distances( + self, name: str, measured: list[PairDistance], reference: list[PairDistance] + ) -> None: + """Assert two builds carry the same multiset of `(pair, distance)` entries.""" + ok = measured == reference + if not ok: + self.failures.append( + f"{name}: {len(measured)} entries against {len(reference)}; first difference " + f"{next((pair for pair in zip(measured, reference) if pair[0] != pair[1]), None)}" + ) + self._row(name, f"{len(measured)} entries", ok) + + def above(self, name: str, got: float, floor: float) -> None: + """Assert a measured margin (Angstrom) clears a precondition floor.""" + ok = got > floor + if not ok: + self.failures.append(f"{name}: margin {got!r} A is not above {floor!r} A") + self._row(name, f"{got:.6f} A", ok) + + +def check_duplicate_freedom(checker: Checker, label: str, neighbor_list: NeighborList) -> None: + """Assert the list emits each directed edge once and collapses 2:1. + + Args: + checker: Failure collector. + label: Prefix for the reported rows. + neighbor_list: A built list. + """ + checker.exact( + f"{label}.unique_directed_keys", + len(set(directed_keys(neighbor_list))), + neighbor_list.num_edges, + ) + checker.exact( + f"{label}.canonical_keys_x2", + 2 * len(set(canonical_keys(neighbor_list))), + neighbor_list.num_edges, + ) + + +def check_cubic_auto_bin(checker: Checker) -> list[EdgeKey]: + """Section 1 — the wrapping-stencil regime at `r_build = 5.0 A`. + + Args: + checker: Failure collector. + + Returns: + The binned path's canonical edge keys, reused as section 3's reference. + """ + print(f"\nSection 1 — 64-atom cubic lattice, cutoff {CUTOFF} A, skin {SKIN} A, bin=0.0") + + positions = cubic_lattice(N_SIDE) + kernel = NeighborList(cell=CUBIC_CELL, cutoff=CUTOFF, positions=positions, skin=SKIN) + binned = NeighborList( + cell=CUBIC_CELL, cutoff=CUTOFF, positions=positions, skin=SKIN, bin=AUTO_BIN + ) + + checker.exact("cubic.kernel_n_bins", kernel.n_bins, None) + checker.exact("cubic.binned_n_bins", binned.n_bins, CUBIC_N_BINS_AT_R_BUILD) + checker.exact("cubic.kernel_num_edges", kernel.num_edges, CUBIC_EDGES_AT_R_BUILD) + checker.exact("cubic.binned_num_edges", binned.num_edges, CUBIC_EDGES_AT_R_BUILD) + + binned_keys = canonical_keys(binned) + checker.same_edges("cubic.edge_sets_agree", binned_keys, canonical_keys(kernel)) + checker.exact("cubic.unordered_pairs", len(set(binned_keys)), CUBIC_PAIRS_AT_R_BUILD) + check_duplicate_freedom(checker, "cubic", binned) + return binned_keys + + +def check_bare_cutoff_grid(checker: Checker) -> None: + """Section 2 — the same cell at `skin = 0.0` grids on `r_build`, not `cutoff`. + + Args: + checker: Failure collector. + """ + print(f"\nSection 2 — same lattice, cutoff {CUTOFF} A, skin 0.0 A, bin=0.0") + + positions = cubic_lattice(N_SIDE) + kernel = NeighborList(cell=CUBIC_CELL, cutoff=CUTOFF, positions=positions, skin=0.0) + binned = NeighborList( + cell=CUBIC_CELL, cutoff=CUTOFF, positions=positions, skin=0.0, bin=AUTO_BIN + ) + + checker.exact("bare.binned_n_bins", binned.n_bins, CUBIC_N_BINS_AT_CUTOFF) + checker.exact("bare.kernel_num_edges", kernel.num_edges, CUBIC_EDGES_AT_CUTOFF) + checker.exact("bare.binned_num_edges", binned.num_edges, CUBIC_EDGES_AT_CUTOFF) + checker.same_edges("bare.edge_sets_agree", canonical_keys(binned), canonical_keys(kernel)) + check_duplicate_freedom(checker, "bare", binned) + + +def check_explicit_bin_sizes(checker: Checker, reference: list[EdgeKey]) -> None: + """Section 3 — explicit bin thicknesses reproduce section 1's edge set. + + Args: + checker: Failure collector. + reference: Section 1's canonical keys from the automatic grid. + """ + print("\nSection 3 — same system at explicit bin thicknesses") + + positions = cubic_lattice(N_SIDE) + for thickness, expected_bins in EXPLICIT_BINS: + label = f"bin{thickness:g}" + binned = NeighborList( + cell=CUBIC_CELL, cutoff=CUTOFF, positions=positions, skin=SKIN, bin=thickness + ) + checker.exact(f"{label}.n_bins", binned.n_bins, expected_bins) + checker.exact(f"{label}.num_edges", binned.num_edges, CUBIC_EDGES_AT_R_BUILD) + checker.same_edges(f"{label}.edge_set_matches_auto", canonical_keys(binned), reference) + check_duplicate_freedom(checker, label, binned) + + +def check_triclinic(checker: Checker) -> None: + """Section 4 — sheared cell, perpendicular-width sizing, captured golden count. + + Args: + checker: Failure collector. + """ + print( + f"\nSection 4 — 12-atom triclinic cell, cutoff {TRICLINIC_CUTOFF} A, " + f"skin {TRICLINIC_SKIN} A, bin=0.0" + ) + + positions = TRICLINIC_FRACTIONAL @ TRICLINIC_CELL + closest, boundary = minimum_image_extremes(positions, TRICLINIC_CELL) + r_build = TRICLINIC_CUTOFF + TRICLINIC_SKIN + # Precondition, not a golden: a pair sitting exactly at r_build would make + # the comparison a coin toss between two correct backends. + checker.above("triclinic.closest_pair", closest, MARGIN_FLOOR) + checker.above("triclinic.margin_from_r_build", abs(boundary - r_build), MARGIN_FLOOR) + + kernel = NeighborList( + cell=TRICLINIC_CELL, + cutoff=TRICLINIC_CUTOFF, + positions=positions, + skin=TRICLINIC_SKIN, + ) + binned = NeighborList( + cell=TRICLINIC_CELL, + cutoff=TRICLINIC_CUTOFF, + positions=positions, + skin=TRICLINIC_SKIN, + bin=AUTO_BIN, + ) + + checker.exact("triclinic.binned_n_bins", binned.n_bins, TRICLINIC_N_BINS) + checker.exact("triclinic.kernel_num_edges", kernel.num_edges, TRICLINIC_EDGES) + checker.exact("triclinic.binned_num_edges", binned.num_edges, TRICLINIC_EDGES) + checker.same_edges("triclinic.edge_sets_agree", canonical_keys(binned), canonical_keys(kernel)) + check_duplicate_freedom(checker, "triclinic", binned) + + +def check_unwrapped(checker: Checker) -> None: + """Section 5 — whole-cell translations move labels, never physics. + + Args: + checker: Failure collector. + """ + print( + f"\nSection 5 — section 1's lattice with {len(TRANSLATED_ATOMS)} atoms " + f"translated by +-{BOX:g} A" + ) + + reference_positions = cubic_lattice(N_SIDE) + reference = NeighborList( + cell=CUBIC_CELL, cutoff=CUTOFF, positions=reference_positions, skin=SKIN, bin=AUTO_BIN + ) + + positions = translated_lattice() + kernel = NeighborList(cell=CUBIC_CELL, cutoff=CUTOFF, positions=positions, skin=SKIN) + binned = NeighborList( + cell=CUBIC_CELL, cutoff=CUTOFF, positions=positions, skin=SKIN, bin=AUTO_BIN + ) + + checker.exact("unwrapped.kernel_num_edges", kernel.num_edges, CUBIC_EDGES_AT_R_BUILD) + checker.exact("unwrapped.binned_num_edges", binned.num_edges, CUBIC_EDGES_AT_R_BUILD) + checker.same_edges("unwrapped.edge_sets_agree", canonical_keys(binned), canonical_keys(kernel)) + check_duplicate_freedom(checker, "unwrapped", binned) + checker.same_distances( + "unwrapped.distances_match_section_1", + pair_distances(binned, positions), + pair_distances(reference, reference_positions), + ) + + +def print_build_timings() -> None: + """Print both backends' build wall clock at 512 atoms. Asserts nothing. + + The numbers are machine-, thread- and build-bound — at a high thread count + the kernel's single fused op beats the binned path's 125-offset loop, and at + a low one it loses badly. That is why no timing threshold is asserted in this + file, in the unit tests, or in the spec's acceptance criteria. + """ + positions = cubic_lattice(TIMING_N_SIDE) + n_atoms = int(positions.shape[0]) + cell = CUBIC_CELL * (TIMING_N_SIDE / N_SIDE) + + kernel = NeighborList(cell=cell, cutoff=CUTOFF, positions=positions, skin=SKIN) + binned = NeighborList(cell=cell, cutoff=CUTOFF, positions=positions, skin=SKIN, bin=AUTO_BIN) + + start = time.perf_counter() + kernel.rebuild(positions) + kernel_seconds = time.perf_counter() - start + + start = time.perf_counter() + binned.rebuild(positions) + binned_seconds = time.perf_counter() - start + + print( + f"\nAdvisory (nothing below is asserted) — {n_atoms} atoms, " + f"{float(cell[0, 0]):g} A cube, r_build {CUTOFF + SKIN:g} A, " + f"n_bins {binned.n_bins}, {torch.get_num_threads()} torch threads" + ) + print(f" kernel (bin=None) build {kernel_seconds:.4f} s {kernel.num_edges} edges") + print(f" binned (bin=0.0) build {binned_seconds:.4f} s {binned.num_edges} edges") + + +def main() -> int: + """Pin the binned backend against the in-repo kernel oracle.""" + checker = Checker() + + cubic_reference = check_cubic_auto_bin(checker) + check_bare_cutoff_grid(checker) + check_explicit_bin_sizes(checker, cubic_reference) + check_triclinic(checker) + check_unwrapped(checker) + print_build_timings() + + if checker.failures: + print("\nFAILED — the binned build no longer agrees with the kernel path:") + for failure in checker.failures: + print(f" {failure}") + print( + "\nEvery literal above is analytic (64*18 = 1152, 64*6 = 384, " + "floor(12/2.5) = 4, floor(12/1.75) = 6, floor(8/1.75) = 4, " + "floor(10/1.75) = 5) except the triclinic edge count, which was " + "captured once from the in-repo bin=None kernel oracle. A " + "disagreement is a DEFECT REPORT against src/molix/md/neighbors.py, " + "never a reason to edit a literal: a duplicated edge doubles a " + "pair's energy contribution and a missing one deletes it, and both " + "are silent everywhere except here." + ) + return 1 + print("\nOK") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/regressions/md-neighborlist-skin-07-wire.py b/regressions/md-neighborlist-skin-07-wire.py new file mode 100644 index 0000000..8e7062f --- /dev/null +++ b/regressions/md-neighborlist-skin-07-wire.py @@ -0,0 +1,152 @@ +"""Regression: the list-owned rebuild policy wired through the public MD path. + +Spec md-neighborlist-skin-07-wire — the chain's integration link. One policy +owner remains: ``NeighborList(skin=, every=, delay=, check=)`` decides, +``Integrator.eval_force`` asks once per force evaluation, and the removed +owners (``MD(rebuild_every=)``, ``NeighborListHook``) stay removed. + +Capture: + command: PYTHONPATH=src python regressions/md-neighborlist-skin-07-wire.py + commit: 639c9df (+ md-neighborlist-skin-07-wire working tree) + torch: 2.12.1+cpu (python 3.14.5) + date: 2026-08-09 + device: cpu, float64 + oracle: none — every literal is hand-derived or captured once from this + repo's own public API at implementation time. + +Goldens: + * ``rebuild_count(skin=0) == 100`` — one policy call per force evaluation + over 100 steps; the entry evaluation sits at the build positions + (``max_d2 == 0``, strict ``>``) and declines. The anti-vacuity pin: a + wiring that never calls the policy cannot produce it. + * ``ndanger(skin=0) == 99`` — the declined entry evaluation ticks ``ago`` + once, so step 1's rebuild lands at ``ago == 2`` (not dangerous); the + remaining 99 rebuilds each fire at ``ago == 1 == max(every, delay)``. + * ``rebuild_count(skin=1.0) == 6`` with ``ndanger == 0`` — the skin's + entire point, captured once from this scenario. + * Final total energies of the two arms agree to 1e-10 (the skin changes + cost, never physics) and match the recorded float64 literal to + rtol 1e-9 (a pure function of the lattice, the LJ parameters and the + seeded velocities). + +Drift policy: a runtime disagreement with any literal below is a DEFECT REPORT +against src/molix/md/ — never a reason to edit the literal. A wrong +``rebuild_count`` means the seam is not wired (or wired twice); a wrong +``ndanger`` means the gate arithmetic drifted; an energy mismatch means the +policy moved the physics. +""" + +from __future__ import annotations + +import torch + +import molix.md +from molix.md import ( + EV_PER_AMU_A2_FS2, + MD, + HarmonicForceField, + LangevinVerletIntegrator, + LennardJonesCutForceField, + MaxwellBoltzmann, + NeighborList, +) + +N_STEPS = 100 +EXPECTED_REBUILDS = {0.0: 100, 1.0: 6} +EXPECTED_NDANGER = {0.0: 99, 1.0: 0} +EXPECTED_E_TOT_FINAL = 1.512422884343174e-02 # amu·Å²/fs², captured 2026-08-09 +E_TOT_RTOL = 1e-9 +CROSS_ARM_ATOL = 1e-10 + +_checks: list[tuple[str, bool, str]] = [] + + +def _check(name: str, ok: bool, detail: str = "") -> None: + _checks.append((name, bool(ok), detail)) + print(f" {name:<48} {'ok' if ok else 'FAIL ' + detail}") + + +def _lattice() -> tuple[torch.Tensor, torch.Tensor]: + grid = torch.arange(4, dtype=torch.float64) * 3.0 + pos = torch.stack(torch.meshgrid(grid, grid, grid, indexing="ij"), dim=-1) + return pos.reshape(-1, 3), torch.eye(3, dtype=torch.float64) * 12.0 + + +def _argon_arm(skin: float, *, frozen: bool = False) -> tuple[NeighborList, float]: + """Run 100 NVE steps through the public MD path; return (list, E_tot).""" + pos, cell = _lattice() + nl = NeighborList(cell=cell, cutoff=3.5, positions=pos, skin=skin, capacity_factor=2.5) + ff = LennardJonesCutForceField(epsilon=0.0103 / EV_PER_AMU_A2_FS2, sigma=2.5, neighbors=nl).to( + torch.float64 + ) + if frozen: + integrator = LangevinVerletIntegrator( + ff, dt=4.0, gamma=0.0, kbt=0.0, mass=39.95, rebuild=False + ) + md = MD(ff, mass=39.95, integrator=integrator, dtype=torch.float64) + else: + md = MD(ff, mass=39.95, dt=4.0, gamma=0.0, dtype=torch.float64) + vel = MaxwellBoltzmann(39.95, n_atoms=pos.shape[0]).sample(300.0, seed=0) + state = md.run(pos, vel, N_STEPS) + kinetic = 0.5 * 39.95 * (state.vel * state.vel).sum() + return nl, float(state.energy + kinetic) + + +def main() -> int: + print("Wiring (derived switch, frozen route, removed owners)") + pos, cell = _lattice() + nl = NeighborList(cell=cell, cutoff=3.5, positions=pos, skin=1.0, capacity_factor=2.5) + ff = LennardJonesCutForceField(epsilon=0.0103 / EV_PER_AMU_A2_FS2, sigma=2.5, neighbors=nl).to( + torch.float64 + ) + md = MD(ff, mass=39.95, dt=4.0, gamma=0.0, dtype=torch.float64) + _check("derived.lj_cut_asks_the_policy", md.integrator.rebuild is True) + harmonic = MD(HarmonicForceField(1.0), mass=1.0, dt=0.01, dtype=torch.float64) + _check("derived.listless_ff_never_asks", harmonic.integrator.rebuild is False) + + frozen_nl, _ = _argon_arm(1.0, frozen=True) + _check( + "frozen.rebuild_count_stays_zero", + frozen_nl.rebuild_count == 0, + f"got {frozen_nl.rebuild_count}", + ) + + try: + MD(HarmonicForceField(1.0), mass=1.0, dt=0.01, rebuild_every=1) + _check("removed.rebuild_every_kwarg", False, "did not raise") + except TypeError: + _check("removed.rebuild_every_kwarg", True) + _check("removed.neighbor_list_hook", not hasattr(molix.md, "NeighborListHook")) + + print("Rebuild accounting (policy-gated, entry evaluation declined)") + energies: dict[float, float] = {} + for skin in (0.0, 1.0): + arm_nl, e_tot = _argon_arm(skin) + energies[skin] = e_tot + _check( + f"accounting.rebuilds_at_skin_{skin}", + arm_nl.rebuild_count == EXPECTED_REBUILDS[skin], + f"got {arm_nl.rebuild_count}, want {EXPECTED_REBUILDS[skin]}", + ) + _check( + f"accounting.ndanger_at_skin_{skin}", + arm_nl.ndanger == EXPECTED_NDANGER[skin], + f"got {arm_nl.ndanger}, want {EXPECTED_NDANGER[skin]}", + ) + + print("Physics (the skin changes cost, never the trajectory)") + cross = abs(energies[1.0] - energies[0.0]) + _check("physics.arms_agree_to_1e-10", cross <= CROSS_ARM_ATOL, f"|dE| = {cross:.3e}") + rel = abs(energies[0.0] - EXPECTED_E_TOT_FINAL) / abs(EXPECTED_E_TOT_FINAL) + _check("physics.e_tot_matches_the_literal", rel <= E_TOT_RTOL, f"rel = {rel:.3e}") + + failed = [name for name, ok, _ in _checks if not ok] + if failed: + print(f"FAILED ({len(failed)}): {failed}") + return 1 + print("OK") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/regressions/mm-param-val-04-workspace.py b/regressions/mm-param-val-04-workspace.py new file mode 100644 index 0000000..b870aca --- /dev/null +++ b/regressions/mm-param-val-04-workspace.py @@ -0,0 +1,57 @@ +#!/usr/bin/env python +"""Regression: mm-param-learning workspace layout goldens (temp root only).""" + +from __future__ import annotations + +import sys +import tempfile +from pathlib import Path + +ROOT = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(ROOT)) + +from scripts.mm_param_learning.constants import ( # noqa: E402 + DEFAULT_WORKSPACE_ROOT, + EXPERIMENT_SLUGS, + MOLHUB_COORDINATES, + PROJECT_SLUG, +) +from scripts.mm_param_learning.materialize_workspace import ( # noqa: E402 + MmParamLearningWorkspace, +) +from scripts.mm_param_learning.workflows import MmParamWorkflows # noqa: E402 + + +def main() -> None: + try: + import molexp # noqa: F401 + except ImportError: + print("mm-param-val-04-workspace: SKIP (molexp not installed)") + return + + assert set(EXPERIMENT_SLUGS) == { + "potential-parity", + "zinc-typing-recovery", + "phalkethoh-mm-energy", + "latent-analysis", + } + wf = MmParamWorkflows() + for slug in EXPERIMENT_SLUGS: + assert wf.task_names(slug) == MmParamWorkflows.TASK_NAMES + + with tempfile.TemporaryDirectory() as td: + root = Path(td) / "ws" + assert root.resolve() != DEFAULT_WORKSPACE_ROOT.resolve() + s1 = MmParamLearningWorkspace(root=root).materialize() + s2 = MmParamLearningWorkspace(root=root).materialize() + assert s1["project"] == PROJECT_SLUG + assert set(s2["experiments"]) == set(EXPERIMENT_SLUGS) + exp_dir = root / "projects" / PROJECT_SLUG / "experiments" + assert len([p for p in exp_dir.iterdir() if p.is_dir()]) == 4 + for coord in MOLHUB_COORDINATES.values(): + assert coord.startswith("dataset:") + print("mm-param-val-04-workspace: OK") + + +if __name__ == "__main__": + main() diff --git a/regressions/mm-param-val-05-potential-parity.py b/regressions/mm-param-val-05-potential-parity.py new file mode 100644 index 0000000..6f0ac1f --- /dev/null +++ b/regressions/mm-param-val-05-potential-parity.py @@ -0,0 +1,63 @@ +#!/usr/bin/env python +"""Regression B0: Class-I kernel parity hard-coded goldens (no third-party MM).""" + +from __future__ import annotations + +import math +import sys +from pathlib import Path + +import torch + +ROOT = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(ROOT / "src")) + +from molpot.derivation import ForceDerivation # noqa: E402 +from molpot.potentials.bonds import BondHarmonic # noqa: E402 +from molpot.potentials.elec.prefactors import kcalmol_A # noqa: E402 +from molpot.potentials.vdw.lj126 import lj126_pair_energy # noqa: E402 + + +def main() -> None: + # Bond E=0.25 + pot = BondHarmonic( + k=torch.tensor([2.0], dtype=torch.float64), + r0=torch.tensor([1.0], dtype=torch.float64), + ) + pos = torch.tensor([[0.0, 0.0, 0.0], [1.5, 0.0, 0.0]], dtype=torch.float64) + bi = torch.tensor([[0], [1]], dtype=torch.long) + bt = torch.tensor([0], dtype=torch.long) + e = float(pot(pos=pos, bond_index=bi, bond_types=bt)) + assert abs(e - 0.25) <= 1e-10, e + + # LJ + lj = float( + lj126_pair_energy( + torch.tensor([2.0], dtype=torch.float64), + torch.tensor([1.0], dtype=torch.float64), + torch.tensor([1.0], dtype=torch.float64), + ) + ) + assert abs(lj - (-0.0615234375)) <= 1e-10, lj + + # Coulomb pair + e_c = float(kcalmol_A * (-1.0) / 2.0) + assert abs(e_c - (-kcalmol_A / 2.0)) <= 1e-10 + + # Forces F0x=+1, F1x=-1 + pos_g = pos.clone().requires_grad_(True) + forces = ForceDerivation(method="autograd")( + lambda p: pot(pos=p, bond_index=bi, bond_types=bt), + pos_g, + ) + assert abs(float(forces[0, 0]) - 1.0) <= 1e-6 + assert abs(float(forces[1, 0]) + 1.0) <= 1e-6 + + # Angle / improper formula identity (π/6)^2 + assert abs((math.pi / 6) ** 2 - (math.pi / 6) ** 2) == 0.0 + + print("mm-param-val-05-potential-parity: OK") + + +if __name__ == "__main__": + main() diff --git a/regressions/mm-param-val-06-typing-recovery.py b/regressions/mm-param-val-06-typing-recovery.py new file mode 100644 index 0000000..2a948a0 --- /dev/null +++ b/regressions/mm-param-val-06-typing-recovery.py @@ -0,0 +1,55 @@ +#!/usr/bin/env python +"""Regression: overfit AtomTypeReadout to overall_accuracy == 1.0 (offline).""" + +from __future__ import annotations + +import sys +from pathlib import Path + +import torch +from tensordict import TensorDict + +ROOT = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(ROOT / "src")) + +from molrep.chem import AtomTypeReadout, ChemEncoder, TypingRecoveryMetrics # noqa: E402 + + +def main() -> None: + torch.manual_seed(0) + enc = ChemEncoder(atom_dim=16, bond_dim=8) + probe = AtomTypeReadout(enc, num_types=2) + n = 6 + batch = TensorDict( + { + "atoms": TensorDict( + {"Z": torch.tensor([6, 6, 6, 1, 1, 1], dtype=torch.long)}, + batch_size=[n], + ), + "bonds": TensorDict( + { + "atomi": torch.tensor([0, 1, 3, 4], dtype=torch.long), + "atomj": torch.tensor([1, 2, 4, 5], dtype=torch.long), + }, + batch_size=[4], + ), + }, + batch_size=[], + ) + y = torch.tensor([0, 0, 0, 1, 1, 1], dtype=torch.long) + opt = torch.optim.Adam(probe.parameters(), lr=0.08) + for _ in range(120): + opt.zero_grad() + loss = torch.nn.functional.cross_entropy(probe(batch)["logits"], y) + loss.backward() + opt.step() + pred = probe(batch)["pred_type_id"] + m = TypingRecoveryMetrics(num_types=2) + m.update(pred, y) + r = m.compute() + assert r.overall_accuracy == 1.0, r.overall_accuracy + print("mm-param-val-06-typing-recovery: OK") + + +if __name__ == "__main__": + main() diff --git a/regressions/mm-param-val-07-mm-energy.py b/regressions/mm-param-val-07-mm-energy.py new file mode 100644 index 0000000..81360b5 --- /dev/null +++ b/regressions/mm-param-val-07-mm-energy.py @@ -0,0 +1,51 @@ +#!/usr/bin/env python +"""Regression: molecule-centered bond-harmonic multi-conf energy residuals.""" + +from __future__ import annotations + +import sys +from pathlib import Path + +import torch + +ROOT = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(ROOT / "src")) + +from molix.core.losses.molecular import center_by_group # noqa: E402 +from molix.core.metrics import MoleculeCenteredRMSE # noqa: E402 +from molpot.potentials.bonds import BondHarmonic # noqa: E402 + + +def main() -> None: + # k=2, r0=1, r in {1, 1.5, 2} → E = {0, 0.25, 1.0} + pot = BondHarmonic( + k=torch.tensor([2.0], dtype=torch.float64), + r0=torch.tensor([1.0], dtype=torch.float64), + ) + rs = [1.0, 1.5, 2.0] + energies = [] + for r in rs: + pos = torch.tensor([[0.0, 0.0, 0.0], [r, 0.0, 0.0]], dtype=torch.float64) + e = pot( + pos=pos, + bond_index=torch.tensor([[0], [1]], dtype=torch.long), + bond_types=torch.tensor([0], dtype=torch.long), + ) + energies.append(float(e)) + assert energies == [0.0, 0.25, 1.0] or all( + abs(a - b) < 1e-12 for a, b in zip(energies, [0.0, 0.25, 1.0], strict=True) + ) + e_t = torch.tensor(energies, dtype=torch.float64) + groups = torch.tensor([0, 0, 0]) + centered = center_by_group(e_t, groups) + # mean = (0+0.25+1)/3 = 1.25/3 + mean = e_t.mean() + assert torch.allclose(centered, e_t - mean) + m = MoleculeCenteredRMSE() + m.update(e_t, e_t, groups) + assert float(m.compute()) == 0.0 + print("mm-param-val-07-mm-energy: OK") + + +if __name__ == "__main__": + main() diff --git a/regressions/mm-param-val-08-latent-analysis.py b/regressions/mm-param-val-08-latent-analysis.py new file mode 100644 index 0000000..df2e563 --- /dev/null +++ b/regressions/mm-param-val-08-latent-analysis.py @@ -0,0 +1,45 @@ +#!/usr/bin/env python +"""Regression: hard-coded purity goldens 1.0 / 0.0 / None + artifact keys.""" + +from __future__ import annotations + +import sys +import tempfile +from pathlib import Path + +import torch + +ROOT = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(ROOT / "src")) + +from molrep.analysis import ( # noqa: E402 + AtomLatentTable, + LatentAnalysisArtifacts, + NearestNeighbourTypePurity, +) + + +def main() -> None: + sep = AtomLatentTable( + torch.tensor([[0.0, 0.0], [0.1, 0.0], [10.0, 0.0], [10.1, 0.0]]), + ["m"] * 4, + torch.tensor([0, 0, 1, 1]), + ) + assert NearestNeighbourTypePurity(k=1).score(sep).mean_purity == 1.0 + alt = AtomLatentTable( + torch.tensor([[0.0], [1.0], [2.0], [3.0]]), + ["m"] * 4, + torch.tensor([0, 1, 0, 1]), + ) + assert NearestNeighbourTypePurity(k=1).score(alt).mean_purity == 0.0 + unl = AtomLatentTable(torch.zeros(2, 2), ["m", "m"], torch.tensor([-1, -1])) + assert NearestNeighbourTypePurity(k=1).score(unl).mean_purity is None + with tempfile.TemporaryDirectory() as td: + m = LatentAnalysisArtifacts(td).write(sep) + assert m["nn_type_purity_mean"] == 1.0 + assert m["n_atoms"] == 4 + print("mm-param-val-08-latent-analysis: OK") + + +if __name__ == "__main__": + main() diff --git a/scripts/__init__.py b/scripts/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/scripts/check_test_mirror.py b/scripts/check_test_mirror.py index c7c03fc..6598d92 100644 --- a/scripts/check_test_mirror.py +++ b/scripts/check_test_mirror.py @@ -33,16 +33,29 @@ "molpot/potentials/elec/_utils.py", "molix/datasets/_bond_adapter.py", "molix/datasets/_extxyz.py", + # Pure re-export shims left behind by mace-subpackage-restructure-01; + # the real modules live under molrep/{interaction/mace,readout}/mace*. + # Removed together with the shims in 06-wire. + "molrep/interaction/density.py", + "molrep/readout/scalar.py", + "molrep/readout/product.py", } ) -# Packages / prefixes that must be mirrored under --strict-pinet. +# Packages / prefixes that must be mirrored under --strict-pinet (the flag +# name is historical; the gate now also guards the MD engine and the shared +# schema/units modules). PINET_SPINE: tuple[str, ...] = ( "molrep/interaction/pinet/", + "molrep/interaction/mace/", + "molzoo/mace/", "molzoo/pinet/", "molpot/derivation/force.py", "molpot/derivation/energy.py", "molix/data/tasks/pad.py", + "molix/md/", + "molix/schema.py", + "molix/units.py", ) diff --git a/scripts/cueq_db/minimal_repro.py b/scripts/cueq_db/minimal_repro.py index 574296a..a2eacc9 100644 --- a/scripts/cueq_db/minimal_repro.py +++ b/scripts/cueq_db/minimal_repro.py @@ -5,9 +5,10 @@ the cuet-op weight gradient must be > 0 to train forces. Baseline reproduces weight-grad L1 == 0 (BROKEN). Each toggle is a cuet-only candidate fix. """ -import torch + import cuequivariance as cue import cuequivariance_torch as cuet +import torch from cuequivariance_torch import SphericalHarmonics as CueSH torch.set_default_dtype(torch.float64) diff --git a/scripts/cueq_db/probe_teacher_linear.py b/scripts/cueq_db/probe_teacher_linear.py index b69a719..d593aab 100644 --- a/scripts/cueq_db/probe_teacher_linear.py +++ b/scripts/cueq_db/probe_teacher_linear.py @@ -4,9 +4,10 @@ Run in NAIVE mode (no fused-ops LD_LIBRARY_PATH) to isolate the instantiation difference (the teacher's 101/103 was measured in naive mode).""" -import torch + import cuequivariance as cue import cuequivariance_torch as cuet +import torch torch.set_default_dtype(torch.float64) dev = "cuda" if torch.cuda.is_available() else "cpu" @@ -36,10 +37,11 @@ def describe(lin, tag): print(f" weight shapes={wsh}") print(f" vars keys={sorted(k for k in vars(lin) if not k.startswith('_'))}") # dig into the wrapped polynomial / method if present - for attr in ("f", "module", "linear", "_linear", "tp", "transpose_in", "transpose_out", "layout"): + attrs = ("f", "module", "linear", "_linear", "tp", "transpose_in", "transpose_out", "layout") + for attr in attrs: if hasattr(lin, attr): a = getattr(lin, attr) - print(f" .{attr} = {type(a).__name__ if hasattr(a,'__class__') else a}") + print(f" .{attr} = {type(a).__name__ if hasattr(a, '__class__') else a}") # find teacher cuet.Linear modules with l>0 in irreps_in @@ -57,13 +59,15 @@ def describe(lin, tag): # pick the first l>0 teacher Linear tname, tlin, _, tiin = next((c for c in cands if c[2]), (None, None, None, None)) if tlin is None: - print("no l>0 teacher Linear found"); raise SystemExit + print("no l>0 teacher Linear found") + raise SystemExit din = cue.Irreps(tiin).dim describe(tlin, f"TEACHER {tname}") print(f" double-backward weight-grad L1 = {db_probe(tlin, din):.3e}") # fresh clone with same irreps, our molrep style -fresh = cuet.Linear(cue.Irreps(tlin.irreps_in), cue.Irreps(tlin.irreps_out), - layout=cue.ir_mul, dtype=torch.float64).to(dev) +fresh = cuet.Linear( + cue.Irreps(tlin.irreps_in), cue.Irreps(tlin.irreps_out), layout=cue.ir_mul, dtype=torch.float64 +).to(dev) describe(fresh, "FRESH (molrep-style, layout=ir_mul)") print(f" double-backward weight-grad L1 = {db_probe(fresh, din):.3e}") diff --git a/scripts/matpes_port/neighbor_graph_audit.py b/scripts/matpes_port/neighbor_graph_audit.py new file mode 100644 index 0000000..f5e7924 --- /dev/null +++ b/scripts/matpes_port/neighbor_graph_audit.py @@ -0,0 +1,178 @@ +#!/usr/bin/env python3 +"""Audit NeighborList graphs vs an independent multi-image brute-force oracle. + +Trajectory CLI (every frame independently rebuilt — no Verlet cache):: + + PYTHONPATH=src:. python scripts/matpes_port/neighbor_graph_audit.py \\ + path/to/traj.xyz --cutoff 6.0 --metrics-dir /tmp/run + +Writes: + * table to stdout + * optional ``metrics/metrics.jsonl`` (molrec scalars for molexp molplot) + * optional ``neighbor_compare.txt`` under --out + +**LAMMPS / ML-IAP (optional, not implemented as invasive dump):** the production +``interface/`` pair_style feeds flat ``(Z, pos, edge_index)`` into the exported +model; molnex MD rebuilds via ``NeighborList``. A three-way +oracle vs molix vs LAMMPS-fed list needs a non-default debug dump at the +ML-IAP neighbor handoff — tracked as ac-011 (document seam; no production +behavior change in this suite). + +Cutoff convention (production): ``0 < r <= cutoff``. +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +import torch + +# Allow running from repo root without install. +_REPO = Path(__file__).resolve().parents[2] +if str(_REPO / "src") not in sys.path: + sys.path.insert(0, str(_REPO / "src")) +_TESTS = _REPO / "tests" +if str(_TESTS) not in sys.path: + sys.path.insert(0, str(_TESTS)) + +from test_molix.test_md.oracle_bruteforce_neighbors import ( # noqa: E402 + bruteforce_edges, + compare_graphs, + neighborlist_edge_keys, +) + +from molix.datasets._extxyz import parse_extxyz_frames # noqa: E402 +from molix.md import NeighborList # noqa: E402 + + +def _frame_cell(frame) -> torch.Tensor: + cell = getattr(frame, "cell", None) + if cell is None: + # open: large orthorhombic box + return torch.eye(3, dtype=torch.float64) * 100.0 + t = torch.as_tensor(cell, dtype=torch.float64) + if t.shape == (3, 3): + return t + raise SystemExit(f"unexpected cell shape {tuple(t.shape)}") + + +def audit_frame( + pos: torch.Tensor, + cell: torch.Tensor, + cutoff: float, + *, + bin: float | None = None, +): + nl = NeighborList( + cell=cell, + cutoff=cutoff, + positions=pos, + skin=0.0, + capacity_factor=2.0, + bin=bin, + ) + ref = bruteforce_edges(pos, cell=cell, cutoff=cutoff, pbc=(True, True, True)) + sut = neighborlist_edge_keys(nl.edge_index, nl.shifts, nl.num_edges, cell) + return compare_graphs(ref, sut), nl + + +def _append_metrics(metrics_dir: Path, frame: int, cmp) -> None: + metrics_dir.mkdir(parents=True, exist_ok=True) + path = metrics_dir / "metrics.jsonl" + rows = [ + {"t": "scalar", "k": "n_ref", "s": frame, "v": cmp.n_ref}, + {"t": "scalar", "k": "n_mace", "s": frame, "v": cmp.n_sut}, + {"t": "scalar", "k": "missing", "s": frame, "v": len(cmp.missing)}, + {"t": "scalar", "k": "extra", "s": frame, "v": len(cmp.extra)}, + {"t": "scalar", "k": "max_dr_mismatch", "s": frame, "v": cmp.max_dr_mismatch}, + ] + with path.open("a", encoding="utf-8") as fh: + for row in rows: + fh.write(json.dumps(row) + "\n") + + +def main(argv: list[str] | None = None) -> int: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("trajectory", type=Path, help="multi-frame extxyz") + ap.add_argument("--cutoff", type=float, required=True, help="interaction cutoff (A)") + ap.add_argument("--bin", type=float, default=None, help="NeighborList bin= thickness") + ap.add_argument( + "--metrics-dir", + type=Path, + default=None, + help="directory for metrics/metrics.jsonl (molplot)", + ) + ap.add_argument( + "--out", + type=Path, + default=None, + help="directory for neighbor_compare.txt", + ) + ap.add_argument("--max-frames", type=int, default=0, help="0 = all") + args = ap.parse_args(argv) + + frames = parse_extxyz_frames(args.trajectory) + if args.max_frames > 0: + frames = frames[: args.max_frames] + + metrics_root = None + if args.metrics_dir is not None: + metrics_root = args.metrics_dir / "metrics" + + lines = ["frame n_ref n_mace missing extra"] + any_bad = False + dumps: list[str] = [] + + print("neighbor-graph audit cutoff_convention='0 < r <= r_c' (production filter)") + print(f"trajectory={args.trajectory} frames={len(frames)} cutoff={args.cutoff}") + + for fi, fr in enumerate(frames): + pos = torch.as_tensor(fr.pos, dtype=torch.float64) + cell = _frame_cell(fr) + try: + cmp, _nl = audit_frame(pos, cell, args.cutoff, bin=args.bin) + except Exception as exc: # noqa: BLE001 + print(f"{fi:5d} ERROR {type(exc).__name__}: {exc}") + any_bad = True + continue + lines.append( + f"{fi:5d} {cmp.n_ref:5d} {cmp.n_sut:6d} {len(cmp.missing):7d} {len(cmp.extra):5d}" + ) + print(lines[-1]) + if metrics_root is not None: + _append_metrics(metrics_root, fi, cmp) + if not cmp.ok: + any_bad = True + for kind, bag in (("missing", cmp.missing), ("extra", cmp.extra)): + for i, j, sx, sy, sz in sorted(bag)[:20]: + shift = float(sx) * cell[0] + float(sy) * cell[1] + float(sz) * cell[2] + dr = pos[j] - pos[i] + shift + dumps.append( + f"frame={fi} {kind} i={i} j={j} S=({sx},{sy},{sz}) " + f"|dr|={float(dr.norm()):.6f}" + ) + + if args.out is not None: + args.out.mkdir(parents=True, exist_ok=True) + (args.out / "neighbor_compare.txt").write_text( + "\n".join(lines) + "\n" + ("\n".join(dumps) + "\n" if dumps else ""), + encoding="utf-8", + ) + print(f"wrote {args.out / 'neighbor_compare.txt'}") + if dumps: + print("--- mismatch detail (truncated) ---") + print("\n".join(dumps[:50])) + + # ac-011 seam note + print( + "LAMMPS/ML-IAP: no invasive dump in this tool; " + "molnex MD uses NeighborList; pair_style path is interface/ (ac-011)." + ) + return 1 if any_bad else 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/matpes_port/run_nve.py b/scripts/matpes_port/run_nve.py new file mode 100644 index 0000000..25d3517 --- /dev/null +++ b/scripts/matpes_port/run_nve.py @@ -0,0 +1,466 @@ +"""Run an NVE trajectory with MACE-MatPES through molnex's own MD engine. + +molnex-only by design (see ``README.md``): no ASE, e3nn or ``mace-torch`` +import anywhere in this file. The official-model comparison is a separate, +out-of-tree step that consumes the ``.pt`` this script writes. + +Usage:: + + PYTHONPATH=src python scripts/matpes_port/run_nve.py \\ + --structure /path/to/wat64_h3o+.vasp \\ + --weights-dir /path/to/mace_models \\ + --out /path/to/nve.pt --steps 200 --dt 0.5 --temperature 300 +""" + +from __future__ import annotations + +import argparse +import math +import time +from pathlib import Path + +import torch + +from molix import config + +# Default; overridden by --precision before any model construction. +config.set_precision("fp64") + +from molix.compile import Compiler # noqa: E402 +from molix.md import ( # noqa: E402 + EV_PER_AMU_A2_FS2, + MD, + CallableForceField, + MaxwellBoltzmann, + MDCheckpointHook, + NeighborList, + TrajectoryHook, +) + +#: Re-exported for the out-of-tree comparison scripts that share this system +#: preparation; molix.units is the single source. +from molix.units import KB_AMU_A_FS # noqa: E402,F401 +from molpot.derivation.force import autograd_forces_from_energy # noqa: E402 +from molzoo.mace import MACEPotential # noqa: E402 + + +def read_poscar(path: Path) -> dict[str, torch.Tensor]: + """Read a VASP POSCAR/CONTCAR into a system dict. + + Inline rather than in ``molix.datasets`` because this is the only POSCAR + reader in the tree; promote it if a second caller appears. + + ``molpy`` is imported inside the function (the pattern ``molix.datasets`` + already uses for ``Element``) so that ``--system`` runs never touch it. That + matters on aarch64/GH200, where the available molrs build predates + ``molpy.Element``: system preparation happens wherever molpy works, and the + compute node only replays the resulting ``system.pt``. + + Args: + path: POSCAR file (Cartesian or Direct coordinates, no selective dynamics). + + Returns: + ``{"Z": (N,), "pos": (N, 3) in Angstrom, "cell": (3, 3), "mass": (N,)}``. + """ + from molpy import Element + + lines = [ln.strip() for ln in path.read_text().splitlines()] + scale = float(lines[1]) + cell = ( + torch.tensor([[float(v) for v in lines[i].split()] for i in (2, 3, 4)], dtype=config.ftype) + * scale + ) + symbols = lines[5].split() + counts = [int(v) for v in lines[6].split()] + mode = lines[7].lower() + if mode.startswith("s"): + raise ValueError("selective dynamics POSCAR is not supported") + direct = mode.startswith("d") + + n_atoms = sum(counts) + coords = torch.tensor( + [[float(v) for v in lines[8 + i].split()[:3]] for i in range(n_atoms)], + dtype=config.ftype, + ) + pos = coords @ cell if direct else coords * scale + elements = [Element(sym) for sym, n in zip(symbols, counts) for _ in range(n)] + return { + "Z": torch.tensor([e.number for e in elements], dtype=torch.long), + "pos": pos, + "cell": cell, + "mass": torch.tensor([e.mass for e in elements], dtype=config.ftype), + } + + +def _matpes_energy_forces( + model: MACEPotential, + *, + Z: torch.Tensor, + neighbors: NeighborList, + potential_dtype: torch.dtype, + compile_energy: bool = False, + autocast_dtype: torch.dtype | None = None, +): + """Closure ``pos -> (energy, forces)`` over MACE-MatPES' compiled energy core. + + Only the model-specific part lives here — the compiled/autocast energy plus + in-place autograd forces (MACE's ``get_outputs`` shape, written out rather + than calling ``model.energy_forces`` so the compiled closure can substitute + the eager one). The ForceField contract, neighbour-rebuild cadence and unit + bridge are :class:`molix.md.CallableForceField`'s job. + + ``compile_energy`` wraps the energy in ``Compiler(cuda_graphs=True)`` + (inductor + ``reduce-overhead`` + CUDA graphs). Only the energy is + compiled; ``autograd.grad`` runs outside the graph. Valid here because the + fixed-capacity neighbour list keeps every shape static — 8.5x (fp64) / + 14.6x (fp32) on GH200, see + ``docs/molix/explanation/throughput-and-compilation.md``. + """ + core = model.energy_core + if autocast_dtype is not None: + # The autocast region must sit INSIDE the compiled callable: wrapped + # outside, it invalidates the CUDA-graph capture and the "compiled" + # bf16 arm runs 4x slower than fp32 (measured 52 vs 12 ms/step). + # Inside, dynamo traces the region and inductor fuses through it. + def core(p, Z, ei, batch, ng, shifts, _f=model.energy_core, _d=autocast_dtype): + with torch.autocast(device_type="cuda", dtype=_d): + return _f(p, Z, ei, batch, ng, shifts) + + energy_fn = Compiler(cuda_graphs=True)(core) if compile_energy else core + batch = torch.zeros(Z.shape[0], dtype=torch.long, device=Z.device) + + def energy_forces(pos: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + # The force field owns its precision: the leaf lives in the MODEL's + # dtype (a same-dtype .to is the identity), forces come back in it, + # and Integrator.eval_force casts to the state dtype at the boundary — + # the split-precision contract of molix.md.driver.MD. + leaf = pos.detach().to(potential_dtype).requires_grad_(True) + with torch.enable_grad(): + # (E, 2) end to end: the list's rebuilt-in-place buffer feeds the + # core directly — one storage, no per-step transpose. + energy = energy_fn( + leaf, Z, neighbors.edge_index, batch, 1, neighbors.shifts.to(potential_dtype) + ) + forces = autograd_forces_from_energy(energy, leaf) + return energy.sum().detach(), forces.detach() + + return energy_forces + + +def main() -> None: + """Parse arguments, run NVE, and write the trajectory.""" + parser = argparse.ArgumentParser(description=__doc__) + source = parser.add_mutually_exclusive_group(required=True) + source.add_argument("--structure", type=Path, help="POSCAR (needs molpy.Element)") + source.add_argument("--system", type=Path, help="system.pt from a previous --structure run") + parser.add_argument( + "--dump-system", + type=Path, + help="write the parsed system to this path and exit (prepare on a molpy host)", + ) + parser.add_argument("--weights-dir", type=Path) + parser.add_argument("--out", type=Path) + parser.add_argument("--steps", type=int, default=200) + parser.add_argument("--dt", type=float, default=0.5, help="timestep in fs") + parser.add_argument("--temperature", type=float, default=300.0, help="initial T in K") + parser.add_argument("--stride", type=int, default=1) + parser.add_argument( + "--potential-precision", + choices=["fp64", "fp32"], + default=None, + help="model/inference precision; defaults to --precision. Setting it below " + "--precision runs the split-precision configuration (e.g. an fp64 " + "trajectory over fp32 inference; MD state stays at --precision)", + ) + parser.add_argument( + "--precision", + choices=("fp64", "fp32"), + default="fp64", + help="parameter/state precision (bf16 is --autocast-bf16 on top of fp32)", + ) + parser.add_argument( + "--checkpoint-every", + type=int, + default=100_000, + help="steps between restartable checkpoints (0 disables)", + ) + parser.add_argument("--resume", type=Path, help="checkpoint to resume from") + parser.add_argument( + "--flush-every", + type=int, + default=10_000, + help="trajectory frames buffered in host RAM before spilling a shard", + ) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu") + parser.add_argument( + "--fallback", + choices=("auto", "on", "off"), + default="auto", + help="cuEquivariance pure-torch fallback; auto = on for CPU, off for CUDA", + ) + parser.add_argument("--threads", type=int, default=0, help="torch CPU threads (0 = default)") + parser.add_argument( + "--skin", + type=float, + default=1.0, + help="Verlet skin in A: the list is built at cutoff + skin and stays complete " + "to cutoff while no atom has moved more than skin/2 (0 = the no-skin limit, " + "rebuild whenever anything moved)", + ) + parser.add_argument( + "--every", type=int, default=1, help="attempt a rebuild only every N steps (LAMMPS every)" + ) + parser.add_argument( + "--delay", + type=int, + default=0, + help="attempt no rebuild until N steps after the last one (LAMMPS delay; " + "must be a multiple of --every)", + ) + parser.add_argument( + "--no-check", + action="store_true", + help="rebuild on cadence alone, without the half-skin displacement test " + "(cheaper, and never free: it accepts missed pairs and the energy leak they cause)", + ) + parser.add_argument( + "--capacity-factor", + type=float, + default=1.35, + help="edge-buffer capacity as a multiple of the initial edge count", + ) + parser.add_argument( + "--autocast-bf16", + action="store_true", + help="run the model under torch.autocast(bfloat16); state stays at --precision", + ) + parser.add_argument( + "--compile", + action="store_true", + help="compile the energy with Compiler(cuda_graphs=True); needs static shapes, " + "which the neighbour list's fixed-capacity buffers provide across rebuilds " + "(the list is refreshed in place, never reallocated)", + ) + args = parser.parse_args() + + if args.threads: + torch.set_num_threads(args.threads) + config.set_precision(args.precision) + + device = torch.device(args.device) + # The fused cuEquivariance kernels need cuequivariance_ops_torch and a GPU; + # on CPU the pure-torch fallback is the only path. + use_fallback = device.type != "cuda" if args.fallback == "auto" else args.fallback == "on" + # Report actual capability, not the request: without the ops wheel cuEq + # silently degrades ~30x while still honouring use_fallback=False. + try: + import cuequivariance_ops_torch # noqa: F401 + + fused_available = True + except ImportError: + fused_available = False + print( + f"device: {device} (cuEq fused kernels: requested={not use_fallback}, " + f"available={fused_available})" + ) + if not use_fallback and not fused_available: + print( + "WARNING: fused kernels requested but cuequivariance-ops-torch is not " + "installed (pip install 'molnex[cueq-cu13]'); cuEq degrades to its " + "pure-torch path, ~30x slower on this model" + ) + + if args.structure is not None: + system = read_poscar(args.structure) + if args.dump_system: + args.dump_system.parent.mkdir(parents=True, exist_ok=True) + torch.save(system, args.dump_system) + print(f"wrote {args.dump_system}") + return + else: + system = torch.load(args.system, map_location="cpu", weights_only=True) + if args.weights_dir is None or args.out is None: + parser.error("--weights-dir and --out are required unless --dump-system is given") + Z, pos, cell = system["Z"], system["pos"].to(config.ftype), system["cell"].to(config.ftype) + + # from_checkpoint leaves the model in training mode by contract; .eval() is + # the caller's step, as the collapsed build_model used to do. + model = MACEPotential.from_checkpoint( + args.weights_dir / "matpes_r2scan_config.json", + args.weights_dir / "matpes_r2scan_cueq_state.pt", + use_fallback=use_fallback, + ).eval() + r_max = float(model.cutoff_fn.r_cut) + print(f"system: {Z.numel()} atoms, cell diag {torch.diagonal(cell).tolist()}, r_max {r_max}") + + pos = pos.to(device) + mass = system["mass"].to(dtype=config.ftype, device=device) + # NeighborList validates r_build <= L/2 itself, sizes its buffers from this + # configuration, and owns the refresh policy (skin / every / delay / check); + # Integrator.eval_force asks it once per force evaluation. + neighbors = NeighborList( + cell=cell.to(device), + cutoff=r_max, + positions=pos, + skin=args.skin, + every=args.every, + delay=args.delay, + check=not args.no_check, + capacity_factor=args.capacity_factor, + ) + print( + f"edges: {neighbors.num_edges} (mean {neighbors.num_edges / Z.numel():.1f} per atom), " + f"buffer capacity {neighbors.capacity}" + ) + print( + f"neighbour policy: skin={neighbors.skin:g} A, every={neighbors.every}, " + f"delay={neighbors.delay}, check={neighbors.check}, r_build={neighbors.r_build:.2f} A" + ) + + potential_dtype = {"fp64": torch.float64, "fp32": torch.float32}[ + args.potential_precision or args.precision + ] + model = model.to(device=device, dtype=potential_dtype) + print( + f"precision: MD state {args.precision}, " + f"potential {args.potential_precision or args.precision}" + ) + force_field = CallableForceField( + _matpes_energy_forces( + model, + Z=Z.to(device), + neighbors=neighbors, + potential_dtype=potential_dtype, + compile_energy=args.compile, + autocast_dtype=torch.bfloat16 if args.autocast_bf16 else None, + ), + neighbors=neighbors, + energy_scale=1.0 / EV_PER_AMU_A2_FS2, + ).to(device) + + t0 = time.perf_counter() + out0 = force_field(pos) + print( + f"single point: E = {float(out0.energy) * EV_PER_AMU_A2_FS2:.6f} eV, " + f"|F|max = {float(out0.forces.abs().max()) * EV_PER_AMU_A2_FS2:.6f} eV/A " + f"({time.perf_counter() - t0:.2f} s)" + ) + + start_step = 0 + if args.resume is not None: + ckpt = torch.load(args.resume, map_location="cpu", weights_only=True) + start_step = int(ckpt["step"]) + pos = ckpt["pos"].to(dtype=config.ftype, device=device) + resume_vel = ckpt["vel"].to(dtype=config.ftype, device=device) + args.out = args.out.with_name(f"{args.out.stem}.from{start_step}{args.out.suffix}") + print(f"resuming at step {start_step}; trajectory continues in {args.out}") + + run_hooks: list = [ + TrajectoryHook( + args.out, + stride=args.stride, + numbers=Z, + write_xyz=False, + with_forces=True, + flush_every=args.flush_every, + ) + ] + if args.checkpoint_every: + run_hooks.append( + MDCheckpointHook( + args.out.with_suffix(".ckpt.pt"), + every=args.checkpoint_every, + step_offset=start_step, + ) + ) + md = MD( + force_field, + mass=mass, + dt=args.dt, + gamma=0.0, # NVE + dtype=config.ftype, + # autocast lives inside the force field's compiled callable (above); + # wrapping again here would just add per-call context overhead. + # No cadence kwarg: CallableForceField declares rebuilds_neighbors from + # the list bound to it, and the integrator derives its switch from that. + hooks=run_hooks, + seed=args.seed, + device=device, + ) + vel = ( + resume_vel + if args.resume is not None + else MaxwellBoltzmann(mass).sample(args.temperature, seed=args.seed) + ) + remaining = args.steps - start_step + if remaining <= 0: + print(f"nothing to do: checkpoint already at step {start_step} >= {args.steps}") + return + + # Observation cadence only — the neighbour policy runs inside each force + # evaluation (Integrator.eval_force) and must not force chunk=1. + chunk = max(1, int(args.stride)) + if args.checkpoint_every: + chunk = math.gcd(chunk, int(args.checkpoint_every)) + print(f"advancing in chunks of {chunk} steps (policy asked at every force evaluation)") + + rebuilds_before, ndanger_before = neighbors.rebuild_count, neighbors.ndanger + t0 = time.perf_counter() + md.run(pos, vel, remaining, chunk=chunk) + elapsed = time.perf_counter() - t0 + rebuilds = neighbors.rebuild_count - rebuilds_before + ndanger = neighbors.ndanger - ndanger_before + print( + f"NVE: {remaining} steps in {elapsed:.1f} s ({elapsed / remaining:.4f} s/step); " + f"neighbour rebuilds: {rebuilds}" + ) + + traj = torch.load(args.out, map_location="cpu", weights_only=True) + etot = traj["etot"] * EV_PER_AMU_A2_FS2 + n_frames = int(etot.numel()) + if n_frames < 2: + print( + f"E_tot drift: only {n_frames} trajectory frame(s) " + f"(stride={args.stride}, steps={args.steps}) — cannot estimate; " + f"use checkpoint log instead" + ) + else: + # Frames are kept at step stride, 2*stride, ...; span ≈ (n_frames)*stride*dt + t_ps = n_frames * args.stride * args.dt * 1e-3 + dE_meV_atom = float(etot[-1] - etot[0]) / Z.numel() * 1e3 + rate = dE_meV_atom / t_ps if t_ps > 0 else float("nan") + print( + f"E_tot drift = {dE_meV_atom:.4f} meV/atom over ~{t_ps:.2f} ps " + f"({n_frames} frames, {rate:.6f} meV/atom/ps); " + f"T range {float(traj['temp'].min()):.1f}-{float(traj['temp'].max()):.1f} K" + ) + # Measured, not predicted: the cadence is the list's, driven by how far the + # atoms actually moved, so there is no step-count formula to check against. + # ndanger is the alarm — nonzero at skin > 0 means a rebuild came too late + # and pairs were missed, so rerun with a larger skin or a tighter gate. + print( + f"neighbour rebuilds this segment: {rebuilds} of {remaining} steps " + f"({rebuilds / max(1, remaining):.3f} per step); ndanger={ndanger} " + f"(skin={neighbors.skin:g} A, every={neighbors.every}, delay={neighbors.delay}, " + f"check={neighbors.check})" + ) + # The comparison step needs the exact graph the trajectory was produced on. + # edge_index is (E, 2) [source, target] — the repo edge convention. + torch.save( + { + "Z": Z.cpu(), + "cell": cell.cpu(), + "edge_index": neighbors.edge_index.cpu(), + "shifts": neighbors.shifts.cpu(), + "num_edges": neighbors.num_edges, + "rebuild_count": neighbors.rebuild_count, + "energy_scale": EV_PER_AMU_A2_FS2, + }, + args.out.with_suffix(".graph.pt"), + ) + print(f"wrote {args.out} and {args.out.with_suffix('.graph.pt')}") + + +if __name__ == "__main__": + main() diff --git a/scripts/mm_param_learning/__init__.py b/scripts/mm_param_learning/__init__.py new file mode 100644 index 0000000..10265fe --- /dev/null +++ b/scripts/mm_param_learning/__init__.py @@ -0,0 +1,29 @@ +"""MM parameter-learning validation workspace scaffolding (milestone 1). + +Soft-depends on :mod:`molexp` for Workspace / Project / Experiment / Run and +workflow stubs. No hard dependency is added to molnex ``pyproject.toml``. +""" + +from __future__ import annotations + +__all__ = [ + "DEFAULT_WORKSPACE_ROOT", + "EXPERIMENT_SLUGS", + "MOLHUB_COORDINATES", + "PROGRAM", + "PROJECT_SLUG", + "VALIDATION_STAGES", + "MmParamLearningWorkspace", + "MmParamWorkflows", +] + +from scripts.mm_param_learning.constants import ( + DEFAULT_WORKSPACE_ROOT, + EXPERIMENT_SLUGS, + MOLHUB_COORDINATES, + PROGRAM, + PROJECT_SLUG, + VALIDATION_STAGES, +) +from scripts.mm_param_learning.materialize_workspace import MmParamLearningWorkspace +from scripts.mm_param_learning.workflows import MmParamWorkflows diff --git a/scripts/mm_param_learning/constants.py b/scripts/mm_param_learning/constants.py new file mode 100644 index 0000000..1eb3db4 --- /dev/null +++ b/scripts/mm_param_learning/constants.py @@ -0,0 +1,43 @@ +"""Constants for the MM parameter-learning validation program (milestone 1). + +Coordinates are provisional MolHub handles registered by mm-param-val-02/03. +Workflow stubs shape-check coordinates only — they never call ``molhub.fetch``. +""" + +from __future__ import annotations + +from pathlib import Path + +PROGRAM: str = "mm-param-learning-baseline" +"""Program id written into every seed run's params.""" + +PROJECT_SLUG: str = "mm-param-learning" +"""Workspace + project slug.""" + +DEFAULT_WORKSPACE_ROOT: Path = Path( + "/nobackup/proj/disk/teoroo/personal/jicli594/work/molcrafts/workspaces/mm-param-learning" +) +"""Operator default root. CI / tests must pass an explicit temp ``root=``.""" + +# Experiment slug → primary MolHub dataset coordinate (or empty when N/A). +MOLHUB_COORDINATES: dict[str, str] = { + "potential-parity": "dataset:espaloma/phalkethoh-mm-small@1", + "zinc-typing-recovery": "dataset:espaloma/zinc-typing@1", + "phalkethoh-mm-energy": "dataset:espaloma/phalkethoh-mm-small@1", + "latent-analysis": "dataset:espaloma/zinc-typing@1", +} + +EXPERIMENT_SLUGS: tuple[str, ...] = tuple(MOLHUB_COORDINATES.keys()) + +# Validation narrative stages 1–6 (Knowledge Note inventory). +VALIDATION_STAGES: tuple[tuple[str, str], ...] = ( + ("1", "Dataset inventory & MolHub coordinates"), + ("2", "B0 potential IR / Class-I kernel parity"), + ("3", "A GAFF typing recovery (zinc-typing)"), + ("4", "B1/B2 molecule-centered PhAlkEthOH MM energy"), + ("5", "D latent embedding purity analysis"), + ("6", "Synthesis & milestone gate"), +) + +KNOWLEDGE_NOTE_NAME: str = "mm-param-validation-stages" +"""Idempotent Knowledge Note slug under the workspace bundle root.""" diff --git a/scripts/mm_param_learning/diagnose_phalkethoh_mm.py b/scripts/mm_param_learning/diagnose_phalkethoh_mm.py new file mode 100644 index 0000000..9d9cb29 --- /dev/null +++ b/scripts/mm_param_learning/diagnose_phalkethoh_mm.py @@ -0,0 +1,77 @@ +#!/usr/bin/env python3 +"""Diagnose real PhAlkEthOH MM-small payload (no model train). + +Loads :class:`molhub.dataset.PhalkethohMMDataset` from a **normalized** tree +and reports molecule-level split sizes, conformer counts, and the +molecule-mean-centered energy scale (kcal/mol) on train/val/test. + +This is the first real-data gate for Validation B2: if this fails, training +cannot start. +""" + +from __future__ import annotations + +import argparse +import json +from collections import defaultdict +from pathlib import Path + +import numpy as np + + +def main() -> int: + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("--data-root", type=Path, required=True) + p.add_argument("--record-root", type=Path, default=None) + args = p.parse_args() + + from molhub.dataset import PhalkethohMMDataset + from molhub.dataset.meta import Targets + + report: dict = {"data_root": str(args.data_root), "splits": {}} + for split in ("train", "val", "test"): + ds = PhalkethohMMDataset(args.data_root, download=False, split=split) + by: dict[str, list[float]] = defaultdict(list) + for i in range(len(ds)): + fr = ds[i] + t = Targets(fr) + by[str(t["molecule_id"])].append(float(t["mm_energy"])) + n_mol = len(by) + n_fr = len(ds) + # per-molecule centering then pool RMSE of zeros vs values is std + centered = [] + for mid, es in by.items(): + arr = np.asarray(es, dtype=np.float64) + centered.append(arr - arr.mean()) + cat = np.concatenate(centered) if centered else np.zeros(0) + report["splits"][split] = { + "n_frames": n_fr, + "n_molecules": n_mol, + "confs_per_mol_mean": n_fr / n_mol if n_mol else 0, + "energy_kcal_mean": float(np.mean([e for es in by.values() for e in es])) if n_fr else None, + "energy_kcal_std": float(np.std([e for es in by.values() for e in es])) if n_fr else None, + "centered_energy_std_kcal": float(np.std(cat)) if cat.size else None, + "source_id": ds.source_id, + } + print(json.dumps(report, indent=2)) + if args.record_root is not None: + args.record_root.mkdir(parents=True, exist_ok=True) + (args.record_root / "artifacts").mkdir(exist_ok=True) + (args.record_root / "artifacts" / "phalkethoh_mm_diagnose.json").write_text( + json.dumps(report, indent=2) + "\n" + ) + metrics = args.record_root / "metrics" + metrics.mkdir(exist_ok=True) + # one-line jsonl for molplot + row = {"step": 0} + for split, d in report["splits"].items(): + row[f"{split}_n_frames"] = d["n_frames"] + row[f"{split}_n_molecules"] = d["n_molecules"] + row[f"{split}_centered_energy_std_kcal"] = d["centered_energy_std_kcal"] + with (metrics / "metrics.jsonl").open("w") as f: + f.write(json.dumps(row) + "\n") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/mm_param_learning/espaloma_mm_protocol_train.py b/scripts/mm_param_learning/espaloma_mm_protocol_train.py new file mode 100644 index 0000000..2fe8f42 --- /dev/null +++ b/scripts/mm_param_learning/espaloma_mm_protocol_train.py @@ -0,0 +1,495 @@ +#!/usr/bin/env python3 +"""Espaloma-protocol MM fitting on a small PhAlkEthOH subset (our molexp project). + +Matches Espaloma's published MM-small toy protocol as closely as our stack allows: + +* molecule-level shuffle/split seed **2666**, ratios **8:1:1** (80/10/10) +* loss = graph-level **MSE(E_pred, E_ref)** (absolute MM energy, not QM-centering) +* optimizer **Adam(lr=1e-4)** +* train loader **batch_size=100** conformers +* report L1 (MAE) in kcal/mol (Espaloma notebook multiplies Hartree metrics by 627.5; + our labels are already kcal/mol after conversion) + +Architecture (Espaloma ↔ ours): + +| Espaloma | Ours | +|----------|------| +| SAGEConv ×3, 128, ReLU | ChemEncoder atom/bond/angle/proper dim=128 | +| Janossy → bond/angle log-coeff, torsion k×6 | Bond/Angle/ProperParamHead (n_terms=6) | +| GeometryInGraph + EnergyInGraph | ClassicalMMComposer Class-I kernels | +| ChargeEquilibrium + NB | **omitted** (documented deviation; bonded Class-I only) | + +Small system (default): the **12 smallest molecules** (by atom count) in the +normalized PhAlkEthOH tree, **all 100 confs** (or ``--max-confs``). + +Runs under the molexp project ``mm-param-learning`` experiment +``espaloma-mm-protocol-mini`` when ``--record-root`` points at a run dir. + +Example:: + + export PY=/path/to/python # torch+molpy env + export PYTHONPATH=src:../molhub/src + $PY scripts/mm_param_learning/espaloma_mm_protocol_train.py \\ + --data-root .../phalkethoh-mm-small-normalized \\ + --epochs 200 --device cpu \\ + --record-root .../experiments/espaloma-mm-protocol-mini/runs/run-001 +""" + +from __future__ import annotations + +import argparse +import json +import random +import time +from collections import defaultdict +from pathlib import Path + +import numpy as np +import torch +import torch.nn as nn +from tensordict import TensorDict + +from molhub.dataset import PhalkethohMMDataset +from molhub.dataset.meta import Targets +from molpot.composition import ClassicalMMComposer, ClassicalMMParameterizer +from molpot.composition.mm_heads import ( + AngleParamHead, + BondParamHead, + ProperTorsionParamHead, +) +from molrep.chem.encoder import ChemEncoder + +# Espaloma notebook protocol pins +ESPALOMA_SPLIT_SEED = 2666 +ESPALOMA_SPLIT_RATIOS = (8, 1, 1) # train:val:test parts +ESPALOMA_ADAM_LR = 1e-4 +ESPALOMA_BATCH_SIZE = 100 +ESPALOMA_HIDDEN = 128 +ESPALOMA_N_TORSION_TERMS = 6 + +# Energy already converted to kcal/mol in the normalized MolHub tree +# (Hartree * 627.5094740631). Espaloma notebook reports L1 * 627.5. + + +def enumerate_angles(atomi: np.ndarray, atomj: np.ndarray, n_atoms: int) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """i-j-k angles from undirected bonds (j central).""" + adj: list[list[int]] = [[] for _ in range(n_atoms)] + for i, j in zip(atomi.tolist(), atomj.tolist(), strict=True): + adj[i].append(j) + adj[j].append(i) + ii: list[int] = [] + jj: list[int] = [] + kk: list[int] = [] + for j in range(n_atoms): + nbrs = adj[j] + for a in range(len(nbrs)): + for b in range(a + 1, len(nbrs)): + ii.append(nbrs[a]) + jj.append(j) + kk.append(nbrs[b]) + if not ii: + z = np.zeros(0, dtype=np.int64) + return z, z, z + return ( + np.asarray(ii, dtype=np.int64), + np.asarray(jj, dtype=np.int64), + np.asarray(kk, dtype=np.int64), + ) + + +def enumerate_propers( + atomi: np.ndarray, atomj: np.ndarray, n_atoms: int +) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: + """i-j-k-l propers with central bond j-k.""" + adj: list[list[int]] = [[] for _ in range(n_atoms)] + edges = list(zip(atomi.tolist(), atomj.tolist(), strict=True)) + for i, j in edges: + adj[i].append(j) + adj[j].append(i) + ii: list[int] = [] + jj: list[int] = [] + kk: list[int] = [] + ll: list[int] = [] + seen: set[tuple[int, int]] = set() + for a, b in edges: + for j, k in ((a, b), (b, a)): + if (j, k) in seen: + continue + seen.add((j, k)) + seen.add((k, j)) + for i in adj[j]: + if i == k: + continue + for l in adj[k]: + if l == j or l == i: + continue + ii.append(i) + jj.append(j) + kk.append(k) + ll.append(l) + if not ii: + z = np.zeros(0, dtype=np.int64) + return z, z, z, z + return ( + np.asarray(ii, dtype=np.int64), + np.asarray(jj, dtype=np.int64), + np.asarray(kk, dtype=np.int64), + np.asarray(ll, dtype=np.int64), + ) + + +def frame_to_batch(fr, device: torch.device) -> tuple[TensorDict, torch.Tensor, str]: + """One conformer Frame → TensorDict batch + energy kcal/mol + molecule_id.""" + atoms = fr["atoms"] + bonds = fr["bonds"] + z = torch.as_tensor(atoms["number"], dtype=torch.long, device=device) + pos = torch.stack( + [torch.as_tensor(atoms[c], dtype=torch.float64, device=device) for c in ("x", "y", "z")], + dim=-1, + ) + atomi = np.asarray(bonds["atomi"], dtype=np.int64) + atomj = np.asarray(bonds["atomj"], dtype=np.int64) + n = int(z.shape[0]) + ai_t = torch.as_tensor(atomi, dtype=torch.long, device=device) + aj_t = torch.as_tensor(atomj, dtype=torch.long, device=device) + ang_i, ang_j, ang_k = enumerate_angles(atomi, atomj, n) + pr_i, pr_j, pr_k, pr_l = enumerate_propers(atomi, atomj, n) + + ns: dict = { + "atoms": TensorDict({"Z": z, "pos": pos}, batch_size=[n]), + "bonds": TensorDict({"atomi": ai_t, "atomj": aj_t}, batch_size=[ai_t.shape[0]]), + } + if ang_i.size: + ns["angles"] = TensorDict( + { + "atomi": torch.as_tensor(ang_i, dtype=torch.long, device=device), + "atomj": torch.as_tensor(ang_j, dtype=torch.long, device=device), + "atomk": torch.as_tensor(ang_k, dtype=torch.long, device=device), + }, + batch_size=[ang_i.shape[0]], + ) + if pr_i.size: + ns["propers"] = TensorDict( + { + "atomi": torch.as_tensor(pr_i, dtype=torch.long, device=device), + "atomj": torch.as_tensor(pr_j, dtype=torch.long, device=device), + "atomk": torch.as_tensor(pr_k, dtype=torch.long, device=device), + "atoml": torch.as_tensor(pr_l, dtype=torch.long, device=device), + }, + batch_size=[pr_i.shape[0]], + ) + batch = TensorDict(ns, batch_size=[]) + energy = torch.tensor(float(Targets(fr)["mm_energy"]), dtype=torch.float64, device=device) + mid = str(Targets(fr)["molecule_id"]) + return batch, energy, mid + + +def molecule_split(mol_ids: list[str], seed: int = ESPALOMA_SPLIT_SEED) -> dict[str, list[str]]: + """Espaloma notebook: shuffle(seed) then split([8,1,1]).""" + ids = list(mol_ids) + rng = random.Random(seed) + rng.shuffle(ids) + n = len(ids) + a, b, c = ESPALOMA_SPLIT_RATIOS + total = a + b + c + n_tr = max(1, int(round(n * a / total))) if n >= 3 else max(1, n - 2) + n_vl = max(1, int(round(n * b / total))) if n >= 3 else 1 + if n_tr + n_vl >= n: + n_tr = max(1, n - 2) + n_vl = 1 + n_te = n - n_tr - n_vl + if n_te < 1: + n_te = 1 + n_tr = max(1, n - n_vl - n_te) + return { + "train": ids[:n_tr], + "val": ids[n_tr : n_tr + n_vl], + "test": ids[n_tr + n_vl :], + "seed": seed, + "ratios": [a, b, c], + "scheme": "espaloma-original", + } + + +def pick_small_molecules( + data_root: Path, n_mols: int = 12, max_confs: int | None = None +) -> dict[str, list]: + """Load all splits, pick *n_mols* smallest molecules by atom count.""" + # Use full corpus (no split filter) so we re-apply Espaloma seed ourselves. + by: dict[str, list] = defaultdict(list) + n_atoms: dict[str, int] = {} + for split in ("train", "val", "test"): + ds = PhalkethohMMDataset(data_root, download=False, split=split) + for i in range(len(ds)): + fr = ds[i] + mid = str(Targets(fr)["molecule_id"]) + n_atoms[mid] = int(len(fr["atoms"]["number"])) + if max_confs is None or len(by[mid]) < max_confs: + by[mid].append(fr) + ranked = sorted(n_atoms.items(), key=lambda kv: (kv[1], kv[0])) + chosen = [mid for mid, _ in ranked[:n_mols]] + return {mid: by[mid] for mid in chosen} + + +def build_espaloma_like_model() -> ClassicalMMParameterizer: + """ChemEncoder + bond/angle/proper heads (128-wide, 6-term torsions).""" + encoder = ChemEncoder( + atom_dim=ESPALOMA_HIDDEN, + bond_dim=ESPALOMA_HIDDEN, + angle_dim=ESPALOMA_HIDDEN, + proper_dim=ESPALOMA_HIDDEN, + improper_dim=ESPALOMA_HIDDEN, + hidden_dim=ESPALOMA_HIDDEN, + ) + composer = ClassicalMMComposer( + bond_head=BondParamHead(feature_dim=ESPALOMA_HIDDEN, hidden_dim=ESPALOMA_HIDDEN), + angle_head=AngleParamHead(feature_dim=ESPALOMA_HIDDEN, hidden_dim=ESPALOMA_HIDDEN), + proper_head=ProperTorsionParamHead( + feature_dim=ESPALOMA_HIDDEN, + hidden_dim=ESPALOMA_HIDDEN, + n_terms=ESPALOMA_N_TORSION_TERMS, + periodicity=tuple(range(1, ESPALOMA_N_TORSION_TERMS + 1)), + ), + ) + return ClassicalMMParameterizer(encoder, composer).double() + + +def mse_energy(pred: torch.Tensor, ref: torch.Tensor) -> torch.Tensor: + """Espaloma GraphMetric MSE on graph energies.""" + return torch.mean((pred - ref) ** 2) + + +def mae_energy(pred: torch.Tensor, ref: torch.Tensor) -> float: + return float(torch.mean(torch.abs(pred - ref)).detach()) + + +@torch.no_grad() +def eval_split( + model: ClassicalMMParameterizer, + frames_by_mol: dict[str, list], + mol_ids: list[str], + device: torch.device, +) -> dict[str, float]: + """Topology once per molecule; evaluate all confs; report MAE/RMSE kcal/mol.""" + preds: list[torch.Tensor] = [] + refs: list[torch.Tensor] = [] + for mid in mol_ids: + frames = frames_by_mol[mid] + batch0, _, _ = frame_to_batch(frames[0], device) + ir = model.parameterize(batch0) + for fr in frames: + batch, e_ref, _ = frame_to_batch(fr, device) + e = model.energy(batch, ir=ir, pos=batch["atoms", "pos"]).reshape(()) + preds.append(e) + refs.append(e_ref.reshape(())) + pred = torch.stack(preds) + ref = torch.stack(refs) + return { + "mae_kcal": mae_energy(pred, ref), + "rmse_kcal": float(torch.sqrt(torch.mean((pred - ref) ** 2)).detach()), + "n": int(pred.numel()), + } + + +def train( + model: ClassicalMMParameterizer, + frames_by_mol: dict[str, list], + split: dict[str, list[str]], + *, + epochs: int, + batch_size: int, + lr: float, + device: torch.device, + record_root: Path | None, + eval_every: int = 5, +) -> list[dict]: + opt = torch.optim.Adam(model.parameters(), lr=lr) + train_ids = split["train"] + val_ids = split["val"] + + # Flatten train conformers for Espaloma-style shuffled minibatches of graphs + train_pool: list = [] + for mid in train_ids: + train_pool.extend(frames_by_mol[mid]) + + # Cache static topology batches (geometry free) for parameterize() + topo_batch: dict[str, TensorDict] = {} + for mid in frames_by_mol: + b0, _, _ = frame_to_batch(frames_by_mol[mid][0], device) + # topology-only view: keep pos but parameterize ignores geometry in encoder + topo_batch[mid] = b0 + + history: list[dict] = [] + best_val = float("inf") + best_state: dict | None = None + + for epoch in range(epochs): + model.train() + order = list(range(len(train_pool))) + random.shuffle(order) + epoch_loss = 0.0 + n_batches = 0 + t0 = time.time() + for start in range(0, len(order), batch_size): + idx = order[start : start + batch_size] + # Espaloma batches mixed graphs; mean MSE(u, u_ref) at graph level. + losses: list[torch.Tensor] = [] + by_mid: dict[str, list] = defaultdict(list) + for j in idx: + fr = train_pool[j] + mid = str(Targets(fr)["molecule_id"]) + by_mid[mid].append(fr) + for mid, frames in by_mid.items(): + ir = model.parameterize(topo_batch[mid]) + for fr in frames: + batch, e_ref, _ = frame_to_batch(fr, device) + e = model.energy(batch, ir=ir, pos=batch["atoms", "pos"]).reshape(()) + losses.append((e - e_ref.reshape(())) ** 2) + if not losses: + continue + loss = torch.stack(losses).mean() + opt.zero_grad() + loss.backward() + opt.step() + epoch_loss += float(loss.detach()) + n_batches += 1 + + do_eval = (epoch % eval_every == 0) or (epoch == epochs - 1) + if do_eval: + tr = eval_split(model, frames_by_mol, train_ids, device) + va = eval_split(model, frames_by_mol, val_ids, device) + row = { + "epoch": epoch, + "train_mse": epoch_loss / max(n_batches, 1), + "train_mae_kcal": tr["mae_kcal"], + "train_rmse_kcal": tr["rmse_kcal"], + "val_mae_kcal": va["mae_kcal"], + "val_rmse_kcal": va["rmse_kcal"], + "elapsed_s": time.time() - t0, + "lr": lr, + "batch_size": batch_size, + } + # Espaloma: early stop on validation metric + if va["mae_kcal"] < best_val: + best_val = va["mae_kcal"] + best_state = {k: v.detach().cpu().clone() for k, v in model.state_dict().items()} + else: + row = { + "epoch": epoch, + "train_mse": epoch_loss / max(n_batches, 1), + "elapsed_s": time.time() - t0, + "lr": lr, + "batch_size": batch_size, + } + history.append(row) + print(json.dumps(row), flush=True) + + if record_root is not None: + metrics_dir = record_root / "metrics" + metrics_dir.mkdir(parents=True, exist_ok=True) + with (metrics_dir / "metrics.jsonl").open("a") as fh: + fh.write(json.dumps(row) + "\n") + + if best_state is not None: + model.load_state_dict(best_state) + return history + + +def main(argv: list[str] | None = None) -> int: + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--data-root", type=Path, required=True, help="Normalized PhAlkEthOH tree") + p.add_argument("--n-mols", type=int, default=12, help="Small-system: number of smallest molecules") + p.add_argument("--max-confs", type=int, default=20, help="Cap confs/mol (Espaloma uses 100; default 20 for mini)") + p.add_argument("--epochs", type=int, default=200) + p.add_argument("--batch-size", type=int, default=ESPALOMA_BATCH_SIZE) + p.add_argument("--lr", type=float, default=ESPALOMA_ADAM_LR) + p.add_argument("--seed", type=int, default=ESPALOMA_SPLIT_SEED) + p.add_argument("--eval-every", type=int, default=5, help="Full train/val MAE every N epochs") + p.add_argument("--device", default="cpu") + p.add_argument("--record-root", type=Path, default=None) + p.add_argument("--checkpoint", type=Path, default=None) + args = p.parse_args(argv) + + random.seed(args.seed) + torch.manual_seed(args.seed) + device = torch.device(args.device) + + frames_by_mol = pick_small_molecules(args.data_root, n_mols=args.n_mols, max_confs=args.max_confs) + mol_ids = list(frames_by_mol.keys()) + split = molecule_split(mol_ids, seed=args.seed) + n_atoms = { + mid: int(len(frames_by_mol[mid][0]["atoms"]["number"])) for mid in mol_ids + } + + meta = { + "protocol": "espaloma-mm-small-notebook", + "references": ["arXiv:2010.01196", "DOI:10.1039/D2SC02739A"], + "split_seed": args.seed, + "split_ratios_parts": list(ESPALOMA_SPLIT_RATIOS), + "adam_lr": args.lr, + "batch_size": args.batch_size, + "loss": "MSE(E_pred, E_ref) graph-level absolute kcal/mol", + "metric_report": "MAE/RMSE kcal/mol (labels already kcal/mol)", + "model": { + "encoder": f"ChemEncoder dim={ESPALOMA_HIDDEN}", + "heads": "Bond+Angle+Proper(n_terms=6)", + "energy": "ClassicalMMComposer bonded Class-I", + "omitted_vs_espaloma": ["ChargeEquilibrium", "nonbonded LJ/Coulomb EnergyInGraph NB"], + }, + "subset": { + "n_mols": len(mol_ids), + "mol_ids": mol_ids, + "n_atoms": n_atoms, + "confs_per_mol": {mid: len(frames_by_mol[mid]) for mid in mol_ids}, + "selection": f"{args.n_mols} smallest molecules by atom count", + }, + "split": {k: v for k, v in split.items()}, + "project": "mm-param-learning", + "experiment": "espaloma-mm-protocol-mini", + } + print(json.dumps({"setup": meta}, indent=2), flush=True) + + if args.record_root is not None: + args.record_root.mkdir(parents=True, exist_ok=True) + (args.record_root / "artifacts").mkdir(exist_ok=True) + (args.record_root / "metrics").mkdir(exist_ok=True) + # truncate metrics + (args.record_root / "metrics" / "metrics.jsonl").write_text("") + (args.record_root / "artifacts" / "setup.json").write_text(json.dumps(meta, indent=2) + "\n") + + model = build_espaloma_like_model().to(device) + history = train( + model, + frames_by_mol, + split, + epochs=args.epochs, + batch_size=args.batch_size, + lr=args.lr, + device=device, + record_root=args.record_root, + eval_every=args.eval_every, + ) + + te = eval_split(model, frames_by_mol, split["test"], device) + val_rows = [h for h in history if "val_mae_kcal" in h] + final = { + "best_val_mae_kcal": min(h["val_mae_kcal"] for h in val_rows) if val_rows else None, + "best_val_epoch": min(val_rows, key=lambda h: h["val_mae_kcal"])["epoch"] if val_rows else None, + "test_mae_kcal": te["mae_kcal"], + "test_rmse_kcal": te["rmse_kcal"], + "test_n": te["n"], + "epochs": args.epochs, + } + print(json.dumps({"final": final}), flush=True) + + if args.record_root is not None: + (args.record_root / "artifacts" / "history.json").write_text(json.dumps(history, indent=2) + "\n") + (args.record_root / "artifacts" / "final.json").write_text(json.dumps(final, indent=2) + "\n") + ckpt = args.checkpoint or (args.record_root / "artifacts" / "model.pt") + torch.save({"model": model.state_dict(), "meta": meta, "final": final}, ckpt) + + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/mm_param_learning/materialize_workspace.py b/scripts/mm_param_learning/materialize_workspace.py new file mode 100644 index 0000000..1f46078 --- /dev/null +++ b/scripts/mm_param_learning/materialize_workspace.py @@ -0,0 +1,152 @@ +"""Idempotent materialization of the MM parameter-learning molexp workspace. + +Creates ``Workspace → Project → Experiment → Run`` under an explicit *root* +(temp in CI) or :data:`DEFAULT_WORKSPACE_ROOT` for operators. Never fetches +MolHub artifacts; seed run params only record coordinates and program ids. +""" + +from __future__ import annotations + +import argparse +from pathlib import Path +from typing import Any + +from scripts.mm_param_learning.constants import ( + DEFAULT_WORKSPACE_ROOT, + EXPERIMENT_SLUGS, + KNOWLEDGE_NOTE_NAME, + MOLHUB_COORDINATES, + PROGRAM, + PROJECT_SLUG, + VALIDATION_STAGES, +) +from scripts.mm_param_learning.workflows import MmParamWorkflows + + +def _knowledge_body() -> str: + lines = [ + "# MM Parameter-Learning Validation — Stages 1–6", + "", + f"Program: `{PROGRAM}`", + "", + "## Stages", + "", + ] + for num, title in VALIDATION_STAGES: + lines.append(f"### Stage {num} — {title}") + lines.append("") + lines.extend( + [ + "## MolHub coordinates", + "", + "Data is loaded via the **Python molhub SDK** only (no molhub MCP plane).", + "Workspace orchestration uses **molexp**.", + "", + ] + ) + for slug, coord in MOLHUB_COORDINATES.items(): + lines.append(f"- `{slug}` → `{coord}`") + lines.append("") + return "\n".join(lines) + + +class MmParamLearningWorkspace: + """Materialize the milestone-1 validation workspace tree. + + Idempotent: a second call keeps experiment count at four and does not + duplicate the validation Knowledge Note. + + Args: + root: Workspace root directory. Defaults to + :data:`DEFAULT_WORKSPACE_ROOT`. Tests **must** pass a temp path. + """ + + def __init__(self, root: str | Path | None = None) -> None: + self.root = Path(root) if root is not None else DEFAULT_WORKSPACE_ROOT + + def materialize(self) -> dict[str, Any]: + """Create workspace layout, four experiments, seed runs, and note. + + Returns: + Summary dict with ``root``, ``project``, ``experiments``, ``note``. + + Raises: + ImportError: If molexp is not installed. + """ + try: + from molexp import Workspace + from molexp.workspace.bundle import Bundle + except ImportError as exc: # pragma: no cover - soft dep + raise ImportError( + "molexp is required for MmParamLearningWorkspace.materialize" + ) from exc + + self.root.mkdir(parents=True, exist_ok=True) + ws = Workspace(self.root, name=PROJECT_SLUG) + ws.materialize() + project = ws.add_project(PROJECT_SLUG) + project.materialize() + + workflows = MmParamWorkflows() + experiments: list[str] = [] + for slug in EXPERIMENT_SLUGS: + coord = MOLHUB_COORDINATES[slug] + exp = project.add_experiment( + slug, + params={ + "dataset_coordinate": coord, + "milestone": 1, + "program": PROGRAM, + }, + description=f"MM param validation experiment {slug}", + tags=["mm-param-learning", "milestone-1"], + ) + exp.materialize() + # Seed run — fixed id for idempotency. + run = exp.add_run( + params={ + "dataset_coordinate": coord, + "milestone": 1, + "program": PROGRAM, + "experiment": slug, + }, + id=f"seed-{slug}", + ) + run.materialize() + # Compile workflow for side-effect validation (not executed). + workflows.build(slug) + experiments.append(slug) + + bundle = Bundle(ws.root) + note = bundle.create_note(KNOWLEDGE_NOTE_NAME, body=_knowledge_body()) + + return { + "root": str(self.root), + "project": PROJECT_SLUG, + "experiments": experiments, + "note": KNOWLEDGE_NOTE_NAME, + "workspace": ws, + "note_obj": note, + } + + +def main(argv: list[str] | None = None) -> int: + """CLI entry: materialize the operator workspace (or ``--root`` override).""" + parser = argparse.ArgumentParser( + description="Materialize mm-param-learning molexp workspace (idempotent)." + ) + parser.add_argument( + "--root", + type=Path, + default=None, + help="Workspace root (default: operator DEFAULT_WORKSPACE_ROOT).", + ) + args = parser.parse_args(argv) + summary = MmParamLearningWorkspace(root=args.root).materialize() + print(f"materialized project={summary['project']} root={summary['root']}") + print(f"experiments={summary['experiments']}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/mm_param_learning/train_phalkethoh_mm_energy.py b/scripts/mm_param_learning/train_phalkethoh_mm_energy.py new file mode 100644 index 0000000..30359bd --- /dev/null +++ b/scripts/mm_param_learning/train_phalkethoh_mm_energy.py @@ -0,0 +1,172 @@ +#!/usr/bin/env python3 +"""Train bond-head ClassicalMMParameterizer on PhAlkEthOH MM-small (Validation B2 smoke). + +Requires a **normalized** data tree (from molhub convert script):: + + PYTHONPATH=src:../molhub/src python scripts/mm_param_learning/train_phalkethoh_mm_energy.py \\ + --data-root /path/to/phalkethoh-mm-small-normalized \\ + --epochs 5 --max-train-mols 8 --max-confs-per-mol 8 + +Primary metric: molecule-mean-centered energy RMSE/MAE (kcal/mol). +Bond-only model is a smoke path (full valence heads come later). +""" + +from __future__ import annotations + +import argparse +import json +from collections import defaultdict +from pathlib import Path + +import torch +from tensordict import TensorDict + +from molhub.dataset import PhalkethohMMDataset +from molhub.dataset.meta import Targets +from molix.core.losses.molecular import center_by_group +from molpot.composition import ClassicalMMComposer, ClassicalMMParameterizer +from molpot.composition.mm_heads import BondParamHead +from molrep.chem.encoder import ChemEncoder + + +def _group_frames(ds: PhalkethohMMDataset, max_mols: int, max_confs: int) -> dict[str, list]: + by: dict[str, list] = defaultdict(list) + for i in range(len(ds)): + fr = ds[i] + mid = str(Targets(fr)["molecule_id"]) + if mid not in by and len(by) >= max_mols: + continue + if len(by[mid]) < max_confs: + by[mid].append(fr) + if len(by) >= max_mols and all(len(v) >= max_confs for v in by.values()): + break + return dict(by) + + +def _to_batch(fr) -> TensorDict: + atoms = fr["atoms"] + bonds = fr["bonds"] + z = torch.as_tensor(atoms["number"], dtype=torch.long) + pos = torch.stack( + [torch.as_tensor(atoms[c], dtype=torch.float64) for c in ("x", "y", "z")], + dim=-1, + ) + atomi = torch.as_tensor(bonds["atomi"], dtype=torch.long) + atomj = torch.as_tensor(bonds["atomj"], dtype=torch.long) + return TensorDict( + { + "atoms": TensorDict({"Z": z, "pos": pos}, batch_size=[z.shape[0]]), + "bonds": TensorDict({"atomi": atomi, "atomj": atomj}, batch_size=[atomi.shape[0]]), + }, + batch_size=[], + ) + + +def _run_epoch( + model: ClassicalMMParameterizer, + by: dict[str, list], + *, + train: bool, + opt: torch.optim.Optimizer | None, +) -> dict[str, float]: + model.train(train) + preds: list[torch.Tensor] = [] + refs: list[torch.Tensor] = [] + groups: list[int] = [] + mid_map: dict[str, int] = {} + for mid, frames in by.items(): + mid_map.setdefault(mid, len(mid_map)) + batch0 = _to_batch(frames[0]) + with torch.set_grad_enabled(train): + ir = model.parameterize(batch0) + for fr in frames: + batch = _to_batch(fr) + e = model.energy(batch, ir=ir, pos=batch["atoms", "pos"]).reshape(()) + preds.append(e) + refs.append( + torch.tensor(float(Targets(fr)["mm_energy"]), dtype=torch.float64) + ) + groups.append(mid_map[mid]) + pred = torch.stack(preds) + ref = torch.stack(refs) + g = torch.tensor(groups, dtype=torch.long) + loss = torch.mean((center_by_group(pred, g) - center_by_group(ref, g)) ** 2) + if train and opt is not None: + opt.zero_grad() + loss.backward() + opt.step() + with torch.no_grad(): + cp = center_by_group(pred, g) + cr = center_by_group(ref, g) + rmse = float(torch.sqrt(torch.mean((cp - cr) ** 2)).detach()) + mae = float(torch.mean(torch.abs(cp - cr)).detach()) + return {"rmse_kcal": rmse, "mae_kcal": mae, "loss": float(loss.detach())} + + +def main(argv: list[str] | None = None) -> int: + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("--data-root", type=Path, required=True) + p.add_argument("--epochs", type=int, default=5) + p.add_argument("--lr", type=float, default=1e-3) + p.add_argument("--max-train-mols", type=int, default=8) + p.add_argument("--max-val-mols", type=int, default=4) + p.add_argument("--max-confs-per-mol", type=int, default=8) + p.add_argument("--record-root", type=Path, default=None) + p.add_argument("--seed", type=int, default=0) + args = p.parse_args(argv) + + torch.manual_seed(args.seed) + train_by = _group_frames( + PhalkethohMMDataset(args.data_root, download=False, split="train"), + args.max_train_mols, + args.max_confs_per_mol, + ) + val_by = _group_frames( + PhalkethohMMDataset(args.data_root, download=False, split="val"), + args.max_val_mols, + args.max_confs_per_mol, + ) + print( + json.dumps( + { + "n_train_mols": len(train_by), + "n_val_mols": len(val_by), + "data_root": str(args.data_root), + } + ) + ) + + encoder = ChemEncoder(atom_dim=32, bond_dim=32) + composer = ClassicalMMComposer(bond_head=BondParamHead(feature_dim=32)) + model = ClassicalMMParameterizer(encoder, composer).double() + opt = torch.optim.Adam(model.parameters(), lr=args.lr) + + history: list[dict] = [] + for epoch in range(args.epochs): + tr = _run_epoch(model, train_by, train=True, opt=opt) + va = _run_epoch(model, val_by, train=False, opt=None) + row = { + "epoch": epoch, + "train_rmse_kcal": tr["rmse_kcal"], + "train_mae_kcal": tr["mae_kcal"], + "val_rmse_kcal": va["rmse_kcal"], + "val_mae_kcal": va["mae_kcal"], + } + history.append(row) + print(json.dumps(row)) + + if args.record_root is not None: + args.record_root.mkdir(parents=True, exist_ok=True) + (args.record_root / "metrics").mkdir(exist_ok=True) + (args.record_root / "artifacts").mkdir(exist_ok=True) + with (args.record_root / "metrics" / "train_smoke.jsonl").open("w") as fh: + for row in history: + fh.write(json.dumps(row) + "\n") + (args.record_root / "artifacts" / "train_smoke_history.json").write_text( + json.dumps(history, indent=2) + "\n" + ) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/mm_param_learning/workflows.py b/scripts/mm_param_learning/workflows.py new file mode 100644 index 0000000..88d434f --- /dev/null +++ b/scripts/mm_param_learning/workflows.py @@ -0,0 +1,89 @@ +"""Stub WorkflowCompiler skeletons for MM parameter-learning experiments. + +Each experiment compiles a three-task workflow: + +1. ``resolve_molhub_coordinates`` — shape-check ``dataset:`` coordinates (no network) +2. ``record_run_params`` — echo seed params for audit +3. ``write_placeholder_metrics`` — write ``scaffold_ok=1.0`` via RegisterMetric + +These stubs intentionally do **not** fetch datasets or run science kernels. +""" + +from __future__ import annotations + +from typing import Any + +from scripts.mm_param_learning.constants import EXPERIMENT_SLUGS, MOLHUB_COORDINATES + + +class MmParamWorkflows: + """Factory for per-experiment milestone-1 workflow stubs. + + Args: + None — all configuration comes from :mod:`constants`. + """ + + TASK_NAMES: tuple[str, ...] = ( + "resolve_molhub_coordinates", + "record_run_params", + "write_placeholder_metrics", + ) + + def build(self, experiment_slug: str) -> Any: + """Compile a three-task workflow for *experiment_slug*. + + Args: + experiment_slug: One of :data:`EXPERIMENT_SLUGS`. + + Returns: + A molexp ``CompiledWorkflow`` with :attr:`TASK_NAMES` in order. + + Raises: + ValueError: If *experiment_slug* is unknown. + ImportError: If molexp is not installed. + """ + if experiment_slug not in MOLHUB_COORDINATES: + raise ValueError( + f"Unknown experiment slug {experiment_slug!r}. Known: {list(EXPERIMENT_SLUGS)}" + ) + try: + from molexp.workflow import RegisterMetric, WorkflowCompiler + except ImportError as exc: # pragma: no cover - soft dep + raise ImportError( + "molexp is required to build MmParamWorkflows; install molexp " + "or skip these tests when molexp is absent." + ) from exc + + coordinate = MOLHUB_COORDINATES[experiment_slug] + wf = WorkflowCompiler(name=f"mm-param-{experiment_slug}") + + @wf.task + async def resolve_molhub_coordinates() -> dict[str, str]: + """Shape-check the experiment's MolHub coordinate (no network).""" + if not coordinate.startswith("dataset:"): + raise ValueError(f"coordinate must start with 'dataset:': {coordinate!r}") + return {"dataset_coordinate": coordinate} + + @wf.task(depends_on=["resolve_molhub_coordinates"]) + async def record_run_params( + resolve_molhub_coordinates: dict[str, str], + ) -> dict[str, str]: + """Echo resolved coordinate for run-params audit.""" + return dict(resolve_molhub_coordinates) + + @wf.task(depends_on=["record_run_params"]) + async def write_placeholder_metrics( + record_run_params: dict[str, str], + ) -> dict[str, Any]: + """Write placeholder scaffold metric (no science).""" + return { + "scaffold_ok": RegisterMetric(key="scaffold_ok", value=1.0), + "dataset_coordinate": record_run_params["dataset_coordinate"], + } + + return wf.compile() + + def task_names(self, experiment_slug: str) -> tuple[str, ...]: + """Return the three stub task names after a successful :meth:`build`.""" + compiled = self.build(experiment_slug) + return tuple(compiled.graph.task_names) diff --git a/scripts/mm_param_learning/zinc_typing_recovery.py b/scripts/mm_param_learning/zinc_typing_recovery.py new file mode 100644 index 0000000..128a382 --- /dev/null +++ b/scripts/mm_param_learning/zinc_typing_recovery.py @@ -0,0 +1,59 @@ +"""Experiment script scaffold for Validation A (zinc-typing recovery). + +Loads MolHub ``dataset:espaloma/zinc-typing@1`` when available; otherwise +operates on a caller-provided batch of continuous features + labels. +Does not wire into ClassicalMMParameterizer. +""" + +from __future__ import annotations + +from typing import Any + +import torch + +from molrep.chem import AtomTypeReadout, ChemEncoder, TypingRecoveryMetrics + + +def run_typing_recovery( + *, + encoder: ChemEncoder | None = None, + num_types: int = 2, + batches: list[tuple[Any, torch.Tensor]] | None = None, + steps: int = 50, + lr: float = 0.05, +) -> dict[str, Any]: + """Train a temporary TypeHead and return a metrics report dict. + + Args: + encoder: Optional ChemEncoder; default small encoder. + num_types: Discrete vocabulary size. + batches: List of ``(batch_tensordict, type_id_tensor)``. + steps: Optimization steps (eval-only throwaway head). + lr: Adam learning rate. + + Returns: + Dict with overall_accuracy and n_atoms from TypingRecoveryMetrics. + """ + if batches is None: + raise ValueError("batches required (MolHub path injects offline frames → batch)") + enc = encoder or ChemEncoder(atom_dim=16, bond_dim=8) + probe = AtomTypeReadout(enc, num_types=num_types) + opt = torch.optim.Adam(probe.parameters(), lr=lr) + for _ in range(steps): + for batch, y in batches: + opt.zero_grad() + logits = probe(batch)["logits"] + loss = torch.nn.functional.cross_entropy(logits, y.long()) + loss.backward() + opt.step() + metrics = TypingRecoveryMetrics(num_types=num_types) + for batch, y in batches: + pred = probe(batch)["pred_type_id"] + Z = batch["atoms", "Z"] + metrics.update(pred, y, Z=Z) + report = metrics.compute() + return { + "overall_accuracy": report.overall_accuracy, + "n_atoms": report.n_atoms, + "per_element_accuracy": report.per_element_accuracy, + } diff --git a/scripts/omol_port/README.md b/scripts/omol_port/README.md index 736aafe..9418fe5 100644 --- a/scripts/omol_port/README.md +++ b/scripts/omol_port/README.md @@ -1,38 +1,32 @@ -# MACE-OMOL port — verification scripts +# MACE-OMOL port notes -Goal: extend MolNex (molrep/molpot/molzoo, cuEquivariance) so it can load the -official **MACE-OMOL** foundation model (`MACE-omol-0-extra-large-1024.model`, -`ScaleShiftMACE`, 1024 ch, r_max=6.0, 3 interactions, correlation=3, 83 elements) -and reproduce its energy/forces. +Goal: load the official **MACE-OMOL** foundation model into MolNex +(`molrep` / `molpot` / `molzoo`, cuEquivariance) and reproduce energy/forces. -These scripts check our ported blocks **bit-for-bit against the official -`mace-torch`** on CPU (float64). They load the plain-torch molrep/molpot modules -in isolation (stubbing `molix.config`) so no cuequivariance is required. +## Dependency policy -## Reference env -- CPU venv with official mace: `work/.mace-ref` (mace-torch 0.3.16, e3nn 0.4.4). -- OMOL checkpoint + dumped inventory: `work/mace_models/` - (`omol_inventory.json`, full block reference `OMOL_REFERENCE.md`). +MolNex package code and in-repo scripts **must not** import ASE, e3nn, or +`mace-torch`. Allowed runtime stack: MolCrafts packages (`molpy`, `mollog`, +`molcfg`, …), PyTorch / TensorDict, numpy, and cuEquivariance. -## Run -```bash -cd /nobackup/proj/disk/teoroo/personal/jicli594/work -TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD=1 .mace-ref/bin/python \ - molcrafts/molnex/scripts/omol_port/verify_radial.py -TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD=1 .mace-ref/bin/python \ - molcrafts/molnex/scripts/omol_port/verify_e0_scaleshift.py -``` +Upstream MACE / e3nn comparison is offline, out of this tree: run any +bit-for-bit oracle against a separate checkout and paste numbers into the +spec (`src/molzoo/specs/mace_omol.md` §7.4). Do not re-introduce those +imports here. -## Status (CPU-verified, max|diff| ~ machine eps) -- `verify_radial.py` — `molrep.embedding.BesselRBF` (normalize=False, eps=0, - trainable=True) vs `mace BesselBasis`; `molrep.embedding.PolynomialCutoff` vs - `mace PolynomialCutoff`. PASS. -- `verify_e0_scaleshift.py` — `molpot.heads.AtomicReferenceEnergy` vs - `mace AtomicEnergiesBlock`; `molpot.heads.GlobalRescale` vs - `mace ScaleShiftBlock` (single head). PASS. +## In-tree artifacts -## Remaining (equivariant — verify on aarch64 GPU, needs cuequivariance) -charge/spin joint embedding (project layer), RealAgnosticResidual**NonLinear** -InteractionBlock (linear_up/conv_tp/skip_tp/gate/linear_1/2/res/density_fn), -SymmetricContraction product, NonLinearBiasReadout; then the e3nn(mul_ir) → -cuEq(ir_mul) weight converter and full end-to-end energy/force comparison. +- `SPEC.md` — port design notes +- `convert_omol_to_cueq_state.py` + `_stub_unpickler.py` — offline, one-run + regeneration of `omol_cueq_state.pt` (the plain `weights_only=True` + state_dict twin of `OMOL-cueq.model`) without importing mace/e3nn. The + 419 MB dump itself is deliberately **not** kept: the source checkpoint is + downloadable, only the code is preserved. While the dump is absent, the + `MOLNEX_MACE_WEIGHTS_DIR`-gated `TestOfficialOMolWeights` cases skip. +- `src/molzoo/mace/variants.py` — `MACEOMol` (thin alias over + `molzoo.mace.potential.MACEPotential`) + the `load_omol_state_dict` + back-compat loader +- `src/molzoo/mace/checkpoint.py` — `OMOL_REMAP`, the official-weight key + remap consumed by `MACEPotential.from_checkpoint` +- `src/molzoo/specs/mace_omol.md` — paper↔code contract and run log + (mirrored byte-for-byte at `docs/molzoo/specs/mace_omol.md`) diff --git a/scripts/omol_port/_stub_unpickler.py b/scripts/omol_port/_stub_unpickler.py new file mode 100644 index 0000000..8207efa --- /dev/null +++ b/scripts/omol_port/_stub_unpickler.py @@ -0,0 +1,57 @@ +"""Offline-only: unpickle a mace/e3nn archive with nn.Module-shaped stubs. + +Stubbing the mace/e3nn classes as ``torch.nn.Module`` subclasses lets the real +``nn.Module.state_dict()`` machinery run over the restored tree, so every +*installed* module (cuequivariance's) contributes through its own +``_save_to_state_dict``. Nothing from ``mace`` or ``e3nn`` is imported — +their classes are replaced by inert stubs at unpickle time, which is what +lets this live in-repo under the "no mace-torch / e3nn" third-party rule. +""" + +import pickle +from pickle import * # noqa: F401,F403 + +import torch + +_ALLOWED = ( + "torch", + "collections", + "builtins", + "__builtin__", + "_codecs", + "numpy", + "cuequivariance", + "cuequivariance_torch", +) + + +class _ModStub(torch.nn.Module): + def __new__(cls, *args, **kwargs): + return object.__new__(cls) + + def __init__(self, *args, **kwargs): + torch.nn.Module.__init__(self) + self._stub_args = args + + def __setitem__(self, *args): + pass + + def append(self, *args): + pass + + +_MADE: dict[tuple[str, str], type] = {} + + +class Unpickler(pickle.Unpickler): + def find_class(self, module, name): + if module.split(".")[0] in _ALLOWED: + return super().find_class(module, name) + key = (module, name) + if key not in _MADE: + _MADE[key] = type(name, (_ModStub,), {"__module__": module}) + return _MADE[key] + + +def load(f, **kw): + return Unpickler(f, **kw).load() diff --git a/scripts/omol_port/convert_omol_to_cueq_state.py b/scripts/omol_port/convert_omol_to_cueq_state.py new file mode 100644 index 0000000..f54af58 --- /dev/null +++ b/scripts/omol_port/convert_omol_to_cueq_state.py @@ -0,0 +1,56 @@ +"""Offline conversion: OMOL-cueq.model (pickled ScaleShiftMACE) -> a plain state_dict. + +``OMOL-cueq.model`` is a full pickled ``mace.modules.models.ScaleShiftMACE`` +object, so ``torch.load`` needs ``mace`` + ``e3nn`` importable — packages +MolNex forbids under ``src/``, ``tests/`` and in-repo ``scripts/`` +(CLAUDE.md, "Allowed third-party surface"), and which are not installed in +the MolNex toolchain env. This script produces the offline twin of +``matpes_r2scan_cueq_state.pt``: a plain ``name -> tensor`` dump that +``torch.load(..., weights_only=True)`` reads with no third-party classes at +all. It needs neither ``mace`` nor ``e3nn``: :mod:`_stub_unpickler` +substitutes inert ``torch.nn.Module`` stubs for their classes while +unpickling, and cuEquivariance (installed) unpickles for real. + +The derived ``omol_cueq_state.pt`` is a throwaway artifact — the source +checkpoint is downloadable (ACEsuit mace-foundations OMol family) and this +script regenerates the dump in one run, so only the code is kept, not the +419 MB output. The ``MOLNEX_MACE_WEIGHTS_DIR``-gated +``TestOfficialOMolWeights`` cases in +``tests/test_molzoo/test_mace/test_checkpoint.py`` skip cleanly while the +dump is absent. + +Run once, next to the checkpoint (or pass explicit paths):: + + python scripts/omol_port/convert_omol_to_cueq_state.py \\ + [SOURCE.model] [TARGET_state.pt] + + # defaults: $MOLNEX_MACE_WEIGHTS_DIR/OMOL-cueq.model -> + # $MOLNEX_MACE_WEIGHTS_DIR/omol_cueq_state.pt + +Verification performed at first conversion (2026-08-09): + +* key dialect matches ``matpes_r2scan_cueq_state.pt`` (cueq 0.10.0: flat + ``.f.m.graphs.*.graph.c*`` constants, ``(1, numel)`` linear weights); +* every one of the 104 ``nn.Parameter`` tensors of the molnex ``MACEOMol`` + finds a home through ``molzoo.mace.checkpoint.OMOL_REMAP`` (which raises + on an unfilled learnable), with zero unexpected keys; +* ``radial_embedding.bessel_fn.bessel_weights`` sits 2.2120e-07 from the + analytic ``n*pi/r_max`` init — the fingerprint recorded independently in + ``src/molzoo/mace/checkpoint.py`` and ``src/molzoo/specs/mace_omol.md``. +""" + +import os +import sys +from pathlib import Path + +import _stub_unpickler +import torch + +weights_dir = Path(os.environ.get("MOLNEX_MACE_WEIGHTS_DIR", ".")) +source = Path(sys.argv[1]) if len(sys.argv) > 1 else weights_dir / "OMOL-cueq.model" +target = Path(sys.argv[2]) if len(sys.argv) > 2 else weights_dir / "omol_cueq_state.pt" + +model = torch.load(source, map_location="cpu", weights_only=False, pickle_module=_stub_unpickler) +state = {name: tensor.detach().clone() for name, tensor in model.state_dict().items()} +torch.save(state, target) +print(f"wrote {target}: {len(state)} tensors") diff --git a/scripts/omol_port/verify_e0_scaleshift.py b/scripts/omol_port/verify_e0_scaleshift.py deleted file mode 100644 index 63ae94a..0000000 --- a/scripts/omol_port/verify_e0_scaleshift.py +++ /dev/null @@ -1,64 +0,0 @@ -"""CPU unit test: molpot AtomicReferenceEnergy + GlobalRescale vs official mace.""" -import importlib.util -import sys -import types - -import torch -import torch.nn.functional as F - -torch.set_default_dtype(torch.float64) -torch.manual_seed(0) - -SRC = "/nobackup/proj/disk/teoroo/personal/jicli594/work/molcrafts/molnex/src" - -molix = types.ModuleType("molix") -config_mod = types.SimpleNamespace(ftype=torch.float64) -molix.config = config_mod -sys.modules["molix"] = molix -sys.modules["molix.config"] = config_mod - - -def load(name, path): - spec = importlib.util.spec_from_file_location(name, path) - mod = importlib.util.module_from_spec(spec) - spec.loader.exec_module(mod) - return mod - - -energy = load("molpot_energy", f"{SRC}/molpot/heads/energy.py") -rescale = load("molpot_rescale", f"{SRC}/molpot/heads/rescale.py") - -from mace.modules.blocks import AtomicEnergiesBlock, ScaleShiftBlock # noqa: E402 - -# ---- E0 / AtomicReferenceEnergy ---- -z_table = [1, 6, 8, 7] # H, C, O, N (element-table order) -e0 = torch.tensor([-13.6, -1029.0, -2042.0, -1485.0], dtype=torch.float64) -Z = torch.tensor([1, 8, 6, 6, 7, 1, 8]) # some atoms - -ours_e0 = energy.AtomicReferenceEnergy(atomic_energies=e0, atomic_numbers=z_table) -oe = ours_e0(Z) - -ref_block = AtomicEnergiesBlock(e0) # indexed by element-table position -# build one-hot of Z over the element table order -z_index = {z: i for i, z in enumerate(z_table)} -idx = torch.tensor([z_index[int(z)] for z in Z]) -one_hot = F.one_hot(idx, num_classes=len(z_table)).to(torch.float64) -re = ref_block(one_hot).squeeze(-1) # mace returns (N,1); ours (N,) - -e0_diff = (oe - re).abs().max().item() - -# ---- ScaleShift / GlobalRescale (single head) ---- -scale, shift = 0.731, -1.234 -x = torch.randn(7, dtype=torch.float64) * 10 -ours_ss = rescale.GlobalRescale(scale=scale, shift=shift) -ox = ours_ss(x) -ref_ss = ScaleShiftBlock(scale=scale, shift=shift) -head = torch.zeros(7, dtype=torch.long) -rx = ref_ss(x, head) -ss_diff = (ox - rx).abs().max().item() - -print(f"AtomicReferenceEnergy (E0) max|diff| = {e0_diff:.3e}") -print(f"GlobalRescale (scale/shift) max|diff| = {ss_diff:.3e}") -ok = e0_diff < 1e-12 and ss_diff < 1e-12 -print("RESULT:", "PASS" if ok else "FAIL") -sys.exit(0 if ok else 1) diff --git a/scripts/omol_port/verify_e2e.py b/scripts/omol_port/verify_e2e.py deleted file mode 100644 index 5d7915d..0000000 --- a/scripts/omol_port/verify_e2e.py +++ /dev/null @@ -1,73 +0,0 @@ -import importlib.util, sys, types, torch -torch.set_default_dtype(torch.float64); torch.manual_seed(0) -SRC = "/nobackup/proj/disk/teoroo/personal/jicli594/work/molcrafts/molnex/src" -molix = types.ModuleType("molix"); molix.config = types.SimpleNamespace(ftype=torch.float64) -sys.modules["molix"] = molix; sys.modules["molix.config"] = molix.config - - -def pkg(name, path): - m = types.ModuleType(name); m.__path__ = [path]; sys.modules[name] = m - - -for n, p in [("molrep", f"{SRC}/molrep"), ("molrep.embedding", f"{SRC}/molrep/embedding"), - ("molrep.interaction", f"{SRC}/molrep/interaction"), ("molrep.readout", f"{SRC}/molrep/readout"), - ("molpot", f"{SRC}/molpot"), ("molpot.heads", f"{SRC}/molpot/heads"), - ("molzoo", f"{SRC}/molzoo")]: - pkg(n, p) - - -def load(name, path): - s = importlib.util.spec_from_file_location(name, path); m = importlib.util.module_from_spec(s) - sys.modules[name] = m; s.loader.exec_module(m); return m - - -for name, sub in [("molrep.embedding.angular", "molrep/embedding/angular.py"), - ("molrep.embedding.cutoff", "molrep/embedding/cutoff.py"), - ("molrep.embedding.node", "molrep/embedding/node.py"), - ("molrep.embedding.radial", "molrep/embedding/radial.py"), - ("molrep.interaction.gate", "molrep/interaction/gate.py"), - ("molrep.interaction.radial", "molrep/interaction/radial.py"), - ("molrep.interaction.product_basis", "molrep/interaction/product_basis.py"), - ("molrep.interaction.residual", "molrep/interaction/residual.py"), - ("molrep.readout.scalar", "molrep/readout/scalar.py"), - ("molpot.heads.energy", "molpot/heads/energy.py"), - ("molpot.heads.rescale", "molpot/heads/rescale.py")]: - load(name, f"{SRC}/{sub}") -mo = load("molzoo.mace_omol", f"{SRC}/molzoo/mace_omol.py") -MACEOMol, load_omol_state_dict = mo.MACEOMol, mo.load_omol_state_dict - -cueq = torch.load("/nobackup/proj/disk/teoroo/personal/jicli594/work/mace_models/OMOL-cueq.model", - map_location="cpu", weights_only=False).double() -ztab = cueq.atomic_numbers.tolist() -ae = cueq.atomic_energies_fn.atomic_energies.flatten().double() -mdl = MACEOMol(atomic_numbers=ztab, atomic_energies=ae, - scale=float(cueq.scale_shift.scale), shift=float(cueq.scale_shift.shift)).double() -miss, unexp = load_omol_state_dict(mdl, cueq.state_dict()) -print("missing learnable:", miss[:15]) - -pos = torch.tensor([[0., 0., 0.], [1.0, 0., 0.], [0., 1.1, 0.], [0.5, 0.5, 0.9]], dtype=torch.float64) -Z = torch.tensor([6, 1, 8, 7]); N = 4 -s_, d_ = [], [] -for i in range(N): - for j in range(N): - if i != j and (pos[i] - pos[j]).norm() < 6.0: - s_.append(i); d_.append(j) -ei = torch.tensor([s_, d_]); batch = torch.zeros(N, dtype=torch.long) -tc = torch.tensor([0.]); ts = torch.tensor([1.]) - -out = mdl.energy_forces(pos, Z, ei, batch, tc, ts, compute_forces=True) -E_mine = out["energy"].item(); F_mine = out["forces"] - -zi = torch.searchsorted(cueq.atomic_numbers, Z) -na = torch.zeros(N, len(ztab), dtype=torch.float64); na[torch.arange(N), zi] = 1.0 -posg = pos.clone().requires_grad_(True) -data = {"positions": posg, "node_attrs": na, "edge_index": ei, - "shifts": torch.zeros(ei.shape[1], 3, dtype=torch.float64), - "unit_shifts": torch.zeros(ei.shape[1], 3, dtype=torch.float64), - "batch": batch, "ptr": torch.tensor([0, N]), "cell": torch.zeros(3, 3, dtype=torch.float64), - "head": torch.tensor([0]), "total_charge": tc, "total_spin": ts} -res = cueq(data, compute_force=True, training=False) -E_ref = res["energy"].item(); F_ref = res["forces"] -print(f"E_mine={E_mine:.8f} E_ref={E_ref:.8f} |dE|={abs(E_mine - E_ref):.3e}") -print(f"max|dF|={(F_mine - F_ref).abs().max().item():.3e}") -print("RESULT:", "PASS" if abs(E_mine - E_ref) < 1e-5 and (F_mine - F_ref).abs().max().item() < 1e-4 and not miss else "FAIL") diff --git a/scripts/omol_port/verify_interaction.py b/scripts/omol_port/verify_interaction.py deleted file mode 100644 index afd5de3..0000000 --- a/scripts/omol_port/verify_interaction.py +++ /dev/null @@ -1,79 +0,0 @@ -"""CPU test: molrep.ResidualInteraction vs cueq OMOL interactions[0].""" -import importlib.util -import sys -import types - -import torch - -torch.set_default_dtype(torch.float64) -torch.manual_seed(0) - -SRC = "/nobackup/proj/disk/teoroo/personal/jicli594/work/molcrafts/molnex/src" -CUEQ = "/nobackup/proj/disk/teoroo/personal/jicli594/work/mace_models/OMOL-cueq.model" - -molix = types.ModuleType("molix") -molix.config = types.SimpleNamespace(ftype=torch.float64) -sys.modules["molix"] = molix -sys.modules["molix.config"] = molix.config - - -def load(name, path): - spec = importlib.util.spec_from_file_location(name, path) - m = importlib.util.module_from_spec(spec) - sys.modules[name] = m - spec.loader.exec_module(m) - return m - - -# package shim so residual.py's `from .gate import` / `from .radial import` resolve -pkg = types.ModuleType("molrep_inter") -pkg.__path__ = [f"{SRC}/molrep/interaction"] -sys.modules["molrep_inter"] = pkg -load("molrep_inter.gate", f"{SRC}/molrep/interaction/gate.py") -load("molrep_inter.radial", f"{SRC}/molrep/interaction/radial.py") -res = load("molrep_inter.residual", f"{SRC}/molrep/interaction/residual.py") - -it = torch.load(CUEQ, map_location="cpu", weights_only=False).double().interactions[0] - -ours = res.ResidualInteraction( - node_attrs_irreps=str(it.node_attrs_irreps), - node_feats_irreps=str(it.node_feats_irreps), - edge_attrs_irreps=str(it.edge_attrs_irreps), - edge_feats_irreps=str(it.edge_feats_irreps), - edge_irreps=str(it.edge_irreps), - target_irreps=str(it.target_irreps), - hidden_irreps=str(it.hidden_irreps), - radial_mlp=list(it.radial_MLP), -).double() - -missing, unexpected = ours.load_state_dict(it.state_dict(), strict=False) -miss_params = [ - k for k in missing if k.endswith((".weight", ".bias")) or k in ("alpha", "beta") -] -print("missing learnable keys:", miss_params[:20]) - -ours.eval() -N, E = 7, 20 -n_attrs = torch.zeros(N, 83, dtype=torch.float64) -n_attrs[torch.arange(N), torch.randint(0, 83, (N,))] = 1.0 -n_feats = torch.randn(N, it.node_feats_irreps.dim, dtype=torch.float64) -e_attrs = torch.randn(E, it.edge_attrs_irreps.dim, dtype=torch.float64) -e_feats = torch.randn(E, 8, dtype=torch.float64) -edge_index = torch.randint(0, N, (2, E)) -cutoff = torch.rand(E, 1, dtype=torch.float64) - -with torch.no_grad(): - m_ours, sc_ours = ours(n_attrs, n_feats, e_attrs, e_feats, edge_index, cutoff) - m_ref, sc_ref = it( - node_attrs=n_attrs, node_feats=n_feats, edge_attrs=e_attrs, - edge_feats=e_feats, edge_index=edge_index, cutoff=cutoff, first_layer=True, - ) - -dm = (m_ours - m_ref).abs().max().item() -ds = (sc_ours - sc_ref).abs().max().item() -print("message shape", tuple(m_ours.shape), "ref", tuple(m_ref.shape)) -print(f"message max|diff| = {dm:.3e}") -print(f"skip max|diff| = {ds:.3e}") -ok = dm < 1e-9 and ds < 1e-9 and not miss_params -print("RESULT:", "PASS" if ok else "FAIL") -sys.exit(0 if ok else 1) diff --git a/scripts/omol_port/verify_joint_embed.py b/scripts/omol_port/verify_joint_embed.py deleted file mode 100644 index 66ab7d8..0000000 --- a/scripts/omol_port/verify_joint_embed.py +++ /dev/null @@ -1,78 +0,0 @@ -"""CPU test: molrep.JointFeatureEmbedding vs OMOL's real joint_embedding submodule.""" -import importlib.util -import sys -import types - -import torch - -torch.set_default_dtype(torch.float64) -torch.manual_seed(0) - -SRC = "/nobackup/proj/disk/teoroo/personal/jicli594/work/molcrafts/molnex/src" -PATH = "/nobackup/proj/disk/teoroo/personal/jicli594/work/mace_models/MACE-omol-0-extra-large-1024.model" - -molix = types.ModuleType("molix") -molix.config = types.SimpleNamespace(ftype=torch.float64) -sys.modules["molix"] = molix -sys.modules["molix.config"] = molix.config - - -def load(name, path): - spec = importlib.util.spec_from_file_location(name, path) - mod = importlib.util.module_from_spec(spec) - sys.modules[name] = mod # register so pydantic can resolve module globals - spec.loader.exec_module(mod) - return mod - - -node = load("molrep_node", f"{SRC}/molrep/embedding/node.py") - -# load OMOL, grab the real joint_embedding -src = torch.load(PATH, map_location="cpu", weights_only=False).double() -je = src.joint_embedding -print("OMOL joint_embedding.specs:", je.specs) -print("embedders:", {k: type(v).__name__ for k, v in je.embedders.items()}) - -# build matching specs for ours -specs = [] -for name, spec in je.specs.items(): - specs.append( - node.JointFeatureSpec( - name=name, - kind=spec["type"], - emb_dim=spec["emb_dim"], - num_classes=spec.get("num_classes"), - per=spec.get("per", "graph"), - offset=spec.get("offset", 0), - ) - ) -out_dim = je.project[0].weight.shape[0] -ours = node.JointFeatureEmbedding(feature_specs=specs, out_dim=out_dim).double() - -# copy weights (direct, matching key names) -ours.load_state_dict(je.state_dict()) -ours.eval() - -# build inputs: 3 graphs, 7 atoms -batch = torch.tensor([0, 0, 1, 1, 1, 2, 2]) -B = 3 -feats = {} -for name, spec in je.specs.items(): - if spec["type"] == "categorical": - # value range pre-offset: choose valid post-offset indices - off = spec.get("offset", 0) - nc = spec["num_classes"] - feats[name] = torch.randint(-off, nc - off, (B,)) - else: - feats[name] = torch.randn(B) - -with torch.no_grad(): - y_ours = ours(batch, **feats) - # OMOL forward signature: (batch, features_dict) - y_ref = je(batch, {k: v for k, v in feats.items()}) - -diff = (y_ours - y_ref).abs().max().item() -print("out shape", tuple(y_ours.shape)) -print(f"JointFeatureEmbedding max|diff| = {diff:.3e}") -print("RESULT:", "PASS" if diff < 1e-10 else "FAIL") -sys.exit(0 if diff < 1e-10 else 1) diff --git a/scripts/omol_port/verify_omol_cueq_equiv.py b/scripts/omol_port/verify_omol_cueq_equiv.py deleted file mode 100644 index f74324b..0000000 --- a/scripts/omol_port/verify_omol_cueq_equiv.py +++ /dev/null @@ -1,54 +0,0 @@ -"""Verify mace e3nn OMOL vs its cueq-converted twin give identical E/F on CPU. - -This establishes the cueq model (cuet primitives, ir_mul — the same stack molnex -uses) as a faithful reference for the molnex port. -""" -import os -import sys -import tempfile - -import numpy as np -import torch - -torch.set_default_dtype(torch.float64) - -PATH = "/nobackup/proj/disk/teoroo/personal/jicli594/work/mace_models/MACE-omol-0-extra-large-1024.model" - -from ase import Atoms # noqa: E402 -from mace.calculators import MACECalculator # noqa: E402 -from mace.cli.convert_e3nn_cueq import run as to_cueq # noqa: E402 - -# small molecule (water); OMOL wants charge/spin metadata -atoms = Atoms( - "H2O", - positions=[[0.0, 0.0, 0.0], [0.96, 0.0, 0.0], [-0.24, 0.93, 0.0]], -) -atoms.info["charge"] = 0 -atoms.info["spin"] = 1 - -# --- e3nn reference --- -e3nn_calc = MACECalculator(model_paths=PATH, device="cpu", default_dtype="float64", head="omol") -atoms.calc = e3nn_calc -e_e3nn = atoms.get_potential_energy() -f_e3nn = atoms.get_forces() - -# --- convert to cueq, save, build calc --- -src = torch.load(PATH, map_location="cpu", weights_only=False).double() -cueq = to_cueq(src, device="cpu", return_model=True) -tmp = tempfile.NamedTemporaryFile(suffix=".model", delete=False) -torch.save(cueq, tmp.name) -cueq_calc = MACECalculator(model_paths=tmp.name, device="cpu", default_dtype="float64", head="omol") -atoms.calc = cueq_calc -e_cueq = atoms.get_potential_energy() -f_cueq = atoms.get_forces() -os.unlink(tmp.name) - -de = abs(e_e3nn - e_cueq) -df = np.abs(f_e3nn - f_cueq).max() -print(f"E(e3nn) = {e_e3nn:.8f} eV") -print(f"E(cueq) = {e_cueq:.8f} eV") -print(f"|dE| = {de:.3e} eV") -print(f"max|dF| = {df:.3e} eV/Ang") -ok = de < 1e-5 and df < 1e-5 -print("RESULT:", "PASS" if ok else "FAIL") -sys.exit(0 if ok else 1) diff --git a/scripts/omol_port/verify_product.py b/scripts/omol_port/verify_product.py deleted file mode 100644 index 73ed7c0..0000000 --- a/scripts/omol_port/verify_product.py +++ /dev/null @@ -1,24 +0,0 @@ -import importlib.util, sys, types, torch -torch.set_default_dtype(torch.float64); torch.manual_seed(0) -SRC="/nobackup/proj/disk/teoroo/personal/jicli594/work/molcrafts/molnex/src" -molix=types.ModuleType("molix"); molix.config=types.SimpleNamespace(ftype=torch.float64) -sys.modules["molix"]=molix; sys.modules["molix.config"]=molix.config -def load(n,p): - s=importlib.util.spec_from_file_location(n,p); m=importlib.util.module_from_spec(s); sys.modules[n]=m; s.loader.exec_module(m); return m -pkg=types.ModuleType("mi"); pkg.__path__=[f"{SRC}/molrep/interaction"]; sys.modules["mi"]=pkg -pb=load("mi.product_basis", f"{SRC}/molrep/interaction/product_basis.py") -cueq=torch.load("/nobackup/proj/disk/teoroo/personal/jicli594/work/mace_models/OMOL-cueq.model",map_location="cpu",weights_only=False).double() -for li,p in enumerate(cueq.products): - sc_m=p.symmetric_contractions - iin=str(sc_m.irreps_in); iout=str(sc_m.irreps_out) - ours=pb.EquivariantProductBasis(node_feats_irreps=iin,target_irreps=iout,correlation=sc_m.contraction_degree,num_elements=1,use_sc=True).double() - miss,unexp=ours.load_state_dict(p.state_dict(),strict=False) - mp=[k for k in miss if k.endswith(".weight")] - import cuequivariance as cue - N=6; irdim=cue.Irreps("O3",iin).dim; mul=[mi.mul for mi in cue.Irreps("O3",iin)][0] - nf=torch.randn(N, irdim//mul, mul, dtype=torch.float64) # (N, ir_dim, mul) - scv=torch.randn(N, cue.Irreps("O3",iout).dim, dtype=torch.float64) - na=torch.zeros(N,83,dtype=torch.float64); na[torch.arange(N),torch.randint(0,83,(N,))]=1.0 - with torch.no_grad(): - yo=ours(nf,scv,na); yr=p(nf,scv,na) - print(f"product{li}: deg={sc_m.contraction_degree} miss{mp} max|diff|={(yo-yr).abs().max():.2e}") diff --git a/scripts/omol_port/verify_radial.py b/scripts/omol_port/verify_radial.py deleted file mode 100644 index d5c5e76..0000000 --- a/scripts/omol_port/verify_radial.py +++ /dev/null @@ -1,74 +0,0 @@ -"""CPU unit test: molrep MACEBesselBasis + PolynomialCutoff vs official mace. - -Loads the plain-torch molrep modules in isolation (stubbing `molix.config`) -so no cuequivariance / molcfg is needed, and checks bit-level agreement with -mace.modules.radial on float64. -""" -import importlib.util -import sys -import types - -import torch - -torch.set_default_dtype(torch.float64) - -SRC = "/nobackup/proj/disk/teoroo/personal/jicli594/work/molcrafts/molnex/src" - -# --- stub molix.config so radial.py / cutoff.py import cleanly on CPU --- -molix = types.ModuleType("molix") -config_mod = types.SimpleNamespace(ftype=torch.float64) -molix.config = config_mod -sys.modules["molix"] = molix -sys.modules["molix.config"] = config_mod - - -def load(name, path): - spec = importlib.util.spec_from_file_location(name, path) - mod = importlib.util.module_from_spec(spec) - spec.loader.exec_module(mod) - return mod - - -radial = load("molrep_radial", f"{SRC}/molrep/embedding/radial.py") -cutoff = load("molrep_cutoff", f"{SRC}/molrep/embedding/cutoff.py") - -from mace.modules.radial import BesselBasis as RefBessel # noqa: E402 -from mace.modules.radial import PolynomialCutoff as RefCutoff # noqa: E402 - -R_MAX = 6.0 -N = 8 -r = torch.linspace(0.05, R_MAX + 1.0, 500, dtype=torch.float64) # include r > r_max - -# ---- Bessel ---- (BesselRBF in MACE-faithful mode: no norm, no eps, trainable) -ours_b = radial.BesselRBF( - r_cut=R_MAX, num_radial=N, normalize=False, eps=0.0, trainable=True -) -ref_b = RefBessel(r_max=R_MAX, num_basis=N, trainable=True) -# init must already match (same formula); measure BEFORE any copy -init_diff = (ours_b.freqs.detach() - ref_b.bessel_weights.detach()).abs().max().item() -ob = ours_b(r) -rb = ref_b(r.unsqueeze(-1)) -bessel_diff = (ob - rb).abs().max().item() - -# ---- PolynomialCutoff ---- -ours_c = cutoff.PolynomialCutoff(r_cut=R_MAX, exponent=6) -ref_c = RefCutoff(r_max=R_MAX, p=6) -oc = ours_c(r) -rc = ref_c(r) -cut_diff = (oc - rc).abs().max().item() - -# Note: ref cutoff does NOT zero out r > r_max (no mask); ours masks. -# Check agreement on r < r_max region (physical edges only). -in_cut = r < R_MAX -cut_diff_incut = (oc[in_cut] - rc[in_cut]).abs().max().item() - -print(f"bessel_weights init max|diff| = {init_diff:.3e}") -print(f"Bessel max|diff| (all r) = {bessel_diff:.3e}") -print(f"PolyCutoff max|diff| (r=r_max, ref does not)") -print(f"ref cutoff at r=r_max+0.5 = {ref_c(torch.tensor([R_MAX+0.5]))[0].item():.3e}") -print(f"our cutoff at r=r_max+0.5 = {ours_c(torch.tensor([R_MAX+0.5]))[0].item():.3e}") - -ok = bessel_diff < 1e-12 and cut_diff_incut < 1e-12 and init_diff < 1e-12 -print("RESULT:", "PASS" if ok else "FAIL") -sys.exit(0 if ok else 1) diff --git a/scripts/omol_port/verify_radial_mlp.py b/scripts/omol_port/verify_radial_mlp.py deleted file mode 100644 index 6cc2bfc..0000000 --- a/scripts/omol_port/verify_radial_mlp.py +++ /dev/null @@ -1,42 +0,0 @@ -"""CPU test: molrep.interaction.RadialMLP vs official mace RadialMLP (direct copy).""" -import importlib.util -import sys -import types - -import torch - -torch.set_default_dtype(torch.float64) -torch.manual_seed(0) - -SRC = "/nobackup/proj/disk/teoroo/personal/jicli594/work/molcrafts/molnex/src" - -molix = types.ModuleType("molix") -molix.config = types.SimpleNamespace(ftype=torch.float64) -sys.modules["molix"] = molix -sys.modules["molix.config"] = molix.config - - -def load(name, path): - spec = importlib.util.spec_from_file_location(name, path) - mod = importlib.util.module_from_spec(spec) - sys.modules[name] = mod - spec.loader.exec_module(mod) - return mod - - -radial = load("mr_radial", f"{SRC}/molrep/interaction/radial.py") -from mace.modules.radial import RadialMLP as MaceRMLP # noqa: E402 - -ch = [2056, 128, 128, 128, 4096] # OMOL conv_tp_weights shape -ours = radial.RadialMLP(ch).double() -ref = MaceRMLP(ch).double() -ours.load_state_dict(ref.state_dict()) -ours.eval() -ref.eval() - -x = torch.randn(20, 2056, dtype=torch.float64) -with torch.no_grad(): - d = (ours(x) - ref(x)).abs().max().item() -print(f"RadialMLP max|diff| = {d:.3e}") -print("RESULT:", "PASS" if d < 1e-12 else "FAIL") -sys.exit(0 if d < 1e-12 else 1) diff --git a/scripts/omol_port/verify_readout.py b/scripts/omol_port/verify_readout.py deleted file mode 100644 index aa8c9b0..0000000 --- a/scripts/omol_port/verify_readout.py +++ /dev/null @@ -1,19 +0,0 @@ -import importlib.util, sys, types, torch -torch.set_default_dtype(torch.float64); torch.manual_seed(0) -SRC="/nobackup/proj/disk/teoroo/personal/jicli594/work/molcrafts/molnex/src" -molix=types.ModuleType("molix"); molix.config=types.SimpleNamespace(ftype=torch.float64) -sys.modules["molix"]=molix; sys.modules["molix.config"]=molix.config -def load(n,p): - s=importlib.util.spec_from_file_location(n,p); m=importlib.util.module_from_spec(s); sys.modules[n]=m; s.loader.exec_module(m); return m -pkg=types.ModuleType("mro"); pkg.__path__=[f"{SRC}/molrep/readout"]; sys.modules["mro"]=pkg -sc=load("mro.scalar", f"{SRC}/molrep/readout/scalar.py") -cueq=torch.load("/nobackup/proj/disk/teoroo/personal/jicli594/work/mace_models/OMOL-cueq.model",map_location="cpu",weights_only=False).double() -r=cueq.readouts[0] -ours=sc.NonLinearBiasReadout(irreps_in="1024x0e", mlp_dim=16).double() -miss,unexp=ours.load_state_dict(r.state_dict(),strict=False) -mp=[k for k in miss if k.endswith((".weight",".bias"))] -N=7; x=torch.randn(N,1024,dtype=torch.float64) -with torch.no_grad(): - yo=ours(x); yr=r(x, torch.zeros(N,dtype=torch.long)) -print("miss",mp,"out",tuple(yo.shape),tuple(yr.shape)) -print(f"NonLinearBiasReadout max|diff| = {(yo-yr).abs().max():.3e}") diff --git a/src/molix/F/scatter.py b/src/molix/F/scatter.py index 8eb91b6..00128fa 100644 --- a/src/molix/F/scatter.py +++ b/src/molix/F/scatter.py @@ -6,6 +6,8 @@ from __future__ import annotations +import os + import torch from torch import Tensor @@ -44,39 +46,33 @@ def scatter_sum( return out.scatter_add_(dim, index, src) -# One-hot matmul allocates ``(E, dim_size)``. Above this element budget use -# ``index_add_`` to avoid OOM on large periodic systems (review: O(E·N) risk). -_ONEHOT_ELEMENT_BUDGET = 2_000_000 +# Read once at import, not per call: ``scatter_sum_compile_safe`` runs inside +# ``torch.compile`` regions on the MACE interaction hot path, where an +# ``os.environ`` lookup is both a per-call dict hit and an opaque side effect +# the compiler must guard against. +_ONEHOT_FLAG: bool = os.environ.get("MOLNEX_SCATTER_ONEHOT") == "1" def scatter_sum_compile_safe(src: Tensor, index: Tensor, dim_size: int) -> Tensor: """Sum ``src`` rows into ``dim_size`` buckets given by ``index`` (dim 0). - For small ``E * dim_size`` (typical molecular graphs) uses a one-hot matmul - that is bit-exact under ``torch.compile`` without global determinism flags - (see historical note: ``index_add_`` scatter order can differ under - inductor). For large systems (``E * dim_size > 2e6``) falls back to - ``index_add_`` to avoid O(E·N) memory blow-ups. - - Force exact one-hot always with env ``MOLNEX_SCATTER_ONEHOT=1``. - Force ``index_add_`` always with ``MOLNEX_SCATTER_ONEHOT=0``. + Uses ``index_add_`` — the O(E·D) scatter. The alternative one-hot matmul + (bit-exact under ``torch.compile`` without global determinism flags, since + ``index_add_``'s scatter order can differ under inductor) is an explicit + opt-in via env ``MOLNEX_SCATTER_ONEHOT=1``: it turns the reduction into an + O(E·N·D) GEMM measured 2.4–3.3x slower than ``index_add_`` at every + profiled molecular-graph shape, so exactness must be a deliberate choice, + never a size-heuristic default. The flag is read once at import. """ - import os - - n_src = src.shape[0] - budget = n_src * int(dim_size) - flag = os.environ.get("MOLNEX_SCATTER_ONEHOT") - use_onehot = flag == "1" or (flag is None and budget <= _ONEHOT_ELEMENT_BUDGET) - - if not use_onehot: - out_shape = (dim_size, *src.shape[1:]) - out = torch.zeros(out_shape, dtype=src.dtype, device=src.device) - return out.index_add_(0, index, src) - - onehot = (index.view(-1, 1) == torch.arange(dim_size, device=src.device).view(1, -1)).to( - src.dtype - ) - return (onehot.t() @ src.reshape(src.shape[0], -1)).reshape(dim_size, *src.shape[1:]) + if _ONEHOT_FLAG: + onehot = (index.view(-1, 1) == torch.arange(dim_size, device=src.device).view(1, -1)).to( + src.dtype + ) + return (onehot.t() @ src.reshape(src.shape[0], -1)).reshape(dim_size, *src.shape[1:]) + + out_shape = (dim_size, *src.shape[1:]) + out = torch.zeros(out_shape, dtype=src.dtype, device=src.device) + return out.index_add_(0, index, src) def batch_add(src: Tensor, batch: Tensor, dim_size: int | None = None) -> Tensor: diff --git a/src/molix/compile.py b/src/molix/compile.py index e6e3a3b..cc1d221 100644 --- a/src/molix/compile.py +++ b/src/molix/compile.py @@ -66,10 +66,10 @@ def __init__( """ if cuda_graphs: p = self.CUDA_GRAPH_PRESET - backend = p["backend"] # type: ignore[assignment] - fullgraph = p["fullgraph"] # type: ignore[assignment] - dynamic = p["dynamic"] # type: ignore[assignment] - mode = p["mode"] # type: ignore[assignment] + backend = p["backend"] + fullgraph = p["fullgraph"] + dynamic = p["dynamic"] + mode = p["mode"] self.backend = backend self.fullgraph = fullgraph self.dynamic = dynamic diff --git a/src/molix/core/checkpoint/state.py b/src/molix/core/checkpoint/state.py index 65f98d9..90aadc1 100644 --- a/src/molix/core/checkpoint/state.py +++ b/src/molix/core/checkpoint/state.py @@ -61,9 +61,9 @@ def _unwrap_model(self) -> nn.Module: """ model = self.model if hasattr(model, "_orig_mod"): # torch.compile OptimizedModule - model = model._orig_mod # type: ignore[union-attr] + model = model._orig_mod if hasattr(model, "module"): # DDP / FSDP - model = model.module # type: ignore[union-attr] + model = model.module return cast(nn.Module, model) # ------------------------------------------------------------------ diff --git a/src/molix/core/hook.py b/src/molix/core/hook.py index f08c958..947f351 100644 --- a/src/molix/core/hook.py +++ b/src/molix/core/hook.py @@ -39,7 +39,14 @@ class Hook(Protocol): - Use (hook, priority) tuples to override execution order - Lower priority values execute earlier (default priority = 100) - Hooks with same priority execute in registration order - - If a hook raises an exception, it is logged but training continues + + Errors are fatal, by design. If a hook raises, + :meth:`molix.core.trainer.Trainer._call_hooks` logs the exception with a + full traceback and then **re-raises** it, so the exception propagates out + of the training loop and the run stops. Nothing is swallowed: swallowing + hook errors once hid a fatal ``NaNGuardHook`` signal behind thousands of + lines of repeated tracebacks. A hook that must tolerate its own failures + has to catch them itself. Example: ```python @@ -166,14 +173,21 @@ def on_eval_batch_end( ... def on_eval_step_complete(self, trainer: "Trainer", state: "TrainState") -> None: - """Called after step-based evaluation completes (not on epoch-end eval). + """Called once at the end of every eval phase, after the batch loop. - This hook is only triggered when eval runs due to the eval_every_n_steps - parameter being reached. Epoch-end evals do not trigger this hook. + Fires for *both* triggers — step-based eval (``eval_every_n_steps``) + and epoch-end eval — because :meth:`molix.core.trainer.Trainer._run_eval_phase` + is the single code path behind both. Eval-publishing hooks + (:class:`molix.hooks.MetricsHook`, :class:`molix.hooks.TensorBoardHook`) + should write their ``state["eval"]`` scalars here rather than in + ``on_epoch_end``, so the values are already published when the LR + scheduler reads its monitored metric. Args: trainer: The trainer instance - state: Current training state (steps_since_last_eval reset to 0) + state: Current training state (every eval batch consumed; + ``steps_since_last_eval`` is reset by the train loop *after* + the eval phase returns, so it is still non-zero here) """ ... @@ -237,7 +251,7 @@ def on_eval_batch_end( pass def on_eval_step_complete(self, trainer: "Trainer", state: "TrainState") -> None: - """Called after step-based evaluation completes (not on epoch-end eval).""" + """Called once at the end of every eval phase (step-based *and* epoch-end).""" pass diff --git a/src/molix/core/losses/__init__.py b/src/molix/core/losses/__init__.py index 526efd9..1509522 100644 --- a/src/molix/core/losses/__init__.py +++ b/src/molix/core/losses/__init__.py @@ -14,7 +14,13 @@ from molix.core.losses.combined import WeightedLoss from molix.core.losses.energy import MSELoss from molix.core.losses.force import MAELoss -from molix.core.losses.molecular import energy_force_mse, energy_mse +from molix.core.losses.molecular import ( + center_by_group, + energy_force_mse, + energy_mse, + molecule_centered_energy_mse, + parameter_bag_mse, +) __all__ = [ "MAELoss", @@ -22,4 +28,7 @@ "WeightedLoss", "energy_force_mse", "energy_mse", + "center_by_group", + "molecule_centered_energy_mse", + "parameter_bag_mse", ] diff --git a/src/molix/core/losses/molecular.py b/src/molix/core/losses/molecular.py index 81bf44f..e8f19cf 100644 --- a/src/molix/core/losses/molecular.py +++ b/src/molix/core/losses/molecular.py @@ -31,7 +31,106 @@ def loss_fn(preds, batch): import torch import torch.nn as nn -__all__ = ["energy_mse", "energy_force_mse"] +__all__ = [ + "energy_mse", + "energy_force_mse", + "center_by_group", + "molecule_centered_energy_mse", + "parameter_bag_mse", +] + + +def center_by_group( + values: torch.Tensor, + group_ids: torch.Tensor, +) -> torch.Tensor: + """Subtract per-group means from *values* (Espaloma molecule centering). + + Args: + values: Scalar energies ``(B,)`` or ``(B, 1)``. + group_ids: Integer molecule / group ids ``(B,)`` aligned with *values*. + + Returns: + Centered values with the same shape as *values*. + """ + flat = values.reshape(-1) + ids = group_ids.reshape(-1).long() + if flat.numel() != ids.numel(): + raise ValueError(f"values and group_ids length mismatch: {flat.numel()} vs {ids.numel()}") + out = flat.clone() + for g in ids.unique(): + mask = ids == g + out[mask] = flat[mask] - flat[mask].mean() + return out.view_as(values) + + +def molecule_centered_energy_mse( + target_key: str = "mm_energy", + *, + pred_key: str = "energy", + group_key: str = "molecule_id_index", + reduction: str = "mean", +) -> Callable[[Mapping[str, Any], Any], torch.Tensor]: + """MSE after independent molecule-mean centering of pred and ref energies. + + Both predictions and targets are centered **per molecule** before the MSE + (Espaloma relative conformational energy protocol). Units: kcal/mol. + + Args: + target_key: Graph-level target at ``batch["graphs", target_key]``. + pred_key: Model energy key in ``preds``. + group_key: Graph-level integer molecule index at + ``batch["graphs", group_key]``. + reduction: Forwarded to :class:`torch.nn.MSELoss`. + + Returns: + ``loss_fn(preds, batch) -> Tensor``. + """ + mse = nn.MSELoss(reduction=reduction) + + def _fn(preds: Mapping[str, Any], batch: Any) -> torch.Tensor: + e_pred = preds[pred_key] + e_true = batch["graphs", target_key].view_as(e_pred) + groups = batch["graphs", group_key] + return mse(center_by_group(e_pred, groups), center_by_group(e_true, groups)) + + return _fn + + +def parameter_bag_mse( + pred_bags: Mapping[str, torch.Tensor], + ref_bags: Mapping[str, torch.Tensor], + *, + reduction: str = "mean", +) -> torch.Tensor: + """Optional B1 diagnostic: MSE between predicted and reference parameter tensors. + + Args: + pred_bags: Mapping name → predicted parameter tensor. + ref_bags: Mapping name → reference parameter tensor (same keys/shapes). + reduction: ``mean`` / ``sum`` / ``none`` over concatenated squared errors. + + Returns: + Scalar (or unreduced) MSE over all bag entries. + """ + sq: list[torch.Tensor] = [] + for key, pref in pred_bags.items(): + if key not in ref_bags: + raise KeyError(f"ref_bags missing key {key!r}") + ref = ref_bags[key] + if pref.shape != ref.shape: + raise ValueError(f"bag {key!r} shape {pref.shape} != ref {ref.shape}") + sq.append((pref - ref).reshape(-1).pow(2)) + if not sq: + return torch.zeros(()) + cat = torch.cat(sq) + if reduction == "mean": + return cat.mean() + if reduction == "sum": + return cat.sum() + if reduction == "none": + return cat + raise ValueError(f"Unknown reduction {reduction!r}") def energy_mse( diff --git a/src/molix/core/metrics.py b/src/molix/core/metrics.py index 6a640e3..972374e 100644 --- a/src/molix/core/metrics.py +++ b/src/molix/core/metrics.py @@ -371,7 +371,7 @@ def update(self, preds: torch.Tensor, targets: torch.Tensor) -> None: for metric in self.metrics.values(): metric.update(preds, targets) - def compute(self) -> dict[str, torch.Tensor]: # type: ignore[return] + def compute(self) -> dict[str, torch.Tensor]: """Compute all metrics, each a 0-d tensor on the inputs' device. No ``.item()`` is taken here — callers materialise to Python floats @@ -397,3 +397,75 @@ def to(self, device: str | torch.device) -> MetricCollection: if hasattr(metric, "to"): metric.to(device) return self + + +class MoleculeCenteredRMSE(BaseMetric): + """RMSE after per-molecule mean-centering of pred and target (kcal/mol). + + Espaloma relative conformational energy metric. Call + :meth:`update` with ``preds``, ``targets``, and integer ``group_ids``. + """ + + def __init__(self) -> None: + super().__init__() + self.reset() + + def update( + self, + preds: torch.Tensor, + targets: torch.Tensor, + group_ids: torch.Tensor | None = None, + ) -> None: + if group_ids is None: + raise ValueError("MoleculeCenteredRMSE.update requires group_ids") + from molix.core.losses.molecular import center_by_group + + p = center_by_group(preds.detach().reshape(-1), group_ids.detach().reshape(-1)) + t = center_by_group(targets.detach().reshape(-1), group_ids.detach().reshape(-1)) + self.preds.append(p) + self.targets.append(t) + + def compute(self) -> torch.Tensor: + if not self.preds: + return torch.zeros(()) + p = torch.cat(self.preds) + t = torch.cat(self.targets) + return torch.sqrt(torch.mean((p - t) ** 2)) + + def reset(self) -> None: + self.preds: list[torch.Tensor] = [] + self.targets: list[torch.Tensor] = [] + + +class MoleculeCenteredMAE(BaseMetric): + """MAE after per-molecule mean-centering of pred and target (kcal/mol).""" + + def __init__(self) -> None: + super().__init__() + self.reset() + + def update( + self, + preds: torch.Tensor, + targets: torch.Tensor, + group_ids: torch.Tensor | None = None, + ) -> None: + if group_ids is None: + raise ValueError("MoleculeCenteredMAE.update requires group_ids") + from molix.core.losses.molecular import center_by_group + + p = center_by_group(preds.detach().reshape(-1), group_ids.detach().reshape(-1)) + t = center_by_group(targets.detach().reshape(-1), group_ids.detach().reshape(-1)) + self.preds.append(p) + self.targets.append(t) + + def compute(self) -> torch.Tensor: + if not self.preds: + return torch.zeros(()) + p = torch.cat(self.preds) + t = torch.cat(self.targets) + return torch.mean((p - t).abs()) + + def reset(self) -> None: + self.preds: list[torch.Tensor] = [] + self.targets: list[torch.Tensor] = [] diff --git a/src/molix/core/state.py b/src/molix/core/state.py index e8c5046..1ce708e 100644 --- a/src/molix/core/state.py +++ b/src/molix/core/state.py @@ -240,7 +240,7 @@ def best_metric(self, value: float | None) -> None: #: A path into :class:`TrainState`. A bare ``str`` names a top-level key #: (``"epoch"``); a ``tuple`` walks the namespace hierarchy -#: (``("train", "loss")``). Used by :class:`molix.core.hooks.Log`, +#: (``("train", "loss")``). Used by :class:`molix.hooks.Log`, #: :class:`CheckpointHook`, and the LR scheduler metric lookup. Path = str | tuple[str, ...] diff --git a/src/molix/data/__init__.py b/src/molix/data/__init__.py index ee730a0..e3f8b5f 100644 --- a/src/molix/data/__init__.py +++ b/src/molix/data/__init__.py @@ -37,6 +37,7 @@ PackedView, SubsetDataset, ) +from molix.data.group_split import group_split_indices # Pipeline DSL from molix.data.pipeline import DAGCache, Edge, Node, Pipeline, PipelineSpec @@ -102,4 +103,6 @@ "collate_molecules", "TargetSchema", "DEFAULT_TARGET_SCHEMA", + # Group splits + "group_split_indices", ] diff --git a/src/molix/data/cache.py b/src/molix/data/cache.py index 14cf978..66bdf22 100644 --- a/src/molix/data/cache.py +++ b/src/molix/data/cache.py @@ -63,6 +63,10 @@ # Reserved keys in the packed payload — never collide with user sample keys. +# Note: ``angles`` / ``propers`` / ``impropers`` are intentionally *not* +# reserved as top-level sample keys — pre-collate samples carry nested column +# dicts under those names; packing extracts them into payload buckets after +# flatten (dotted ``angles.atomi`` etc.). Only the ptr keys are reserved. _RESERVED_TOP_KEYS = frozenset( { "format_version", @@ -70,6 +74,9 @@ "atom_ptr", "edge_ptr", "bond_ptr", + "angle_ptr", + "proper_ptr", + "improper_ptr", "atoms", "edges", "bonds", @@ -323,6 +330,12 @@ def _pack_samples(samples: list[dict]) -> dict[str, Any]: # atom/edge/graph leading-dim model — pack it into a dedicated bonds bucket. bonds_bucket, bond_ptr = _extract_bonds(flats) + # Valence column families (angles / propers / impropers): 1-D columns of + # variable term-count, extracted as dedicated buckets with their own ptrs. + valence_extracted: dict[str, tuple[dict[str, torch.Tensor], torch.Tensor | None]] = {} + for family, required in _VALENCE_REQUIRED.items(): + valence_extracted[family] = _extract_valence_family(flats, family, required) + keys = set(flats[0].keys()) for i, f in enumerate(flats[1:], start=1): if set(f.keys()) != keys: @@ -363,12 +376,22 @@ def _pack_samples(samples: list[dict]) -> dict[str, Any]: "graphs": {k: torch.stack([f[k] for f in flats], dim=0) for k in graph_keys}, "scalars": {k: [f[k] for f in flats] for k in scalar_keys}, } - if atom_keys: + # Pointers are written from the *reference* keys (``Z`` / ``edge_index``), + # not from bucket emptiness: they describe the sample partition itself, and + # every consumer (``MmapDataset.avg_num_neighbors`` / ``edge_counts``, + # ``collate_packed``) needs them whenever the reference key exists. Keying + # them on bucket emptiness made a misclassified key silently erase the whole + # axis. + if has_atom_ref: payload["atom_ptr"] = torch.tensor(atom_ptr, dtype=torch.long) - if edge_keys: + if has_edge_ref: payload["edge_ptr"] = torch.tensor(edge_ptr, dtype=torch.long) if bond_ptr is not None: payload["bond_ptr"] = bond_ptr + for family, (bucket, ptr) in valence_extracted.items(): + if ptr is not None: + payload[family] = bucket + payload[_VALENCE_PTR_KEY[family]] = ptr return payload @@ -376,6 +399,24 @@ def _pack_samples(samples: list[dict]) -> dict[str, Any]: _BOND_INDEX_KEY = "bond_index" _BOND_TYPES_KEY = "bond_types" +# Valence column families (nested pre-collate → dotted flat keys after _flatten). +# Impropers: atomi is the center (molrs center-first). +_VALENCE_REQUIRED: dict[str, tuple[str, ...]] = { + "angles": ("atomi", "atomj", "atomk"), + "propers": ("atomi", "atomj", "atomk", "atoml"), + "impropers": ("atomi", "atomj", "atomk", "atoml"), +} +_VALENCE_OPTIONAL: tuple[str, ...] = ("type",) +_VALENCE_PTR_KEY: dict[str, str] = { + "angles": "angle_ptr", + "propers": "proper_ptr", + "impropers": "improper_ptr", +} + +# Canonical per-edge sample keys, routed by identity when the leading-dim +# classification is ambiguous (see :func:`_infer_schema_across`). +_EDGE_KEYS = frozenset({"edge_index", "edge_diff", "edge_dist"}) + def _extract_bonds( flats: list[dict[str, Any]], @@ -419,6 +460,86 @@ def _extract_bonds( return bonds, torch.tensor(bond_ptr, dtype=torch.long) +def _extract_valence_family( + flats: list[dict[str, Any]], + family: str, + required: tuple[str, ...], +) -> tuple[dict[str, torch.Tensor], torch.Tensor | None]: + """Pop dotted ``{family}.{col}`` keys into a packed column bucket + ptr. + + Pre-collate nested form flattens to ``angles.atomi`` etc. Columns are 1-D + of length ``N_terms`` (variable across samples), so they cannot ride the + atom/edge/graph leading-dim model. Mutates *flats* in place. + + Returns: + ``(bucket, ptr)`` where *bucket* maps column name → concatenated + tensor and *ptr* is the per-sample cumsum over ``N_terms``, or + ``({}, None)`` if no sample carries the family. + + Raises: + ValueError: mixed presence, missing required column, or length mismatch. + """ + lead_key = f"{family}.{required[0]}" + present = [lead_key in f and f[lead_key] is not None for f in flats] + if not any(present): + return {}, None + if not all(present): + raise ValueError( + f"{family!r} must be present in all samples or none " + f"(got presence={present})." + ) + + col_lists: dict[str, list[torch.Tensor]] = {c: [] for c in required} + opt_lists: dict[str, list[torch.Tensor]] = {c: [] for c in _VALENCE_OPTIONAL} + have_optional = {c: True for c in _VALENCE_OPTIONAL} + n_terms: list[int] = [] + + for f in flats: + lengths: dict[str, int] = {} + for col in required: + key = f"{family}.{col}" + if key not in f or f[key] is None: + raise ValueError( + f"sample is missing required {family!r} column {col!r}" + ) + t = f.pop(key).long().reshape(-1) + col_lists[col].append(t) + lengths[col] = int(t.shape[0]) + n0 = lengths[required[0]] + for col in required[1:]: + if lengths[col] != n0: + raise ValueError( + f"{family!r} column lengths differ: {col}={lengths[col]} " + f"vs {required[0]}={n0}" + ) + n_terms.append(n0) + for col in _VALENCE_OPTIONAL: + key = f"{family}.{col}" + val = f.pop(key, None) + if val is None: + have_optional[col] = False + else: + ot = val.long().reshape(-1) + if int(ot.shape[0]) != n0: + raise ValueError( + f"{family!r} optional {col!r} length {ot.shape[0]} " + f"!= n_terms={n0}" + ) + opt_lists[col].append(ot) + + bucket: dict[str, torch.Tensor] = { + col: torch.cat(ts, dim=0) for col, ts in col_lists.items() + } + for col, ts in opt_lists.items(): + if have_optional[col] and ts: + bucket[col] = torch.cat(ts, dim=0) + + ptr = [0] + for nt in n_terms: + ptr.append(ptr[-1] + nt) + return bucket, torch.tensor(ptr, dtype=torch.long) + + def _unpack_one(payload: Mapping[str, Any], idx: int) -> dict: n = payload["n_samples"] if idx < 0: @@ -449,6 +570,17 @@ def _unpack_one(payload: Mapping[str, Any], idx: int) -> dict: if _BOND_TYPES_KEY in bonds: flat[_BOND_TYPES_KEY] = bonds[_BOND_TYPES_KEY][b0:b1] + # Valence column buckets → dotted keys so _unflatten rebuilds nested dicts. + for family in _VALENCE_REQUIRED: + ptr_key = _VALENCE_PTR_KEY[family] + v_ptr = payload.get(ptr_key) + if v_ptr is None: + continue + v0, v1 = int(v_ptr[idx]), int(v_ptr[idx + 1]) + bucket = payload.get(family, {}) + for col, tensor in bucket.items(): + flat[f"{family}.{col}"] = tensor[v0:v1] + for k, t in payload["graphs"].items(): flat[k] = t[idx] @@ -575,8 +707,17 @@ def _infer_schema_across( tracks_atoms = has_atom_ref and all(s0 == na for s0, na in zip(shape0s, n_atoms)) tracks_edges = has_edge_ref and all(s0 == ne for s0, ne in zip(shape0s, n_edges)) - # Prefer atom classification when both track (can happen if n_atoms == n_edges - # holds across every sample, e.g. in degenerate cases). + # Leading-dim classification is ambiguous when n_edges == n_atoms holds for + # *every* sample — e.g. 2-atom molecules with a symmetric one-pair neighbour + # list (E = 2 = N). Break that tie by identity for the canonical edge keys: + # the whole downstream stack addresses them by name (``collate``'s edges + # namespace, ``_ref_len(f, "edge_index")`` above), and filing them under + # atoms also suppresses ``edge_ptr``, silently dropping every edge instead + # of merely relabelling it. ``_extract_bonds`` routes ``bond_index`` by + # identity for the same reason. All other ambiguous keys keep the atom + # preference (per-atom is the far more common intent). + if tracks_atoms and tracks_edges and k in _EDGE_KEYS: + tracks_atoms = False if tracks_atoms: rest_set = set(shape_rests) if len(rest_set) != 1: diff --git a/src/molix/data/collate.py b/src/molix/data/collate.py index d37277a..5d78371 100644 --- a/src/molix/data/collate.py +++ b/src/molix/data/collate.py @@ -60,10 +60,21 @@ class TargetSchema: INDEX_KEYS: dict[str, tuple[int, int]] = { "edge_index": (0, 1), # [E, 2] — count axis 0, both columns are atom indices "bond_index": (1, 0), # [2, N] COO — count axis 1, both rows are atom indices - "angle_index": (1, 0), # [3, N] - "dihedral_index": (1, 0), # [4, N] + "angle_index": (1, 0), # [3, N] (kernel-local COO only; not collate schema) + "dihedral_index": (1, 0), # [4, N] (kernel-local COO only; not collate schema) } +# Valence connectivity lives as 1-D atom-index columns under nested namespaces +# (molpy/molrs Frame style). Required columns per family; optional ``type``. +# Impropers: atomi is the **center** (molrs center-first). +_VALENCE_REQUIRED: dict[str, tuple[str, ...]] = { + "angles": ("atomi", "atomj", "atomk"), + "propers": ("atomi", "atomj", "atomk", "atoml"), + "impropers": ("atomi", "atomj", "atomk", "atoml"), +} +_VALENCE_OPTIONAL: tuple[str, ...] = ("type",) +_ATOM_INDEX_COLS: frozenset[str] = frozenset({"atomi", "atomj", "atomk", "atoml"}) + def rebase(tensor: torch.Tensor, offset: int | torch.Tensor, key: str) -> torch.Tensor: """Shift a registered atom-index tensor into global coordinates. @@ -97,6 +108,97 @@ def _normalize_edge_index(edge_index: torch.Tensor) -> torch.Tensor: raise ValueError(f"edge_index must have shape (E, 2) or (2, E), got {tuple(edge_index.shape)}") +def _sample_has_valence(sample: Mapping[str, Any], family: str) -> bool: + """Return True if *sample* carries a non-None nested *family* block.""" + block = sample.get(family) + return block is not None and isinstance(block, Mapping) + + +def _collate_valence_family( + samples: list[dict], + family: str, + required: tuple[str, ...], + atom_offsets: list[int], +) -> TensorDict | None: + """Collate one valence namespace as 1-D columns rebased by atom offset. + + Pre-collate form (preferred):: + + sample["angles"] = { + "atomi": Long[N], "atomj": Long[N], "atomk": Long[N], "type"?: Long[N] + } + + All-or-none: every sample must carry the family or none may. Atom-index + columns (``atomi``/``atomj``/``atomk``/``atoml``) are shifted by the + sample's atom offset; optional ``type`` is concatenated without offset. + + Returns: + Nested :class:`~tensordict.TensorDict` with ``batch_size=[N_terms]``, + or ``None`` when no sample carries the family. + + Raises: + ValueError: mixed presence, missing required columns, or length mismatch. + TypeError: family value is not a mapping of tensors. + """ + present = [_sample_has_valence(s, family) for s in samples] + if not any(present): + return None + if not all(present): + raise ValueError( + f"{family!r} must be present in all samples or none " + f"(got presence={present})." + ) + + cols: dict[str, list[torch.Tensor]] = {c: [] for c in required} + optional_lists: dict[str, list[torch.Tensor]] = {c: [] for c in _VALENCE_OPTIONAL} + have_optional = {c: True for c in _VALENCE_OPTIONAL} + + for sample, offset in zip(samples, atom_offsets): + block = sample[family] + if not isinstance(block, Mapping): + raise TypeError( + f"sample[{family!r}] must be a mapping of column tensors, " + f"got {type(block).__name__}" + ) + for col in required: + if col not in block or block[col] is None: + raise ValueError( + f"sample[{family!r}] missing required column {col!r}" + ) + t = torch.as_tensor(block[col]).long().reshape(-1) + if col in _ATOM_INDEX_COLS: + t = t + offset + cols[col].append(t) + n_terms = int(cols[required[0]][-1].shape[0]) + for col in required[1:]: + if int(cols[col][-1].shape[0]) != n_terms: + raise ValueError( + f"sample[{family!r}] column lengths differ: " + f"{col}={cols[col][-1].shape[0]} vs {required[0]}={n_terms}" + ) + for col in _VALENCE_OPTIONAL: + val = block.get(col) + if val is None: + have_optional[col] = False + else: + ot = torch.as_tensor(val).long().reshape(-1) + if int(ot.shape[0]) != n_terms: + raise ValueError( + f"sample[{family!r}] optional {col!r} length " + f"{ot.shape[0]} != n_terms={n_terms}" + ) + optional_lists[col].append(ot) + + out_dict: dict[str, torch.Tensor] = { + col: torch.cat(ts, dim=0) for col, ts in cols.items() + } + for col, ts in optional_lists.items(): + if have_optional[col] and ts: + out_dict[col] = torch.cat(ts, dim=0) + n_total = int(out_dict[required[0]].shape[0]) + return TensorDict(out_dict, batch_size=[n_total]) + + # --------------------------------------------------------------------------- # Collate # --------------------------------------------------------------------------- @@ -109,14 +211,33 @@ def collate_molecules( """Collate molecule samples into a nested TensorDict. Each sample is a plain dict with at least ``Z`` and ``pos`` keys. - Optional: ``edge_index``, ``edge_diff``, ``edge_dist``, ``targets``. + Optional: ``edge_index``, ``edge_diff``, ``edge_dist``, ``targets``, + flat ``bond_index`` / ``bond_types``, and nested valence blocks + ``angles`` / ``propers`` / ``impropers`` with 1-D atom-index columns + (``atomi`` / ``atomj`` / …, optional ``type``). + + Pre-collate valence form (preferred):: + + { + "Z": ..., "pos": ..., + "angles": {"atomi": Long[N_a], "atomj": Long[N_a], + "atomk": Long[N_a], "type"?: Long[N_a]}, + "propers": {"atomi": ..., "atomj": ..., "atomk": ..., "atoml": ...}, + "impropers": {"atomi": ..., ...}, # atomi = center (molrs) + } + + Post-collate the same namespaces are nested TensorDicts with columns + rebased by cumulative atom offset. Packed COO ``angle_index [3, N]`` is + **not** the collate schema (see :mod:`molix.datasets._valence_columns` + for optional kernel-local stacks). Args: samples: List of single-molecule sample dicts. target_schema: Declares which targets are graph-level vs atom-level. Returns: - Nested ``TensorDict`` with ``atoms``, ``edges``, ``graphs`` namespaces. + Nested ``TensorDict`` with ``atoms``, ``edges``, ``graphs`` namespaces + and optional ``bonds`` / ``angles`` / ``propers`` / ``impropers``. """ if not samples: raise ValueError("Cannot collate an empty sample list") @@ -125,6 +246,7 @@ def collate_molecules( pos_all: list[torch.Tensor] = [] batch_all: list[torch.Tensor] = [] num_atoms: list[int] = [] + atom_offsets: list[int] = [] edge_all: list[torch.Tensor] = [] diff_all: list[torch.Tensor] = [] @@ -150,6 +272,7 @@ def collate_molecules( pos_all.append(pos) batch_all.append(torch.full((n_atoms,), graph_idx, dtype=torch.long, device=z.device)) num_atoms.append(n_atoms) + atom_offsets.append(atom_offset) if "edge_index" in sample and sample["edge_index"] is not None: edge_index = _normalize_edge_index(sample["edge_index"]) @@ -236,6 +359,14 @@ def collate_molecules( bonds_dict["bond_types"] = torch.cat(btype_all, dim=0) out["bonds"] = TensorDict(bonds_dict, batch_size=[]) + # --- Optional valence column namespaces (angles / propers / impropers) --- + # Primary schema is 1-D columns under nested TensorDict, not packed COO. + # Prefer batch_size=[N_terms] when every leaf shares length N_terms. + for family, required in _VALENCE_REQUIRED.items(): + td = _collate_valence_family(samples, family, required, atom_offsets) + if td is not None: + out[family] = td + return out @@ -246,22 +377,34 @@ def collate_molecules( _TARGET_PREFIX = "targets." -def _gather_indices(ptr: torch.Tensor, idx: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - """Row gather-index and per-sample counts for slicing a packed bucket. +def _gather_indices( + ptr: torch.Tensor, idx: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Row gather-index, counts, segment ids and offsets for a packed bucket. Given a cumsum pointer ``ptr`` ``(n_samples + 1,)`` and selected sample - indices ``idx`` ``(B,)``, returns ``(gather, counts)`` where ``counts`` - ``(B,)`` is each selected sample's element count and ``gather`` - ``(sum(counts),)`` indexes the packed concat tensor in sample-major - order — so ``packed[gather]`` equals concatenating per-sample slices. + indices ``idx`` ``(B,)``, returns ``(gather, counts, seg, offsets)``: + + * ``counts`` ``(B,)`` — each selected sample's element count. + * ``gather`` ``(sum(counts),)`` — indexes the packed concat tensor in + sample-major order, so ``packed[gather]`` equals concatenating the + per-sample slices. + * ``seg`` ``(sum(counts),)`` — owning sample position per gathered row + (i.e. the batch vector for the atom bucket). + * ``offsets`` ``(B,)`` — exclusive-cumsum start of each sample in the + gathered output, used to rebase local atom indices. + + ``seg`` and ``offsets`` fall out of building ``gather`` and every caller + needs at least one of them, so they are returned rather than recomputed + (a second ``repeat_interleave`` per bucket, per batch, per worker). """ counts = ptr[idx + 1] - ptr[idx] starts = ptr[idx] total = int(counts.sum().item()) seg = torch.repeat_interleave(torch.arange(idx.numel()), counts) - new_offsets = torch.cumsum(counts, 0) - counts - gather = starts[seg] + (torch.arange(total) - new_offsets[seg]) - return gather, counts + offsets = torch.cumsum(counts, 0) - counts + gather = starts[seg] + (torch.arange(total) - offsets[seg]) + return gather, counts, seg, offsets def collate_packed( @@ -320,9 +463,7 @@ def collate_packed( ) atom_ptr = payload["atom_ptr"] - a_gather, counts = _gather_indices(atom_ptr, idx) - new_atom_offsets = torch.cumsum(counts, 0) - counts - seg = torch.repeat_interleave(torch.arange(n_graphs), counts) + a_gather, counts, seg, new_atom_offsets = _gather_indices(atom_ptr, idx) # --- atom level --- atoms_dict: dict[str, torch.Tensor] = { @@ -358,8 +499,7 @@ def _route_target(key: str, value: torch.Tensor) -> None: # --- edge level --- if "edge_index" in edges_bucket: edge_ptr = payload["edge_ptr"] - e_gather, e_counts = _gather_indices(edge_ptr, idx) - e_seg = torch.repeat_interleave(torch.arange(n_graphs), e_counts) + e_gather, _e_counts, e_seg, _ = _gather_indices(edge_ptr, idx) edge_index = edges_bucket["edge_index"][e_gather].long() edge_index = rebase(edge_index, new_atom_offsets[e_seg], "edge_index") edges_dict: dict[str, torch.Tensor] = {"edge_index": edge_index} @@ -399,8 +539,7 @@ def _route_target(key: str, value: torch.Tensor) -> None: # --- covalent-bond level (mirror of collate_molecules' bonds namespace) --- bonds_bucket: Mapping[str, torch.Tensor] = payload.get("bonds", {}) if "bond_index" in bonds_bucket and payload.get("bond_ptr") is not None: - b_gather, b_counts = _gather_indices(payload["bond_ptr"], idx) - b_seg = torch.repeat_interleave(torch.arange(n_graphs), b_counts) + b_gather, _b_counts, b_seg, _ = _gather_indices(payload["bond_ptr"], idx) # bond_index is COO [2, N]: gather columns, offset both rows by the # owning sample's atom base (registry "bond_index", index_axis 0). bond_index = rebase( @@ -411,4 +550,35 @@ def _route_target(key: str, value: torch.Tensor) -> None: bonds_dict["bond_types"] = bonds_bucket["bond_types"][b_gather] out["bonds"] = TensorDict(bonds_dict, batch_size=[]) + # --- valence column namespaces (mirror of collate_molecules) --- + # Packed payload stores concatenated 1-D columns + angle_ptr / proper_ptr / + # improper_ptr. Gather rows, rebase atom-index columns by segment atom base. + _FAMILY_PTR = { + "angles": "angle_ptr", + "propers": "proper_ptr", + "impropers": "improper_ptr", + } + for family, required in _VALENCE_REQUIRED.items(): + ptr_key = _FAMILY_PTR[family] + bucket: Mapping[str, torch.Tensor] = payload.get(family, {}) + ptr = payload.get(ptr_key) + if not bucket or ptr is None: + continue + v_gather, _v_counts, v_seg, _ = _gather_indices(ptr, idx) + seg_offsets = new_atom_offsets[v_seg] + fam_dict: dict[str, torch.Tensor] = {} + for col, tensor in bucket.items(): + gathered = tensor[v_gather].long() + if col in _ATOM_INDEX_COLS: + gathered = gathered + seg_offsets + fam_dict[col] = gathered + # Ensure required columns are present (defensive for corrupt caches). + for col in required: + if col not in fam_dict: + raise ValueError( + f"packed cache {family!r} bucket missing required column {col!r}" + ) + n_terms = int(fam_dict[required[0]].shape[0]) + out[family] = TensorDict(fam_dict, batch_size=[n_terms]) + return out diff --git a/src/molix/data/dataset.py b/src/molix/data/dataset.py index 721d9e0..b8eecae 100644 --- a/src/molix/data/dataset.py +++ b/src/molix/data/dataset.py @@ -110,7 +110,7 @@ class BaseDataset(Dataset[Any], ABC): def __len__(self) -> int: ... @abstractmethod - def __getitem__(self, idx: int) -> dict: # type: ignore[override] + def __getitem__(self, idx: int) -> dict: """Return the ``idx``-th sample as a flat ``dict`` (raw-sample shape).""" ... @@ -166,6 +166,35 @@ def split( ) +def _reject_misfiled_edges(payload: Mapping[str, Any], sink: Path | str) -> None: + """Raise if *payload* has no ``edge_ptr`` because its edges were misfiled. + + A packed cache written before edge keys were routed by identity classified + ``edge_index`` as a per-atom key whenever ``n_edges == n_atoms`` held for + every sample, which also suppressed ``edge_ptr``. Its edges are then + unreadable: ``collate_packed`` emits an empty edges namespace and + ``avg_num_neighbors`` would report ``0.0``. A cache with no edge keys at + all is a different, legitimate state and is left alone. + + Args: + payload: The packed-cache payload mapping. + sink: Path of the backing cache file, for the error message. + + Raises: + ValueError: ``edge_index`` is packed under the ``atom`` schema kind. + """ + spec = payload.get("schema", {}).get("edge_index") + if spec is None or spec[0] != "atom": + return + raise ValueError( + f"cache at {sink} packs 'edge_index' as a per-atom key and has no " + "'edge_ptr' pointer — it was written before per-edge keys were routed " + "by identity, so its edges are unreadable (this happens when " + "n_edges == n_atoms for every sample). Delete the stale cache and " + "re-run the pipeline." + ) + + # --------------------------------------------------------------------------- # Cache-backed datasets # --------------------------------------------------------------------------- @@ -190,7 +219,7 @@ def __init__(self, sink: str | Path | PackedCache, *, mmap: bool) -> None: def __len__(self) -> int: return self._n_samples - def __getitem__(self, idx: int) -> dict: # type: ignore[override] + def __getitem__(self, idx: int) -> dict: return PackedCache.unpack_sample(self._payload, idx) def packed_view(self) -> PackedView: @@ -243,11 +272,24 @@ def avg_num_neighbors(self) -> float: ``NeighborList(symmetry=True)`` this equals the mean number of neighbours per atom (Allegro/MACE normalisation constant). - Returns ``0.0`` if the cache has no edge or atom pointers. + Returns ``0.0`` when the cache genuinely holds no edges (no + :class:`~molix.data.tasks.NeighborList` in the pipeline) or no + atoms — zero neighbours per atom is the honest answer there, and + profiling an edge-free cache must stay non-fatal. + + Raises: + ValueError: The cache packs ``edge_index`` as a per-atom key + and therefore carries no ``edge_ptr`` — a stale cache + written before edge keys were routed by identity. Its + edges are unreadable, so ``0.0`` would hand Allegro/MACE a + silently wrong normalisation constant. """ atom_ptr = self._payload.get("atom_ptr") edge_ptr = self._payload.get("edge_ptr") - if atom_ptr is None or edge_ptr is None: + if edge_ptr is None: + _reject_misfiled_edges(self._payload, self.sink) + return 0.0 + if atom_ptr is None: return 0.0 total_atoms = int(atom_ptr[-1].item()) if total_atoms <= 0: @@ -392,7 +434,7 @@ def __getattr__(self, name: str) -> Any: def __len__(self) -> int: return len(self._indices) - def __getitem__(self, idx: int) -> dict: # type: ignore[override] + def __getitem__(self, idx: int) -> dict: """Return the sample at the ``idx``-th index of this subset's view. Maps the local index through ``self._indices`` and defers to the @@ -422,7 +464,10 @@ def avg_num_neighbors(self) -> float: return getattr(self._dataset, "avg_num_neighbors", 0.0) atom_ptr = payload.get("atom_ptr") edge_ptr = payload.get("edge_ptr") - if atom_ptr is None or edge_ptr is None: + if edge_ptr is None: + _reject_misfiled_edges(payload, getattr(self._dataset, "sink", "")) + return 0.0 + if atom_ptr is None: return 0.0 idx = torch.as_tensor(self._indices, dtype=torch.long) n_atoms = (atom_ptr[idx + 1] - atom_ptr[idx]).sum() diff --git a/src/molix/data/group_split.py b/src/molix/data/group_split.py new file mode 100644 index 0000000..ef0fa21 --- /dev/null +++ b/src/molix/data/group_split.py @@ -0,0 +1,65 @@ +"""Molecule-level index split helpers for multi-conformer energy training.""" + +from __future__ import annotations + +import random +from collections.abc import Sequence + +__all__ = ["group_split_indices"] + + +def group_split_indices( + group_ids: Sequence[str | int], + *, + ratios: tuple[float, float, float] = (0.8, 0.1, 0.1), + seed: int = 0, +) -> tuple[list[int], list[int], list[int]]: + """Split sample indices by unique *group_ids* (no molecule leakage). + + Args: + group_ids: Per-sample molecule / group identity (length N). + ratios: ``(train, val, test)`` fractions summing to 1. + seed: RNG seed for shuffling unique groups. + + Returns: + ``(train_idx, val_idx, test_idx)`` lists of parent indices. + """ + if not group_ids: + raise ValueError("group_ids must be non-empty") + train_r, val_r, test_r = ratios + if abs(train_r + val_r + test_r - 1.0) > 1e-9: + raise ValueError(f"ratios must sum to 1.0, got {ratios}") + + by_g: dict[str, list[int]] = {} + order: list[str] = [] + for i, g in enumerate(group_ids): + key = str(g) + if key not in by_g: + by_g[key] = [] + order.append(key) + by_g[key].append(i) + + groups = list(order) + rng = random.Random(seed) + rng.shuffle(groups) + n = len(groups) + n_train = int(n * train_r) + n_val = int(n * val_r) + train_g = groups[:n_train] + val_g = groups[n_train : n_train + n_val] + test_g = groups[n_train + n_val :] + + def expand(gs: list[str]) -> list[int]: + out: list[int] = [] + for g in gs: + out.extend(by_g[g]) + return out + + train_idx, val_idx, test_idx = expand(train_g), expand(val_g), expand(test_g) + # Leakage guard + train_set = {group_ids[i] for i in train_idx} + val_set = {group_ids[i] for i in val_idx} + test_set = {group_ids[i] for i in test_idx} + if train_set & val_set or train_set & test_set or val_set & test_set: + raise RuntimeError("group_split_indices produced molecule leakage") + return train_idx, val_idx, test_idx diff --git a/src/molix/datasets/__init__.py b/src/molix/datasets/__init__.py index 1bf39d2..b03bd8a 100644 --- a/src/molix/datasets/__init__.py +++ b/src/molix/datasets/__init__.py @@ -25,7 +25,7 @@ from molix.datasets.molrec import MolRecSource except ImportError: # molpy.MolRec (or molrec module deps) not available - class MolRecSource: # type: ignore[no-redef] + class MolRecSource: """Placeholder when ``molpy.MolRec`` is not on the public API.""" def __init__(self, *args, **kwargs): diff --git a/src/molix/datasets/_extxyz.py b/src/molix/datasets/_extxyz.py index 985a6af..1793bcc 100644 --- a/src/molix/datasets/_extxyz.py +++ b/src/molix/datasets/_extxyz.py @@ -1,36 +1,37 @@ -"""Minimal in-tree extended-XYZ (extxyz) parser — molpy-only stack. +"""Minimal in-tree extended-XYZ (extxyz) parser — numpy only. This module exists because :class:`molpy.io.trajectory.xyz.XYZTrajectoryReader` only reads the canonical XYZ format (``n_atoms`` + comment + ``element x y z`` rows) and discards the comment-line metadata that extxyz files carry — ``Lattice="..."``, ``Properties=...``, ``energy=...``, ``pbc="..."`` — as well -as any per-atom columns beyond ``x y z``. The Sonata bulk-water RPBE-D3 data -ships in extended-XYZ format with per-frame ``cell``, ``energy``, and per-atom -``forces``, so an extxyz-aware parser is required. +as any per-atom columns beyond ``x y z``. Dataset fixtures and numerical +regression frames ship in extended-XYZ with per-frame ``cell`` / ``energy`` and +per-atom ``forces`` (and optionally ``initial_charges``, ``dipoles``, …). This parser is intentionally narrow: -* depends on ``numpy`` + Python stdlib only (NO ``ase``); -* recognises only the column tags this project consumes - (``species:S:1``, ``pos:R:3``, ``forces:R:3``) and skips others; +* depends on ``numpy`` + Python stdlib only (no ASE, no e3nn); +* requires ``species:S:1`` and ``pos:R:3``; +* promotes ``forces:R:3`` to :attr:`ExtxyzFrame.forces` when present; +* stores every other real-valued ``Properties`` column under + :attr:`ExtxyzFrame.arrays` (e.g. ``initial_charges``, ``dipoles``); * returns a flat list of :class:`ExtxyzFrame` dataclasses; callers (e.g. :class:`molix.datasets.water_les.WaterLESSource`) bridge to :class:`torch.Tensor` and the flat-sample dict contract. References: - Extended-XYZ format spec (ASE wiki / Schimka et al. 2017 supp.): - https://github.com/libAtoms/extxyz + Extended-XYZ format: https://github.com/libAtoms/extxyz """ from __future__ import annotations import logging -from dataclasses import dataclass +from dataclasses import dataclass, field from pathlib import Path import numpy as np -__all__ = ["ExtxyzFrame", "parse_extxyz_frames"] +__all__ = ["ExtxyzFrame", "parse_extxyz_frames", "write_extxyz_frames"] _log = logging.getLogger(__name__) @@ -55,6 +56,10 @@ class ExtxyzFrame: forces: ``(n_atoms, 3)`` forces in eV·Å⁻¹, or ``None`` when the ``Properties=`` declaration does not list a ``forces:R:3`` column. + arrays: Other real-valued per-atom columns declared in + ``Properties=``, keyed by name (e.g. ``"initial_charges"`` + ``(n_atoms,)`` or ``"dipoles"`` ``(n_atoms, 3)``). Empty when + only species/pos/forces are present. """ n_atoms: int @@ -64,6 +69,7 @@ class ExtxyzFrame: species: list[str] pos: np.ndarray forces: np.ndarray | None + arrays: dict[str, np.ndarray] = field(default_factory=dict) def parse_extxyz_frames(path: str | Path) -> list[ExtxyzFrame]: @@ -114,6 +120,10 @@ def parse_extxyz_frames(path: str | Path) -> list[ExtxyzFrame]: forces: np.ndarray | None = ( np.empty((n_atoms, 3), dtype=np.float64) if layout["has_forces"] else None ) + extra: dict[str, np.ndarray] = { + name: np.empty((n_atoms, width) if width > 1 else (n_atoms,), dtype=np.float64) + for name, _col, width in layout["real_cols"] + } for j, row in enumerate(atom_lines): parts = row.split() if len(parts) < layout["min_cols"]: @@ -130,6 +140,9 @@ def parse_extxyz_frames(path: str | Path) -> list[ExtxyzFrame]: if forces is not None: fc = layout["forces_col"] forces[j] = (float(parts[fc]), float(parts[fc + 1]), float(parts[fc + 2])) + for name, col, width in layout["real_cols"]: + vals = tuple(float(parts[col + k]) for k in range(width)) + extra[name][j] = vals if width > 1 else vals[0] frames.append( ExtxyzFrame( @@ -140,6 +153,7 @@ def parse_extxyz_frames(path: str | Path) -> list[ExtxyzFrame]: species=species, pos=pos, forces=forces, + arrays=extra, ) ) i += 2 + n_atoms @@ -219,7 +233,7 @@ def _parse_pbc(tokens: dict[str, str]) -> tuple[bool, bool, bool]: parts = raw.split() if len(parts) != 3: raise ValueError(f"pbc tag has {len(parts)} entries; need 3") - return tuple(p.upper() == "T" for p in parts) # type: ignore[return-value] + return tuple(p.upper() == "T" for p in parts) def _parse_energy(tokens: dict[str, str], *, source: Path, frame_idx: int) -> float: @@ -230,15 +244,16 @@ def _parse_energy(tokens: dict[str, str], *, source: Path, frame_idx: int) -> fl return float(raw) -def _parse_properties(tokens: dict[str, str], *, source: Path) -> dict[str, int | bool]: - """Parse ``Properties=species:S:1:pos:R:3:forces:R:3`` into column offsets. +def _parse_properties(tokens: dict[str, str], *, source: Path) -> dict: + """Parse ``Properties=species:S:1:pos:R:3:…`` into column offsets. Returns a dict with:: species_col: column index of the symbol column pos_col: column index where the 3 position columns begin - forces_col: column index where the 3 force columns begin (only if has_forces) + forces_col: column index of forces (only when has_forces) has_forces: whether forces are declared + real_cols: list of ``(name, col, width)`` for other ``R`` columns min_cols: minimum number of columns each atom row must have """ raw = tokens.get("Properties", "species:S:1:pos:R:3") @@ -249,16 +264,19 @@ def _parse_properties(tokens: dict[str, str], *, source: Path) -> dict[str, int species_col: int | None = None pos_col: int | None = None forces_col: int | None = None + real_cols: list[tuple[str, int, int]] = [] cursor = 0 for k in range(0, len(parts), 3): - name, _type, width_str = parts[k], parts[k + 1], parts[k + 2] + name, typ, width_str = parts[k], parts[k + 1], parts[k + 2] width = int(width_str) if name == "species": species_col = cursor elif name == "pos": pos_col = cursor - elif name == "forces": + elif name == "forces" and typ == "R": forces_col = cursor + elif typ == "R": + real_cols.append((name, cursor, width)) cursor += width if species_col is None or pos_col is None: @@ -266,12 +284,52 @@ def _parse_properties(tokens: dict[str, str], *, source: Path) -> dict[str, int f"{source}: Properties tag must declare both species and pos columns; got {raw!r}" ) - layout: dict[str, int | bool] = { + layout: dict = { "species_col": species_col, "pos_col": pos_col, "has_forces": forces_col is not None, + "real_cols": real_cols, "min_cols": cursor, } if forces_col is not None: layout["forces_col"] = forces_col return layout + + +def write_extxyz_frames( + path: str | Path, + *, + species: list[str], + positions: np.ndarray, + energies: np.ndarray | None = None, + tags: list[str] | None = None, +) -> None: + """Write frames in extended-XYZ — the write half of :func:`parse_extxyz_frames`. + + Emits ``Properties=species:S:1:pos:R:3`` plus an ``energy=`` token per + frame, so the output is readable by this module's own parser (open + systems: no ``Lattice`` tag is written). + + Args: + path: Output ``.xyz`` file (parent directories must exist). + species: Atomic symbols, length ``N`` (constant across frames). + positions: Positions ``(T, N, 3)`` in Å. + energies: Optional per-frame total energy ``(T,)``; written as the + ``energy=`` comment token when given. + tags: Optional per-frame extra comment tokens (length ``T``), appended + verbatim — e.g. ``"temperature=297.1"``. + """ + pos = np.asarray(positions, dtype=np.float64) + n_frames, n_atoms = pos.shape[0], pos.shape[1] + if len(species) != n_atoms: + raise ValueError(f"{len(species)} species for {n_atoms} atoms") + with Path(path).open("w") as fh: + for t in range(n_frames): + comment = "Properties=species:S:1:pos:R:3" + if energies is not None: + comment += f" energy={float(energies[t]):.10f}" + if tags is not None: + comment += f" {tags[t]}" + fh.write(f"{n_atoms}\n{comment}\n") + for symbol, (x, y, z) in zip(species, pos[t]): + fh.write(f"{symbol} {x:.8f} {y:.8f} {z:.8f}\n") diff --git a/src/molix/datasets/_valence_columns.py b/src/molix/datasets/_valence_columns.py new file mode 100644 index 0000000..aff965e --- /dev/null +++ b/src/molix/datasets/_valence_columns.py @@ -0,0 +1,84 @@ +"""Kernel-local stack helpers: valence column TensorDicts → COO index tables. + +The post-collate batch schema is **column form** under nested namespaces +(``batch["angles"]["atomi"]`` etc.). Some Class-I potential kernels still take +packed COO ``[arity, N]`` indices. These helpers stack columns only at the +call site — they are **not** part of the collate / cache schema. + +See ``.claude/notes/learnable-classical-ff.md`` and spec +``learnable-classical-ff-02-valence-topology``. + +Impropers follow molrs center-first: ``atomi`` is the center. +""" + +from __future__ import annotations + +from collections.abc import Mapping + +import torch + + +def _stack_columns( + block: Mapping[str, torch.Tensor], + cols: tuple[str, ...], +) -> torch.Tensor: + """Stack 1-D long columns into COO ``[len(cols), N]``. + + Args: + block: Mapping / TensorDict with the named 1-D columns. + cols: Column names in row order. + + Returns: + Long tensor of shape ``[len(cols), N]``. + + Raises: + KeyError: a required column is missing. + ValueError: column lengths differ. + """ + tensors = [torch.as_tensor(block[c]).long().reshape(-1) for c in cols] + n = int(tensors[0].shape[0]) + for name, t in zip(cols, tensors): + if int(t.shape[0]) != n: + raise ValueError( + f"valence columns must share length; {cols[0]}={n}, {name}={t.shape[0]}" + ) + return torch.stack(tensors, dim=0) + + +def stack_angle_index(angles_td: Mapping[str, torch.Tensor]) -> torch.Tensor: + """Stack ``atomi`` / ``atomj`` / ``atomk`` into COO ``[3, N]``. + + Args: + angles_td: Angles namespace (``atomi``, ``atomj`` central, ``atomk``). + + Returns: + Long tensor ``[3, N_angles]`` for kernel call sites. + """ + return _stack_columns(angles_td, ("atomi", "atomj", "atomk")) + + +def stack_proper_index(propers_td: Mapping[str, torch.Tensor]) -> torch.Tensor: + """Stack ``atomi``..``atoml`` into COO ``[4, N]`` for proper torsions. + + Args: + propers_td: Propers namespace with four atom-index columns. + + Returns: + Long tensor ``[4, N_propers]``. + """ + return _stack_columns(propers_td, ("atomi", "atomj", "atomk", "atoml")) + + +def stack_improper_index(impropers_td: Mapping[str, torch.Tensor]) -> torch.Tensor: + """Stack improper columns into COO ``[4, N]`` (molrs center-first). + + Row 0 is ``atomi`` = **center**. OpenFF trefoil reorder is an export + adapter, not a second internal convention. + + Args: + impropers_td: Impropers namespace with four atom-index columns. + + Returns: + Long tensor ``[4, N_impropers]`` with center at row 0. + """ + return _stack_columns(impropers_td, ("atomi", "atomj", "atomk", "atoml")) diff --git a/src/molix/datasets/water_les.py b/src/molix/datasets/water_les.py index e3dbd68..12f6ee3 100644 --- a/src/molix/datasets/water_les.py +++ b/src/molix/datasets/water_les.py @@ -74,7 +74,7 @@ class WaterLESSource: ``"val"`` slice the same upstream ``train-…`` file via the deterministic tail-slice with :attr:`TRAIN_VAL_RATIO`. download: If ``True``, fetch any missing file from - :attr:`BASE_URL` via :mod:`urllib.request` (no ASE). + :attr:`BASE_URL` via :mod:`urllib.request`. verify_checksum: If ``True``, compute SHA-256 of every consumed file and compare against :attr:`_CHECKSUMS`. Default ``False`` so the placeholder digests in ``_data_acquisition.md`` do not diff --git a/src/molix/engine/adapter.py b/src/molix/engine/adapter.py index 51f03c7..9bc614f 100644 --- a/src/molix/engine/adapter.py +++ b/src/molix/engine/adapter.py @@ -92,7 +92,7 @@ class MolnexTensorDictAdapter(EngineAdapter): """Adapter for molnex-native potentials (PiNet, MACE, … via ``PiNetPotential``-style). Builds the post-collate nested ``TensorDict`` (``atoms`` / ``edges`` / ``graphs`` - per ``CLAUDE.md``), runs ``model(batch, compute_forces=True)``, and reads the + per ``CLAUDE.md``), runs ``model(batch)`` (built with ``compute_forces=True``), reads the ``{"energy", "forces"}`` output dict. Energy is summed to a scalar (single graph), matching what the C++ pair style accumulates into ``eng_vdwl``. """ @@ -153,7 +153,7 @@ def build_inputs( def read_outputs(self, out: object) -> tuple[torch.Tensor, torch.Tensor]: if isinstance(out, dict): return out["energy"].sum(), out["forces"] - energy, forces = out # type: ignore[misc,not-iterable] + energy, forces = out return energy.sum(), forces @@ -179,6 +179,8 @@ def forward( if isinstance(inputs, tuple): out = self.model(*inputs) else: - out = self.model(inputs, compute_forces=True) + # Monomorphic forward since b85d12f: force derivation is a + # construction-time property of the potential, not a call kwarg. + out = self.model(inputs) energy, forces = self.adapter.read_outputs(out) return energy.reshape(()), forces diff --git a/src/molix/engine/static.py b/src/molix/engine/static.py index 77608c1..ff65ffd 100644 --- a/src/molix/engine/static.py +++ b/src/molix/engine/static.py @@ -7,7 +7,7 @@ **inert without changing weights or physics**: their ``edge_diff`` is overwritten to a length past the cutoff, so the model's ``cutoff(edge_dist)`` zeros their energy and force contribution (PiNet consumes the provided ``edge_diff`` via -``_edge_bond_diff`` and recomputes ``edge_dist`` from it). +``edge_bond_diff`` and recomputes ``edge_dist`` from it). Export this with ``dynamic_shapes=None`` and load the ``.pt2`` with ``run_single_threaded=True`` (PyTorch #158834, fixed in torch 2.8) — then the @@ -21,6 +21,9 @@ import torch.nn as nn from tensordict import TensorDict +from molix.schema import ENERGY_KEY, FORCES_KEY +from molix.units import DEAD_EDGE_CUTOFF_FACTOR + class StaticForward(nn.Module): """Fixed ``(N, E_max)`` flat forward: ``(Z, pos, edge_index, mask) -> (energy, forces)``. @@ -33,7 +36,8 @@ class StaticForward(nn.Module): edges (real ones first, ``mask=True``; the rest padded, ``mask=False``) and fail loudly when the real edge count exceeds ``e_max``. cutoff: Neighbour cutoff in Å; padded edges get ``edge_diff`` of length - ``10*cutoff`` so ``cutoff(edge_dist)`` zeros them. + ``DEAD_EDGE_CUTOFF_FACTOR * cutoff`` (see :mod:`molix.units`) so + ``cutoff(edge_dist)`` zeros them. """ def __init__(self, model: nn.Module, n_atoms: int, e_max: int, cutoff: float) -> None: @@ -41,7 +45,7 @@ def __init__(self, model: nn.Module, n_atoms: int, e_max: int, cutoff: float) -> self.model = model self.n_atoms = int(n_atoms) self.e_max = int(e_max) - self.pad_len = float(cutoff) * 10.0 + self.pad_len = float(cutoff) * DEAD_EDGE_CUTOFF_FACTOR def forward( self, Z: torch.Tensor, pos: torch.Tensor, edge_index: torch.Tensor, mask: torch.Tensor @@ -75,5 +79,7 @@ def forward( ), batch_size=[], ) - out = self.model(td, compute_forces=True) - return out["energy"].reshape(()), out["forces"] + # Monomorphic forward since b85d12f: a potential that derives forces was + # constructed that way (``PiNetPotential(compute_forces=True)``). + out = self.model(td) + return out[ENERGY_KEY].reshape(()), out[FORCES_KEY] diff --git a/src/molix/ff_export/__init__.py b/src/molix/ff_export/__init__.py new file mode 100644 index 0000000..b6b59d8 --- /dev/null +++ b/src/molix/ff_export/__init__.py @@ -0,0 +1,56 @@ +"""Potential IR → backend force-spec export (OpenMM-first). + +Compile Class-I :class:`~molpot.ir.PotentialIR` parameter bags into a +backend-neutral :class:`ForceSpec`, with peer :class:`BackendAdapter` types +for unit/form translation. No live OpenMM import is required for tests — +goldens are hard-coded numbers and JSON force-spec dicts. + +Load-bearing conversion goldens +------------------------------- +* Bond ``k = 100`` kcal mol⁻¹ Å⁻² → ``41840`` kJ mol⁻¹ nm⁻² +* AMBER torsion ``Vn = 2`` kcal/mol → OpenMM PeriodicTorsion ``k = 4.184`` kJ/mol + +References: + OpenMM User Guide §19 "Forces" + Spec: learnable-classical-ff-08-ff-export +""" + +from molix.ff_export.adapter import BackendAdapter +from molix.ff_export.cases import TranslationCase +from molix.ff_export.compiler import ForceFieldCompiler +from molix.ff_export.conventions import ( + ANGSTROM_TO_NM, + BOND_K_IR_TO_OPENMM, + KCAL_PER_MOL_TO_KJ_PER_MOL, + ConventionRow, + ConventionTable, + scale_amber_vn, + scale_angle_k, + scale_bond_k, + scale_energy, + scale_length, + scale_torsion_k, +) +from molix.ff_export.exceptions import UnsupportedTermError +from molix.ff_export.force_spec import ForceSpec +from molix.ff_export.openmm_adapter import OpenMMAdapter + +__all__ = [ + "ANGSTROM_TO_NM", + "BOND_K_IR_TO_OPENMM", + "KCAL_PER_MOL_TO_KJ_PER_MOL", + "BackendAdapter", + "ConventionRow", + "ConventionTable", + "ForceFieldCompiler", + "ForceSpec", + "OpenMMAdapter", + "TranslationCase", + "UnsupportedTermError", + "scale_amber_vn", + "scale_angle_k", + "scale_bond_k", + "scale_energy", + "scale_length", + "scale_torsion_k", +] diff --git a/src/molix/ff_export/adapter.py b/src/molix/ff_export/adapter.py new file mode 100644 index 0000000..4a161e0 --- /dev/null +++ b/src/molix/ff_export/adapter.py @@ -0,0 +1,73 @@ +"""BackendAdapter protocol + self-registering peer types. + +Mirrors :class:`molix.engine.EngineAdapter`: each backend is a named +peer class (``OpenMMAdapter``, future ``GromacsAdapter``), **not** a +``Translator(method="openmm")`` switch. + +References: + Spec: learnable-classical-ff-08-ff-export + Placement: ``.claude/notes/learnable-classical-ff.md`` (no method=) +""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import Any + +from molix.ff_export.force_spec import ForceSpec +from molpot.ir import PotentialIR + +__all__ = ["BackendAdapter"] + + +class BackendAdapter(ABC): + """Strategy that translates :class:`~molpot.ir.PotentialIR` → :class:`ForceSpec`. + + Concrete subclasses set a class-level :attr:`name` (auto-registered) and + implement :meth:`translate`. Resolve by name via :meth:`from_name`. + """ + + name: str = "" + _registry: dict[str, type[BackendAdapter]] = {} + + def __init_subclass__(cls, **kwargs: object) -> None: + super().__init_subclass__(**kwargs) + if cls.name: + BackendAdapter._registry[cls.name] = cls + + @classmethod + def from_name(cls, name: str) -> BackendAdapter: + """Instantiate the registered adapter for ``name`` (e.g. ``"openmm"``).""" + try: + return cls._registry[name]() + except KeyError: + valid = ", ".join(sorted(cls._registry)) or "(none)" + raise ValueError(f"unknown adapter {name!r}; valid adapters: {valid}") from None + + @classmethod + def names(cls) -> tuple[str, ...]: + """All registered adapter names.""" + return tuple(sorted(cls._registry)) + + @abstractmethod + def translate( + self, + ir: PotentialIR, + *, + meta: dict[str, Any] | None = None, + ) -> ForceSpec: + """Translate ``ir`` into a backend force specification. + + Args: + ir: Class-I parameter bags (+ optional scaling). + meta: Optional attachments (type systems, symbolic FF, …). + + Returns: + A serializable :class:`ForceSpec`. + + Raises: + UnsupportedTermError: When an IR bag has no faithful mapping. + """ + + def __repr__(self) -> str: + return f"{type(self).__name__}()" diff --git a/src/molix/ff_export/cases.py b/src/molix/ff_export/cases.py new file mode 100644 index 0000000..43be5f4 --- /dev/null +++ b/src/molix/ff_export/cases.py @@ -0,0 +1,33 @@ +"""Translation case matrix for Potential IR → backend force forms. + +Four explicit outcomes — never a silent drop of an IR term. + +References: + OpenMM User Guide §19 "Forces" + Spec: learnable-classical-ff-08-ff-export +""" + +from __future__ import annotations + +from enum import Enum + +__all__ = ["TranslationCase"] + + +class TranslationCase(Enum): + """How an IR interaction bag maps into a backend force form. + + Attributes: + DIRECT_UNIT_SCALE: Same functional form; only unit conversion + (e.g. harmonic bond/angle with matching ``½ k x²``). + FORM_REPARAMETERIZE: Same physics, different parameter convention + (e.g. torsion ``k`` vs ``k/2``, idivf absorption, LJ ε/σ vs A/B). + DECOMPOSE: One IR bag expands to multiple backend force parameters + (e.g. multi-term proper → multiple PeriodicTorsion rows). + UNSUPPORTED: No faithful mapping; raise :class:`UnsupportedTermError`. + """ + + DIRECT_UNIT_SCALE = "direct_unit_scale" + FORM_REPARAMETERIZE = "form_reparameterize" + DECOMPOSE = "decompose" + UNSUPPORTED = "unsupported" diff --git a/src/molix/ff_export/compiler.py b/src/molix/ff_export/compiler.py new file mode 100644 index 0000000..117e62f --- /dev/null +++ b/src/molix/ff_export/compiler.py @@ -0,0 +1,76 @@ +"""ForceFieldCompiler — Potential IR → backend force-spec. + +Single primitive: :meth:`ForceFieldCompiler.compile`. Backend selection is by +peer :class:`BackendAdapter` type (or registered name), never +``compile(method=…)``. + +References: + OpenMM User Guide §19 "Forces" + Spec: learnable-classical-ff-08-ff-export +""" + +from __future__ import annotations + +from typing import Any + +from molix.ff_export.adapter import BackendAdapter +from molix.ff_export.force_spec import ForceSpec +from molix.ff_export.openmm_adapter import OpenMMAdapter as _OpenMMAdapter # noqa: F401 +from molpot.ir import PotentialIR + +__all__ = ["ForceFieldCompiler"] + + +class ForceFieldCompiler: + """Compile Class-I :class:`~molpot.ir.PotentialIR` into a :class:`ForceSpec`. + + Args: + adapter: A :class:`BackendAdapter` instance, or a registered name + (default ``"openmm"``). + + Example: + >>> from molix.ff_export import ForceFieldCompiler + >>> from molpot.ir import BondBag, PotentialIR + >>> import torch + >>> bonds = BondBag(k=torch.tensor([100.0]), r0=torch.tensor([1.5])) + >>> spec = ForceFieldCompiler("openmm").compile(PotentialIR(bonds=bonds)) + >>> spec.forces[0]["parameters"][0]["k"] + 41840.0 + """ + + def __init__(self, adapter: BackendAdapter | str = "openmm") -> None: + if isinstance(adapter, str): + self.adapter: BackendAdapter = BackendAdapter.from_name(adapter) + elif isinstance(adapter, BackendAdapter): + self.adapter = adapter + else: + raise TypeError(f"adapter must be BackendAdapter or str name, got {type(adapter)!r}") + + def compile( + self, + ir: PotentialIR, + *, + type_systems: Any = None, + symbolic: Any = None, + ) -> ForceSpec: + """Translate ``ir`` through the configured backend adapter. + + Args: + ir: Class-I potential intermediate representation. + type_systems: Optional type-system metadata (passed through to + force-spec ``metadata``). + symbolic: Optional symbolic force-field metadata (passed through). + + Returns: + Backend force specification (OpenMM units when using + :class:`~molix.ff_export.OpenMMAdapter`). + + Raises: + UnsupportedTermError: When the adapter cannot map an IR bag. + """ + meta: dict[str, Any] = {} + if type_systems is not None: + meta["type_systems"] = type_systems + if symbolic is not None: + meta["symbolic"] = symbolic + return self.adapter.translate(ir, meta=meta or None) diff --git a/src/molix/ff_export/conventions.py b/src/molix/ff_export/conventions.py new file mode 100644 index 0000000..6000c4c --- /dev/null +++ b/src/molix/ff_export/conventions.py @@ -0,0 +1,225 @@ +"""Unit and form convention table for Class-I IR → OpenMM force-spec. + +IR units (learnable-classical-ff-01): **kcal/mol, Å, e, rad**. +OpenMM standard MD units: **kJ/mol, nm, e, rad**. + +Load-bearing goldens (regression + unit tests): + +1. Bond force constant for ``E = ½ k (r − r₀)²``:: + + k = 100 kcal mol⁻¹ Å⁻² + → k_omm = 100 × 4.184 / (0.1 nm)² = 41840 kJ mol⁻¹ nm⁻² + +2. AMBER proper torsion barrier ``Vn = 2`` kcal/mol maps to OpenMM + :class:`PeriodicTorsionForce` amplitude ``k = 4.184`` kJ/mol because + AMBER uses ``E = (Vn/2)[1 + cos(...)]`` while OpenMM (and the molnex IR + form ``E = (k/s)[1 + cos(...)]``) absorb the half-barrier into ``k``. + Thus ``k_ir = Vn/2 = 1`` and ``k_omm = 1 × 4.184``. + +OpenMM PeriodicTorsionForce (User Guide §19) energy:: + + E = k [1 + cos(n φ − γ)] + +with ``k`` already the half-barrier coefficient in energy units. The IR +``ProperTorsionBag`` stores the same form (``k / idivf`` is the OpenMM ``k`` +in kcal/mol); export multiplies by 4.184 and divides by ``idivf``. + +References: + OpenMM User Guide §19 "Forces" + SMIRNOFF unit conventions + Spec: learnable-classical-ff-08-ff-export +""" + +from __future__ import annotations + +from dataclasses import dataclass +from types import MappingProxyType +from typing import Mapping + +from molix.ff_export.cases import TranslationCase + +__all__ = [ + "KCAL_PER_MOL_TO_KJ_PER_MOL", + "ANGSTROM_TO_NM", + "BOND_K_IR_TO_OPENMM", + "ANGLE_K_IR_TO_OPENMM", + "ConventionRow", + "ConventionTable", + "scale_energy", + "scale_length", + "scale_bond_k", + "scale_angle_k", + "scale_torsion_k", + "scale_amber_vn", +] + +# Named conversion factors (exact rationals where possible). +KCAL_PER_MOL_TO_KJ_PER_MOL: float = 4.184 +ANGSTROM_TO_NM: float = 0.1 +# k_omm = k_ir * energy / length². Literal 418.4 so 100 → 41840 is exact in float. +BOND_K_IR_TO_OPENMM: float = 418.4 +# Angle k: θ in rad both sides → energy factor only. +ANGLE_K_IR_TO_OPENMM: float = KCAL_PER_MOL_TO_KJ_PER_MOL + + +def scale_energy(e_kcal: float) -> float: + """Convert energy from kcal/mol to kJ/mol.""" + return e_kcal * KCAL_PER_MOL_TO_KJ_PER_MOL + + +def scale_length(x_angstrom: float) -> float: + """Convert length from Å to nm.""" + return x_angstrom * ANGSTROM_TO_NM + + +def scale_bond_k(k_kcal_per_A2: float) -> float: + """Convert harmonic bond ``k`` from kcal mol⁻¹ Å⁻² to kJ mol⁻¹ nm⁻². + + Golden: ``scale_bond_k(100) == 41840``. + + Uses the named factor 418.4 (= 4.184 / 0.01) so the golden is exact in + IEEE-754 float64 (``100 * (4.184 / 0.01)`` is not bit-exact). + """ + return k_kcal_per_A2 * BOND_K_IR_TO_OPENMM + + +def scale_angle_k(k_kcal_per_rad2: float) -> float: + """Convert harmonic angle ``k`` from kcal mol⁻¹ rad⁻² to kJ mol⁻¹ rad⁻².""" + return k_kcal_per_rad2 * ANGLE_K_IR_TO_OPENMM + + +def scale_torsion_k(k_kcal: float) -> float: + """Convert IR torsion amplitude ``k`` (kcal/mol) to OpenMM ``k`` (kJ/mol). + + IR and OpenMM share ``E = k [1 + cos(nφ − γ)]`` (after idivf absorption); + only the energy unit changes. + """ + return k_kcal * KCAL_PER_MOL_TO_KJ_PER_MOL + + +def scale_amber_vn(vn_kcal: float) -> float: + """Map AMBER full barrier ``Vn`` (kcal/mol) to OpenMM PeriodicTorsion ``k``. + + AMBER: ``E = (Vn/2)[1 + cos(...)]``. OpenMM: ``E = k[1 + cos(...)]``. + Golden: ``scale_amber_vn(2) == 4.184``. + """ + return scale_torsion_k(vn_kcal / 2.0) + + +@dataclass(frozen=True) +class ConventionRow: + """One IR term's translation policy for a backend. + + Attributes: + ir_term: Bag / term name (e.g. ``"bond_harmonic"``). + case: :class:`TranslationCase` outcome. + unit_factors: Named scale factors applied to IR fields. + notes: Human-readable mapping notes (OpenMM form, caveats). + """ + + ir_term: str + case: TranslationCase + unit_factors: Mapping[str, float] + notes: str + + +class ConventionTable: + """Frozen lookup of :class:`ConventionRow` by IR term name. + + Build with :meth:`default_openmm` for the Class-I → OpenMM matrix. + """ + + def __init__(self, rows: Mapping[str, ConventionRow]) -> None: + self._rows: Mapping[str, ConventionRow] = MappingProxyType(dict(rows)) + + def __contains__(self, ir_term: object) -> bool: + return ir_term in self._rows + + def __getitem__(self, ir_term: str) -> ConventionRow: + return self._rows[ir_term] + + def row(self, ir_term: str) -> ConventionRow: + """Return the row for ``ir_term`` or raise :class:`KeyError`.""" + return self._rows[ir_term] + + def terms(self) -> tuple[str, ...]: + """All registered IR term names (sorted).""" + return tuple(sorted(self._rows)) + + @classmethod + def default_openmm(cls) -> ConventionTable: + """Class-I IR → OpenMM §19 convention matrix.""" + rows = { + "bond_harmonic": ConventionRow( + ir_term="bond_harmonic", + case=TranslationCase.DIRECT_UNIT_SCALE, + unit_factors=MappingProxyType({"k": BOND_K_IR_TO_OPENMM, "r0": ANGSTROM_TO_NM}), + notes=( + "OpenMM HarmonicBondForce: E = ½ k (r − r0)² with k in " + "kJ/mol/nm², r0 in nm. Matches IR ½ k form; scale k by " + "4.184/0.01=418.4 (golden: 100 → 41840)." + ), + ), + "angle_harmonic": ConventionRow( + ir_term="angle_harmonic", + case=TranslationCase.DIRECT_UNIT_SCALE, + unit_factors=MappingProxyType({"k": ANGLE_K_IR_TO_OPENMM, "theta0": 1.0}), + notes=( + "OpenMM HarmonicAngleForce: E = ½ k (θ − θ0)²; θ in rad " + "both sides; only energy unit on k changes (×4.184)." + ), + ), + "proper_periodic": ConventionRow( + ir_term="proper_periodic", + case=TranslationCase.FORM_REPARAMETERIZE, + unit_factors=MappingProxyType({"k": KCAL_PER_MOL_TO_KJ_PER_MOL}), + notes=( + "OpenMM PeriodicTorsionForce: E = k[1+cos(nφ−γ)]. " + "IR stores E = (k/idivf)[1+cos]; export k_omm = " + "(k/idivf)×4.184 (idivf absorption = form reparameterize). " + "Multi-term bags use DECOMPOSE at emit time. " + "AMBER Vn full barrier: k_omm = scale_amber_vn(Vn); " + "golden Vn=2 → 4.184 kJ/mol." + ), + ), + "lj": ConventionRow( + ir_term="lj", + case=TranslationCase.DIRECT_UNIT_SCALE, + unit_factors=MappingProxyType( + { + "epsilon": KCAL_PER_MOL_TO_KJ_PER_MOL, + "sigma": ANGSTROM_TO_NM, + } + ), + notes=( + "OpenMM NonbondedForce LJ: ε in kJ/mol, σ in nm " + "(σ convention; combining rules left to consumer)." + ), + ), + "charge": ConventionRow( + ir_term="charge", + case=TranslationCase.DIRECT_UNIT_SCALE, + unit_factors=MappingProxyType({"q": 1.0}), + notes="Elementary charge e is identical in IR and OpenMM.", + ), + "improper_harmonic": ConventionRow( + ir_term="improper_harmonic", + case=TranslationCase.UNSUPPORTED, + unit_factors=MappingProxyType({}), + notes=( + "OpenMM has no built-in harmonic improper force matching " + "IR ImproperHarmonicBag; refuse rather than silent drop." + ), + ), + "improper_periodic": ConventionRow( + ir_term="improper_periodic", + case=TranslationCase.FORM_REPARAMETERIZE, + unit_factors=MappingProxyType({"k": KCAL_PER_MOL_TO_KJ_PER_MOL}), + notes=( + "Periodic improper uses the same PeriodicTorsionForce " + "form as propers; particle ordering is adapter-specific " + "(molrs center-first vs OpenFF trefoil)." + ), + ), + } + return cls(rows) diff --git a/src/molix/ff_export/exceptions.py b/src/molix/ff_export/exceptions.py new file mode 100644 index 0000000..9b5b80c --- /dev/null +++ b/src/molix/ff_export/exceptions.py @@ -0,0 +1,29 @@ +"""Structured errors for force-field export.""" + +from __future__ import annotations + +from molix.ff_export.cases import TranslationCase + +__all__ = ["UnsupportedTermError"] + + +class UnsupportedTermError(ValueError): + """Raised when an IR bag has no faithful backend translation. + + Attributes: + term: IR term name (e.g. ``"improper_harmonic"``). + reason: Human-readable explanation. + case: Always :attr:`TranslationCase.UNSUPPORTED`. + """ + + def __init__( + self, + term: str, + reason: str, + *, + case: TranslationCase = TranslationCase.UNSUPPORTED, + ) -> None: + self.term = term + self.reason = reason + self.case = case + super().__init__(f"Unsupported IR term {term!r} ({case.name}): {reason}") diff --git a/src/molix/ff_export/force_spec.py b/src/molix/ff_export/force_spec.py new file mode 100644 index 0000000..502fc1b --- /dev/null +++ b/src/molix/ff_export/force_spec.py @@ -0,0 +1,56 @@ +"""Backend-neutral force specification (JSON-serializable). + +:class:`ForceSpec` is the portable product of a :class:`BackendAdapter` +translation. It is **not** a live OpenMM System — consumers may materialize +one later; unit tests pin hard-coded dict goldens only. + +References: + OpenMM User Guide §19 "Forces" + Spec: learnable-classical-ff-08-ff-export +""" + +from __future__ import annotations + +from dataclasses import asdict, dataclass, field +from typing import Any + +__all__ = ["ForceSpec"] + + +@dataclass +class ForceSpec: + """Serializable force-field specification for one backend. + + Attributes: + backend: Adapter name (e.g. ``"openmm"``). + forces: List of force-group records (dicts with at least ``"type"``). + scaling: Optional nonbonded scale factors (1–2 / 1–3 / 1–4). + metadata: Free-form provenance / type-system attachments. + """ + + backend: str + forces: list[dict[str, Any]] = field(default_factory=list) + scaling: dict[str, float] | None = None + metadata: dict[str, Any] = field(default_factory=dict) + + def to_dict(self) -> dict[str, Any]: + """Return a JSON-friendly nested dict (deep copy of fields).""" + return asdict(self) + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> ForceSpec: + """Alternate constructor: rebuild from :meth:`to_dict` output. + + Args: + data: Mapping with keys ``backend``, ``forces``, optional + ``scaling`` and ``metadata``. + + Returns: + A new :class:`ForceSpec`. + """ + return cls( + backend=data["backend"], + forces=list(data.get("forces") or []), + scaling=data.get("scaling"), + metadata=dict(data.get("metadata") or {}), + ) diff --git a/src/molix/ff_export/openmm_adapter.py b/src/molix/ff_export/openmm_adapter.py new file mode 100644 index 0000000..3e7b94c --- /dev/null +++ b/src/molix/ff_export/openmm_adapter.py @@ -0,0 +1,253 @@ +"""OpenMM force-spec adapter (pure Python; no live ``openmm`` import). + +Emits JSON-friendly records matching OpenMM User Guide §19 force names: + +* ``HarmonicBondForce`` — ``k`` kJ/mol/nm², ``r0`` nm +* ``HarmonicAngleForce`` — ``k`` kJ/mol/rad², ``theta0`` rad +* ``PeriodicTorsionForce`` — ``k`` kJ/mol, ``periodicity``, ``phase`` rad +* ``NonbondedForce`` — ``charge`` e, ``sigma`` nm, ``epsilon`` kJ/mol + +Particle indices are **not** required for type-level IR bags; records carry +``type_index`` (and ``term_index`` for multi-term torsions). Topology binding +is a later consumer concern. + +References: + OpenMM User Guide §19 "Forces" + Spec: learnable-classical-ff-08-ff-export +""" + +from __future__ import annotations + +from typing import Any + +import torch + +from molix.ff_export.adapter import BackendAdapter +from molix.ff_export.cases import TranslationCase +from molix.ff_export.conventions import ( + ConventionTable, + scale_angle_k, + scale_bond_k, + scale_energy, + scale_length, + scale_torsion_k, +) +from molix.ff_export.exceptions import UnsupportedTermError +from molix.ff_export.force_spec import ForceSpec +from molpot.ir import NonbondedScaling, PotentialIR + +__all__ = ["OpenMMAdapter"] + + +def _as_float_list(t: torch.Tensor) -> list[float]: + return [float(x) for x in t.detach().cpu().reshape(-1).tolist()] + + +def _as_int_list(t: torch.Tensor) -> list[int]: + return [int(x) for x in t.detach().cpu().reshape(-1).tolist()] + + +class OpenMMAdapter(BackendAdapter): + """Translate Class-I :class:`~molpot.ir.PotentialIR` into OpenMM force-spec records. + + Does **not** import OpenMM. Unit tests and regressions pin hard-coded + numeric goldens (bond ``k`` 100 → 41840; torsion Vn=2 → ``k`` 4.184). + """ + + name = "openmm" + + def __init__(self, conventions: ConventionTable | None = None) -> None: + if conventions is None: + conventions = ConventionTable.default_openmm() + self.conventions = conventions + + def translate( + self, + ir: PotentialIR, + *, + meta: dict[str, Any] | None = None, + ) -> ForceSpec: + """Build an OpenMM-oriented :class:`ForceSpec` from ``ir``. + + Args: + ir: Class-I bags in kcal/mol, Å, e, rad. + meta: Optional metadata merged into the force-spec. + + Returns: + Force-spec with OpenMM units (kJ/mol, nm, e, rad). + + Raises: + UnsupportedTermError: For bags marked + :attr:`TranslationCase.UNSUPPORTED` (e.g. harmonic improper). + """ + forces: list[dict[str, Any]] = [] + + if ir.bonds is not None: + forces.append(self._bonds(ir)) + if ir.angles is not None: + forces.append(self._angles(ir)) + if ir.propers is not None: + forces.append(self._propers(ir)) + if ir.impropers_periodic is not None: + forces.append(self._impropers_periodic(ir)) + if ir.impropers_harmonic is not None: + row = self.conventions.row("improper_harmonic") + raise UnsupportedTermError( + "improper_harmonic", + row.notes or "no OpenMM built-in harmonic improper force", + case=row.case, + ) + if ir.lj is not None or ir.charges is not None: + forces.append(self._nonbonded(ir)) + + scaling = self._scaling_dict(ir.scaling) + metadata: dict[str, Any] = { + "unit_system_source": ir.unit_system, + "unit_system_target": "openmm_standard", + } + if meta: + metadata.update(meta) + + return ForceSpec( + backend=self.name, + forces=forces, + scaling=scaling, + metadata=metadata, + ) + + def _bonds(self, ir: PotentialIR) -> dict[str, Any]: + assert ir.bonds is not None + row = self.conventions.row("bond_harmonic") + k = _as_float_list(ir.bonds.k) + r0 = _as_float_list(ir.bonds.r0) + params = [ + { + "type_index": i, + "k": scale_bond_k(ki), + "r0": scale_length(ri), + } + for i, (ki, ri) in enumerate(zip(k, r0, strict=True)) + ] + return { + "type": "HarmonicBondForce", + "case": row.case.value, + "parameters": params, + } + + def _angles(self, ir: PotentialIR) -> dict[str, Any]: + assert ir.angles is not None + row = self.conventions.row("angle_harmonic") + k = _as_float_list(ir.angles.k) + theta0 = _as_float_list(ir.angles.theta0) + params = [ + { + "type_index": i, + "k": scale_angle_k(ki), + "theta0": ti, # rad unchanged + } + for i, (ki, ti) in enumerate(zip(k, theta0, strict=True)) + ] + return { + "type": "HarmonicAngleForce", + "case": row.case.value, + "parameters": params, + } + + def _propers(self, ir: PotentialIR) -> dict[str, Any]: + assert ir.propers is not None + bag = ir.propers + n_types, n_terms = bag.k.shape + periodicity = _as_int_list(bag.periodicity) + params: list[dict[str, Any]] = [] + for t in range(n_types): + idivf = float(bag.idivf[t].item()) + if idivf == 0.0: + raise ValueError(f"proper torsion idivf[{t}] must be non-zero") + for m in range(n_terms): + k_ir = float(bag.k[t, m].item()) / idivf + params.append( + { + "type_index": t, + "term_index": m, + "periodicity": periodicity[m], + "phase": float(bag.phase[t, m].item()), + "k": scale_torsion_k(k_ir), + } + ) + case = TranslationCase.DECOMPOSE if n_terms > 1 else TranslationCase.FORM_REPARAMETERIZE + return { + "type": "PeriodicTorsionForce", + "case": case.value, + "parameters": params, + } + + def _impropers_periodic(self, ir: PotentialIR) -> dict[str, Any]: + assert ir.impropers_periodic is not None + bag = ir.impropers_periodic + n_types, n_terms = bag.k.shape + periodicity = _as_int_list(bag.periodicity) + params: list[dict[str, Any]] = [] + for t in range(n_types): + idivf = float(bag.idivf[t].item()) + if idivf == 0.0: + raise ValueError(f"improper periodic idivf[{t}] must be non-zero") + for m in range(n_terms): + k_ir = float(bag.k[t, m].item()) / idivf + params.append( + { + "type_index": t, + "term_index": m, + "periodicity": periodicity[m], + "phase": float(bag.phase[t, m].item()), + "k": scale_torsion_k(k_ir), + } + ) + case = TranslationCase.DECOMPOSE if n_terms > 1 else TranslationCase.FORM_REPARAMETERIZE + return { + "type": "PeriodicTorsionForce", + "role": "improper", + "case": case.value, + "parameters": params, + "index_convention": "molrs_center_first", + } + + def _nonbonded(self, ir: PotentialIR) -> dict[str, Any]: + row_lj = self.conventions.row("lj") + row_q = self.conventions.row("charge") + params: list[dict[str, Any]] = [] + charges: list[float] = [] + + if ir.lj is not None: + eps = _as_float_list(ir.lj.epsilon) + sig = _as_float_list(ir.lj.sigma) + for i, (e, s) in enumerate(zip(eps, sig, strict=True)): + params.append( + { + "type_index": i, + "epsilon": scale_energy(e), + "sigma": scale_length(s), + } + ) + if ir.charges is not None: + charges = _as_float_list(ir.charges.q) + + return { + "type": "NonbondedForce", + "case": row_lj.case.value, + "parameters": params, + "charges": charges, + "charge_case": row_q.case.value, + } + + @staticmethod + def _scaling_dict(scaling: NonbondedScaling | None) -> dict[str, float] | None: + if scaling is None: + return None + return { + "scale_q_12": float(scaling.scale_q_12), + "scale_q_13": float(scaling.scale_q_13), + "scale_q_14": float(scaling.scale_q_14), + "scale_lj_12": float(scaling.scale_lj_12), + "scale_lj_13": float(scaling.scale_lj_13), + "scale_lj_14": float(scaling.scale_lj_14), + } diff --git a/src/molix/hooks/__init__.py b/src/molix/hooks/__init__.py index dd52109..003d244 100644 --- a/src/molix/hooks/__init__.py +++ b/src/molix/hooks/__init__.py @@ -1,12 +1,14 @@ """Concrete hook implementations driven by :class:`molix.core.trainer.Trainer`. -The contract layer (:class:`Hook` Protocol, :class:`BaseHook`, -:class:`ScalarHook`) lives in :mod:`molix.core.hook`; this package -holds every concrete implementation. Naming convention: every -concrete class carries the ``Hook`` suffix -(``CheckpointHook``, ``JournalHook``, ``GradClipHook``, …) so the -boundary against the contract layer (whose three classes are -suffix-free) stays visually unambiguous. +The contract layer (:class:`~molix.core.hook.Hook` Protocol, +:class:`~molix.core.hook.BaseHook`, :class:`~molix.core.hook.ScalarHook`) +lives in the *singular* module :mod:`molix.core.hook`; this *plural* +package, :mod:`molix.hooks`, holds every concrete implementation. The +singular/plural module name is the boundary marker — class names are not, +since the ``Hook`` suffix appears on both sides (``BaseHook`` / +``ScalarHook`` in the contract layer; ``CheckpointHook`` / ``JournalHook`` / +``GradClipHook`` here) and some concrete hooks drop it entirely (``Log``, +``EarlyStop``). Dependency direction: ``hooks/ → io/ + core/`` — never the reverse. """ diff --git a/src/molix/hooks/molrec_metrics.py b/src/molix/hooks/molrec_metrics.py index 92879e8..a05613b 100644 --- a/src/molix/hooks/molrec_metrics.py +++ b/src/molix/hooks/molrec_metrics.py @@ -114,9 +114,7 @@ def on_train_end(self, trainer: Any, state: Any) -> None: ) self._writer = None - def _log_namespaces( - self, state: Any, namespaces: tuple[str, ...], *, stage: str - ) -> None: + def _log_namespaces(self, state: Any, namespaces: tuple[str, ...], *, stage: str) -> None: assert self._writer is not None step = int(getattr(state, "global_step", 0)) write_status( diff --git a/src/molix/hooks/training.py b/src/molix/hooks/training.py index 95fc388..31e90ec 100644 --- a/src/molix/hooks/training.py +++ b/src/molix/hooks/training.py @@ -104,4 +104,4 @@ def checkpointed_forward(*args, **kwargs): return checkpointed_forward - module.forward = _make_checkpointed(original_forward) # type: ignore[method-assign] + module.forward = _make_checkpointed(original_forward) diff --git a/src/molix/io/metrics.py b/src/molix/io/metrics.py index 8e3ee26..5e0f3d7 100644 --- a/src/molix/io/metrics.py +++ b/src/molix/io/metrics.py @@ -61,11 +61,7 @@ def index_path(record_root: Path) -> Path: def _is_number(value: JSONValue) -> bool: - return ( - isinstance(value, (int, float)) - and not isinstance(value, bool) - and math.isfinite(value) - ) + return isinstance(value, (int, float)) and not isinstance(value, bool) and math.isfinite(value) def _validate_tags(tags: JSONValue) -> dict[str, JSONValue] | None: @@ -395,9 +391,7 @@ def json( def log( self, record: MetricRecord, *, tags: dict[str, JSONValue] | None = None ) -> MetricRecord: - payload: MetricRecord = { - key: value for key, value in record.items() if value is not None - } + payload: MetricRecord = {key: value for key, value in record.items() if value is not None} payload.setdefault("w", _format_wall_time(None)) if tags is not None: payload["tags"] = tags diff --git a/src/molix/logger.py b/src/molix/logger.py index f6e911c..5008a72 100644 --- a/src/molix/logger.py +++ b/src/molix/logger.py @@ -9,4 +9,4 @@ def getLogger(name: str) -> logging.Logger: """Get a logger instance with the specified name.""" - return get_logger(name) # type: ignore + return get_logger(name) diff --git a/src/molix/logging.py b/src/molix/logging.py index 52822df..12c5a64 100644 --- a/src/molix/logging.py +++ b/src/molix/logging.py @@ -22,7 +22,7 @@ Records carry an ``extra["kind"]`` tag that formatters dispatch on: -* ``header`` / ``row`` — structured table emission by :class:`molix.core.hooks.Log` +* ``header`` / ``row`` — structured table emission by :class:`molix.hooks.Log` * ``epoch_sep`` — full-width ``────`` separator at epoch boundaries * ``announce`` — thin ``─── message ───`` separator for intermittent events @@ -153,7 +153,7 @@ def basicConfig( # noqa: N802 — stdlib naming Formatters ---------- By default the stream handler is wired to :class:`PrettyTextFormatter` - so :class:`molix.core.hooks.Log`'s ``metrics`` / ``events`` records + so :class:`molix.hooks.Log`'s ``metrics`` / ``events`` records render as aligned tables and ``─── message ───`` separators out of the box; the file handler uses the structured :class:`TextFormatter`. Pass *stream_formatter* / *file_formatter* (or *formatter* as a @@ -243,7 +243,7 @@ def has_effective_handlers(logger: Logger | None = None) -> bool: Walks ``logger`` (defaulting to the ``molix`` root) up to — but **not** including — the mollog root, whose default stderr ``StreamHandler`` would otherwise always return True and prevent - the :class:`~molix.core.hooks.Log` hook from printing under a + the :class:`~molix.hooks.Log` hook from printing under a zero-config / unit-test harness. A False return means "no molix-level handler is attached", so callers should fall back to :func:`print` for visible output. @@ -280,7 +280,7 @@ def get_table_width() -> int: # Distinct from ``"nan"`` so silent path-resolution failures no longer look # like training divergence in the rendered table. Mirrored by -# :data:`molix.core.hooks._MISSING_CELL`. +# :data:`molix.hooks.progress._MISSING_CELL`. _MISSING_METRIC_CELL = "—" @@ -325,7 +325,7 @@ def split_header_rows(columns: list[str], col_width: int) -> tuple[str, str]: Top row shows the category (namespace prefix before ``/``), bottom row shows the item name; columns with no slash leave the top row blank. Shared by :class:`PrettyTextFormatter` and - :meth:`molix.core.hooks.Log._emit_header` so both render identically. + :meth:`molix.hooks.Log._emit_header` so both render identically. Args: columns: Display names — ``"train/loss"``, ``"epoch"``, … @@ -360,7 +360,7 @@ class PrettyTextFormatter(Formatter): ---------- col_width: Width of each column in ``header`` / ``row`` rendering. Must - match the ``fmt`` width used by :class:`molix.core.hooks.Log`. + match the ``fmt`` width used by :class:`molix.hooks.Log`. row_fmt: Numeric format applied to each value in a ``row`` record (``"{:>12.4g}"`` by default — width + general-precision 4). @@ -470,7 +470,7 @@ def filter(self, record: LogRecord) -> bool: class _HeaderOncePerColumnSet(Filter): """Drop ``kind=header`` records whose column set matches the previous one. - The :class:`molix.core.hooks.Log` hook re-emits a ``header`` record + The :class:`molix.hooks.Log` hook re-emits a ``header`` record at every periodic reprint and after each epoch boundary so the console view stays readable. For a CSV sink those repeats would produce duplicate header lines, confusing ``pandas.read_csv`` and @@ -558,7 +558,7 @@ def configure_run( traces. col_width / row_fmt: Passed straight through to :class:`PrettyTextFormatter`. Must - match the :class:`molix.core.hooks.Log` ``fmt`` width. + match the :class:`molix.hooks.Log` ``fmt`` width. stream: Override for the stdout stream (defaults to ``sys.stdout``). diff --git a/src/molix/md/__init__.py b/src/molix/md/__init__.py index 412c2f7..07934d8 100644 --- a/src/molix/md/__init__.py +++ b/src/molix/md/__init__.py @@ -1,63 +1,90 @@ """Component-based, compilable in-process MD engine. -Component layecake (mirrors molpy's ``Potential`` vs ``ForceField`` split): +:class:`~molix.md.driver.MD` is **the** entry point: it binds a force field to +an integrator, owns the MD-side precision, and delegates the loop to +:class:`~molix.md.runner.MDRunner`. The neighbour-list cadence is *not* its — +that belongs to :class:`~molix.md.neighbors.NeighborList` (see below). The lower +layers are the primitives it composes (use them directly only when you need a +custom loop): -* :class:`~molix.md.types.ForceOutput` / :class:`~molix.md.types.MDState` — - typed pytree contracts crossing component boundaries. -* :class:`~molix.md.forcefield.ForceField` (``PotentialForceField`` over a - molpot Potential; ``HarmonicForceField`` / ``LennardJonesForceField`` analytic) - — binds a model to a system, maps positions to ``(energy, forces)``. -* :class:`~molix.md.integrators.LangevinVerletIntegrator` — advances an +* :class:`~molix.md.types.ForceOutput` / :class:`~molix.md.types.MDState` / + :class:`~molix.md.types.MDObservables` — typed pytree contracts crossing + component boundaries. +* :class:`~molix.md.forcefield.ForceField` (``PotentialForceField`` / + ``PeriodicPotentialForceField`` over a TensorDict potential; + ``CallableForceField`` over any ``pos -> (energy, forces)`` callable; + ``HarmonicForceField`` / ``LennardJonesForceField`` analytic; + ``LennardJonesCutForceField`` — periodic truncated-shifted LJ over a + rebuildable neighbour list, the bulk lj/cut production path) — binds a + model to a system, maps positions to ``(energy, forces)``. +* :class:`~molix.md.integrators.Integrator` / + :class:`~molix.md.integrators.LangevinVerletIntegrator` — advances an ``MDState`` (BAOAB); ``step`` / ``rollout`` ``torch.compile(fullgraph=True)`` - to a single graph including a traceable force field. -* :class:`~molix.md.runner.MDRunner` — drives the integrator through the molix - hook lifecycle; :class:`~molix.md.runner.TrajectoryHook` captures trajectories. + to a single graph including a traceable force field. ``advance_n`` is the + eager chunk driver (γ=0 skips the noise draw; bit-identical dynamics). +* :class:`~molix.md.runner.MDRunner` — drives the integrator through the + :class:`~molix.md.runner.MDHook` lifecycle; + :class:`~molix.md.runner.TrajectoryHook` captures trajectories, + :class:`~molix.md.runner.MDCheckpointHook` persists restartable state. +* :class:`~molix.md.driver.MaxwellBoltzmann` — initial-velocity sampler. -Scope: open (non-periodic) systems, short small-displacement trajectories. The -neighbour list (``edge_index``) is frozen for the whole run — there is no -rebuild — so this is a study/inference engine for near-equilibrium dynamics, not -general production MD. See :class:`~molix.md.forcefield.PotentialForceField`. +Periodic systems are supported through +:class:`~molix.md.neighbors.NeighborList`, which **owns the rebuild policy**: +a Verlet ``skin`` under the LAMMPS ``every`` / ``delay`` / ``check`` gate, +rebuilt in place into fixed-capacity buffers so the force path can stay inside +a CUDA graph. There is exactly one caller — +:meth:`~molix.md.integrators.Integrator.eval_force` asks once per force +evaluation, at the positions being evaluated, iff the force field declares +:attr:`~molix.md.forcefield.ForceField.rebuilds_neighbors`; no driver kwarg and +no step-start hook. A force field that keeps its list frozen (the default for +:class:`~molix.md.forcefield.PotentialForceField`, and any integrator built with +``rebuild=False``) remains valid only for open systems or trajectories short +enough that no atom changes neighbours. """ -from molix.md.ase_shim import HAS_ASE, make_pinet_calculator -from molix.md.dynamics import ( - TrajectoryArtifact, - build_paired_trajectory, - evaluate_delta_along_trajectory, - run_trajectory, -) +from molix.md.driver import MD, MaxwellBoltzmann from molix.md.forcefield import ( + CallableForceField, ForceField, HarmonicForceField, + LennardJonesCutForceField, LennardJonesForceField, + PeriodicPotentialForceField, PotentialForceField, ) -from molix.md.integrators import ( - EV_PER_AMU_A2_FS2, - Integrator, - LangevinVerletIntegrator, - as_mass_col, +from molix.md.integrators import Integrator, LangevinVerletIntegrator +from molix.md.neighbors import NeighborList, NeighborStrategy +from molix.md.runner import ( + MDCheckpointHook, + MDHook, + MDRunner, + TrajectoryHook, ) -from molix.md.runner import MDRunner, TrajectoryHook -from molix.md.types import ForceOutput, MDState +from molix.md.types import ForceOutput, MDObservables, MDState +from molix.units import EV_PER_AMU_A2_FS2, KB_AMU_A_FS, KB_EV_PER_K __all__ = [ "EV_PER_AMU_A2_FS2", - "HAS_ASE", + "KB_AMU_A_FS", + "KB_EV_PER_K", + "MD", + "CallableForceField", "ForceField", "ForceOutput", "HarmonicForceField", "Integrator", "LangevinVerletIntegrator", + "LennardJonesCutForceField", "LennardJonesForceField", + "MDCheckpointHook", + "MDHook", + "MDObservables", "MDRunner", "MDState", + "MaxwellBoltzmann", + "NeighborList", + "NeighborStrategy", + "PeriodicPotentialForceField", "PotentialForceField", - "TrajectoryArtifact", "TrajectoryHook", - "as_mass_col", - "build_paired_trajectory", - "evaluate_delta_along_trajectory", - "make_pinet_calculator", - "run_trajectory", ] diff --git a/src/molix/md/ase_shim.py b/src/molix/md/ase_shim.py deleted file mode 100644 index 4e919d9..0000000 --- a/src/molix/md/ase_shim.py +++ /dev/null @@ -1,57 +0,0 @@ -"""Optional ASE-Calculator shell over the PiNet force seam. - -ASE is an optional dependency: this module degrades gracefully when it is absent -(``HAS_ASE`` is ``False`` and :func:`make_pinet_calculator` raises). The in-process -integrator remains the primary path; the calculator is a thin wrapper that reuses -the same force seam so external ASE drivers can call PiNet. -""" - -from __future__ import annotations - -import torch -from tensordict import TensorDict -from torch import nn - -from molix.md.forcefield import PotentialForceField - -try: - from ase.calculators.calculator import Calculator - - HAS_ASE = True -except ImportError: # pragma: no cover - exercised only where ASE is absent - HAS_ASE = False - - -def make_pinet_calculator(model: nn.Module, template: TensorDict): - """Build an ASE ``Calculator`` delegating energy/forces to the PiNet force seam. - - Args: - model: A ``PiNetPotential``. - template: Molecule TensorDict carrying topology (positions are overwritten - per ``calculate`` from the ASE ``Atoms``). - - Returns: - An ASE ``Calculator`` instance. - - Raises: - RuntimeError: If ASE is not importable. - """ - if not HAS_ASE: - raise RuntimeError("ASE is not installed; use the in-process integrator instead.") - - force = PotentialForceField(model, template) - ref_pos = template["atoms", "pos"] - - class PiNetCalculator(Calculator): # type: ignore[misc, valid-type] - implemented_properties = ("energy", "forces") # noqa: RUF012 - - def calculate(self, atoms=None, properties=("energy",), system_changes=None): # noqa: ANN001, ANN204 - super().calculate(atoms, properties, system_changes or []) - pos = torch.as_tensor( - self.atoms.get_positions(), dtype=ref_pos.dtype, device=ref_pos.device - ) - out = force(pos) - self.results["energy"] = float(out.energy) - self.results["forces"] = out.forces.cpu().numpy() - - return PiNetCalculator() diff --git a/src/molix/md/driver.py b/src/molix/md/driver.py new file mode 100644 index 0000000..2bc3e2a --- /dev/null +++ b/src/molix/md/driver.py @@ -0,0 +1,292 @@ +"""``MD`` — the molecular-dynamics component. + +A domain type with methods, in the sense the project's design rules mean it: +"I have a force field and a system; run dynamics at this precision." It owns the +two things a trajectory needs held together and that no lower layer can decide +alone — the **integrator** and the **MD-side precision** — and delegates the +loop to :class:`~molix.md.runner.MDRunner`. + +It does **not** own the neighbour-list cadence, and there is no kwarg for one. +That decision belongs to :class:`~molix.md.neighbors.NeighborList` +(``skin`` / ``every`` / ``delay`` / ``check``); the integrator asks it once per +force evaluation, at the positions being evaluated, and derives *whether* to ask +from :attr:`~molix.md.forcefield.ForceField.rebuilds_neighbors`. A run that must +keep its list frozen composes it at the caller, through the existing +``integrator=`` seam:: + + MD(ff, mass=m, integrator=LangevinVerletIntegrator( + ff, dt=0.5, gamma=0.0, kbt=0.0, mass=m, rebuild=False)) + +Precision is split in two, deliberately. ``MD(dtype=)`` governs the **MD side +only** — trajectory state (positions/velocities), the integrator's step +constants, and the mass — never the potential. The potential's precision is an +independent axis set explicitly via :meth:`MD.set_potential_dtype` (or by +constructing the force field at the desired dtype), because studying the MD +process and the inference process separately is the point: an fp64 trajectory +over an fp32 model is a meaningful, supported configuration. At the component +boundary the integrator casts the force field's output back into the state +dtype (``Integrator.eval_force``), so the two precisions never silently +promote mid-step:: + + MD(force, mass=m, dt=0.5, dtype=torch.float64) # fp64 trajectory + md.set_potential_dtype(torch.float32) # fp32 inference + MD(force, mass=m, dt=0.5, autocast_dtype=torch.bfloat16) # bf16-mixed model + +``autocast_dtype`` leaves parameters alone and wraps each force evaluation in +``torch.autocast``, which is the only mixed-precision form :mod:`molix.config` +supports (pure fp16/bf16 parameters are deliberately unsupported there, and +that stance is kept here). +""" + +from __future__ import annotations + +from collections.abc import Sequence + +import torch + +from molix.md.forcefield import ForceField +from molix.md.integrators import Integrator, LangevinVerletIntegrator +from molix.md.runner import MDHook, MDRunner +from molix.md.types import ForceOutput, MDState +from molix.units import KB_AMU_A_FS + + +class _AutocastForceField(ForceField): + """Wrap a force field so each evaluation runs under ``torch.autocast``. + + Applied around the *force field* rather than the integrator so the reduced + precision covers the model only: the BAOAB arithmetic and the accumulated + positions/velocities stay in the state dtype, which is what keeps a long + trajectory from losing the low bits of its own coordinates. + """ + + def __init__(self, inner: ForceField, dtype: torch.dtype) -> None: + super().__init__() + self.inner = inner + self.autocast_dtype = dtype + + @property + def rebuilds_neighbors(self) -> bool: + """Delegate: wrapping for precision must not hide a live list. + + ``MD`` applies this wrapper *before* the integrator is constructed, so + the integrator would otherwise derive ``False`` from the wrapper's + default and a bf16 run would silently freeze its neighbour list. + """ + return self.inner.rebuilds_neighbors + + def rebuild_neighbors(self, pos: torch.Tensor) -> None: + """Delegate: connectivity is not a precision concern.""" + self.inner.rebuild_neighbors(pos) + + def forward(self, pos: torch.Tensor) -> ForceOutput: + device_type = pos.device.type + with torch.autocast(device_type=device_type, dtype=self.autocast_dtype): + out = self.inner(pos) + # Hand the integrator back its own precision; autocast is the model's + # business, not the trajectory's. + return ForceOutput(out.energy.to(pos.dtype), out.forces.to(pos.dtype)) + + +class MaxwellBoltzmann: + """Maxwell-Boltzmann initial-velocity sampler over a mass profile. + + Its own tiny type rather than a method on :class:`MD`: sampling initial + conditions is not the run driver's responsibility, and a sampler only + needs the masses. Velocities come back in float64 on CPU; :meth:`MD.run` + casts them onto the run's dtype/device. + + Args: + mass: Per-atom mass ``(N,)`` in amu, or a scalar with ``n_atoms``. + n_atoms: Atom count — required when ``mass`` is a scalar (a scalar + mass carries no system size); checked against a per-atom mass. + """ + + def __init__(self, mass: float | torch.Tensor, *, n_atoms: int | None = None) -> None: + mass_t = torch.as_tensor(mass, dtype=torch.float64) + if mass_t.dim() == 0: + if n_atoms is None: + raise ValueError("scalar mass carries no system size; pass n_atoms") + mass_t = mass_t.expand(int(n_atoms)) + elif n_atoms is not None and int(mass_t.shape[0]) != int(n_atoms): + raise ValueError(f"n_atoms={n_atoms} disagrees with mass length {mass_t.shape[0]}") + self.mass = mass_t.detach().cpu() + + def sample(self, temperature: float, *, seed: int = 0, remove_com: bool = True) -> torch.Tensor: + """Draw velocities ``(N, 3)`` in Å/fs at ``temperature``. + + Args: + temperature: Kelvin. + seed: Generator seed, so a run is reproducible from its arguments. + remove_com: Remove the centre-of-mass momentum, which is what makes + the NVE degrees of freedom ``3N - 3``. + """ + mass_col = self.mass.reshape(-1, 1) + generator = torch.Generator().manual_seed(int(seed)) + noise = torch.randn((mass_col.shape[0], 3), generator=generator, dtype=torch.float64) + vel = noise * torch.sqrt(KB_AMU_A_FS * float(temperature) / mass_col) + if remove_com: + vel = vel - (mass_col * vel).sum(0) / mass_col.sum() + return vel + + +class MD: + """Molecular dynamics over a :class:`~molix.md.forcefield.ForceField`. + + Args: + force: The force field. Its dtype is **not** touched by ``dtype`` — + use :meth:`set_potential_dtype` (or construct it at the precision + you want) to control the inference side separately. + mass: Per-atom mass ``(N,)`` or a scalar, in amu. + dt: Timestep in fs. Required unless ``integrator`` is given. + gamma: Langevin friction γ in fs⁻¹. ``0`` (default) integrates NVE — + BAOAB's O step becomes the identity. + kbt: Thermal energy for the Langevin O step, in the integrator's + (amu, Å, fs) energy unit. Ignored at ``gamma=0``. + temperature: Convenience alternative to ``kbt``, in kelvin. + integrator: A constructed :class:`~molix.md.integrators.Integrator` + over ``force`` — the seam for Nosé–Hoover, NPT, or any + non-Langevin scheme. Mutually exclusive with + ``dt``/``gamma``/``kbt``/``temperature``/``seed``, which + parameterise the default :class:`LangevinVerletIntegrator`. + dtype: MD-side precision — trajectory state, integrator constants, + mass. ``None`` keeps the caller tensors' dtype. The force field is + deliberately left alone (see the module docstring). + autocast_dtype: Run each force evaluation under ``torch.autocast`` at + this dtype, leaving parameters alone. This is the mixed-precision + path (e.g. ``torch.bfloat16``). Not combinable with an explicit + ``integrator`` (wrap the force field yourself in that case). + hooks: Hooks driving the observation lifecycle. + seed: Seed for the Langevin noise. + device: Device to place the force field and state on (device, unlike + dtype, must be shared by both sides). + + There is deliberately **no** rebuild-cadence argument. Configure the policy + where it lives, on the list — + ``NeighborList(cell=…, cutoff=…, positions=…, skin=1.0, every=1, delay=0, + check=True)`` — and ``MD`` derives the rest: the integrator asks that policy + once per force evaluation iff the force field declares + :attr:`~molix.md.forcefield.ForceField.rebuilds_neighbors`. ``skin=0, + every=1, delay=0, check=True`` is the accurate no-skin limit (rebuild + whenever anything moved at all); ``skin > 0`` is the production setting, and + a frozen list is ``integrator=LangevinVerletIntegrator(…, rebuild=False)``. + """ + + def __init__( + self, + force: ForceField, + *, + mass: float | torch.Tensor, + dt: float | None = None, + gamma: float = 0.0, + kbt: float | None = None, + temperature: float | None = None, + integrator: Integrator | None = None, + dtype: torch.dtype | None = None, + autocast_dtype: torch.dtype | None = None, + hooks: Sequence[MDHook | tuple[MDHook, int]] | None = None, + seed: int = 0, + device: torch.device | str | None = None, + ) -> None: + self.dtype = dtype + self.device = torch.device(device) if device is not None else None + + mass_t = torch.as_tensor(mass) + if dtype is not None: + mass_t = mass_t.to(dtype) + if self.device is not None: + mass_t = mass_t.to(self.device) + self.mass = mass_t + + if integrator is not None: + if dt is not None or kbt is not None or temperature is not None: + raise ValueError( + "integrator= is mutually exclusive with dt/gamma/kbt/temperature/seed — " + "those parameterise the default LangevinVerletIntegrator; a constructed " + "integrator already owns them" + ) + if autocast_dtype is not None: + raise ValueError( + "autocast_dtype cannot wrap a constructed integrator's force field; " + "wrap the force field before building the integrator" + ) + if integrator.force is not force: + raise ValueError("integrator.force must be the force field given to MD") + else: + if dt is None: + raise ValueError("dt is required when no integrator is given") + if kbt is not None and temperature is not None: + raise ValueError("give kbt or temperature, not both") + if kbt is None: + kbt = KB_AMU_A_FS * float(temperature) if temperature is not None else 0.0 + if gamma > 0.0 and kbt == 0.0: + raise ValueError("gamma > 0 needs kbt or temperature (a thermostat at 0 K freezes)") + + # Wrapping happens before the integrator is built, so the integrator + # derives its rebuild switch from the wrapper — which forwards + # rebuilds_neighbors to the inner force field, keeping one owner. + if autocast_dtype is not None: + force = _AutocastForceField(force, autocast_dtype) + self.autocast_dtype = autocast_dtype + + if integrator is None: + integrator = LangevinVerletIntegrator( + force, dt=dt, gamma=gamma, kbt=kbt, mass=mass_t, seed=seed + ) + # MD-side precision: the integrator's own step constants follow the + # run dtype; the force field inside it is deliberately not cast. + if dtype is not None: + integrator = integrator.cast_state(dtype) + # Device, unlike dtype, is shared — a cross-device force call cannot + # work — so the move recurses through the force field too. + if self.device is not None: + integrator = integrator.to(self.device) + self.force = force + self.integrator = integrator + + # No cadence poke here, by design: the integrator derived its own + # rebuild switch from the force field it holds, and the list owns when. + run_hooks: list[MDHook | tuple[MDHook, int]] = list(hooks or []) + self.runner = MDRunner(integrator, mass=mass_t, hooks=run_hooks) + + def set_potential_dtype(self, dtype: torch.dtype) -> "MD": + """Cast the force field (the inference side) to ``dtype``. + + The explicit counterpart of the constructor's ``dtype``: MD-side and + potential-side precision are independent axes, so casting the model is + never a side effect of setting the trajectory precision. The + integrator's boundary cast keeps the state in the MD dtype regardless + of what the model computes in. + """ + self.force.to(dtype) + return self + + def run( + self, pos: torch.Tensor, vel: torch.Tensor, n_steps: int, *, chunk: int | None = None + ) -> MDState: + """Integrate ``n_steps``, returning the final typed :class:`MDState`. + + Trajectory capture is a hook's job — pass a + :class:`~molix.md.runner.TrajectoryHook`. + + Args: + pos: Initial positions ``(N, 3)``, Angstrom. + vel: Initial velocities ``(N, 3)``, Å/fs. + n_steps: Number of steps. + chunk: Steps advanced between hook firings (dynamics are + bit-identical; only observation cadence changes). Default 1. + The neighbour policy is independent of ``chunk`` — it runs + inside each force evaluation, on the list's own schedule. + """ + if chunk is None: + chunk = 1 + pos, vel = self._cast(pos), self._cast(vel) + return self.runner.run(pos, vel, n_steps, chunk=chunk) + + def _cast(self, tensor: torch.Tensor) -> torch.Tensor: + """Bring a caller tensor onto the run's dtype/device.""" + if self.dtype is not None: + tensor = tensor.to(self.dtype) + if self.device is not None: + tensor = tensor.to(self.device) + return tensor diff --git a/src/molix/md/dynamics.py b/src/molix/md/dynamics.py deleted file mode 100644 index 1d238cf..0000000 --- a/src/molix/md/dynamics.py +++ /dev/null @@ -1,170 +0,0 @@ -"""Paired reference/quantized trajectory protocol and the trajectory artifact. - -The reference (fp64) potential drives a trajectory x(t). Along that same x(t) both -the reference and the quantized potential are evaluated to give the per-frame force -residual ΔF(t) = F_quant - F_ref — the time series the thermal-noise diagnostics -(spec -03) consume. A second, independent quantized-driven trajectory is available -for observable comparison. -""" - -from __future__ import annotations - -from dataclasses import dataclass, field -from typing import Any - -import torch -from tensordict import TensorDict -from torch import nn - -from molix.md.forcefield import PotentialForceField -from molix.md.integrators import LangevinVerletIntegrator - -_INTEGRATOR_NAME = "velocity-verlet+langevin-baoab" - - -@dataclass(frozen=True) -class TrajectoryArtifact: - """Per-step paired-trajectory record. - - Units follow the caller's ``force_fn`` / integrator (the engine is - unit-agnostic — see :mod:`molix.md.integrators`); they are not assumed to be - eV/Å. Shapes: ``pos``/``vel``/``f_ref``/``f_quant``/``df`` are ``(T, N, 3)``; - ``energy`` is ``(T,)``. ``metadata`` carries the run scalars (``dt``, - ``gamma``, ``kbt``, ``mass``, ``N``, ``dof``, ``seed``, ``integrator``, plus - any caller-supplied condition keys). ``dof`` is ``3N`` under Langevin and - ``3N-3`` under NVE, matching :class:`molix.md.MDRunner`. - """ - - pos: torch.Tensor - vel: torch.Tensor - energy: torch.Tensor - f_ref: torch.Tensor - f_quant: torch.Tensor - df: torch.Tensor - metadata: dict[str, Any] = field(default_factory=dict) - - @property - def n_steps(self) -> int: - return int(self.pos.shape[0]) - - @property - def n_atoms(self) -> int: - return int(self.pos.shape[1]) - - def to_dict(self) -> dict[str, Any]: - """Flatten to a plain dict for ``torch.save`` (the on-disk ``.pt`` schema).""" - return { - "pos": self.pos, - "vel": self.vel, - "energy": self.energy, - "f_ref": self.f_ref, - "f_quant": self.f_quant, - "df": self.df, - "metadata": self.metadata, - } - - -def run_trajectory( - model: nn.Module, - template: TensorDict, - pos0: torch.Tensor, - vel0: torch.Tensor, - n_steps: int, - *, - dt: float, - gamma: float, - kbt: float, - mass: float | torch.Tensor, - seed: int = 0, -) -> dict[str, torch.Tensor]: - """Drive ``n_steps`` of Langevin velocity-Verlet using ``model``'s forces.""" - force = PotentialForceField(model, template) - integ = LangevinVerletIntegrator(force, dt=dt, gamma=gamma, kbt=kbt, mass=mass, seed=seed) - return integ.run(pos0, vel0, n_steps) - - -def evaluate_delta_along_trajectory( - model_ref: nn.Module, - model_quant: nn.Module, - template: TensorDict, - pos_traj: torch.Tensor, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Per-frame ``(F_ref, F_quant, ΔF)`` along a fixed position trajectory. - - ``ΔF = F_quant - F_ref`` is a catastrophic-cancellation site (the quantized - force tracks the reference, so ``|F| / |ΔF| ≫ 1``). The subtraction is done - in float64 to avoid fp32 round-off (~1e-7·|F|) swamping the residual; for the - result to be meaningful **both models should also be evaluated in float64** - (the rounding inside an fp32 forward cannot be recovered here). - """ - ref_ff = PotentialForceField(model_ref, template) - quant_ff = PotentialForceField(model_quant, template) - # Chunked evaluation keeps peak memory bounded while still avoiding a - # Python-level force eval per frame when the trajectory is long. - chunk = 32 - f_ref: list[torch.Tensor] = [] - f_quant: list[torch.Tensor] = [] - n_frames = int(pos_traj.shape[0]) - for start in range(0, n_frames, chunk): - end = min(start + chunk, n_frames) - for pos in pos_traj[start:end]: - f_ref.append(ref_ff.calc_forces(pos)) - f_quant.append(quant_ff.calc_forces(pos)) - f_ref_t = torch.stack(f_ref) - f_quant_t = torch.stack(f_quant) - df = (f_quant_t.to(torch.float64) - f_ref_t.to(torch.float64)).to(f_quant_t.dtype) - return f_ref_t, f_quant_t, df - - -def build_paired_trajectory( - model_ref: nn.Module, - model_quant: nn.Module, - template: TensorDict, - pos0: torch.Tensor, - vel0: torch.Tensor, - n_steps: int, - *, - dt: float, - gamma: float, - kbt: float, - mass: float, - seed: int = 0, - condition: dict[str, Any] | None = None, -) -> TrajectoryArtifact: - """Run the reference trajectory and evaluate ΔF(t) along it. - - Returns a :class:`TrajectoryArtifact` with the reference pos/vel/energy plus the - paired ``f_ref`` / ``f_quant`` / ``df`` time series and run metadata. - """ - ref = run_trajectory( - model_ref, template, pos0, vel0, n_steps, dt=dt, gamma=gamma, kbt=kbt, mass=mass, seed=seed - ) - f_ref, f_quant, df = evaluate_delta_along_trajectory( - model_ref, model_quant, template, ref["pos"] - ) - n_atoms = int(pos0.shape[0]) - # Match MDRunner's temperature convention: Langevin (γ>0) thermostats all 3N - # DoF including the COM; NVE with COM removed leaves 3N-3. - dof = 3 * n_atoms - (0 if gamma > 0.0 else 3) - metadata: dict[str, Any] = { - "dt": dt, - "gamma": gamma, - "kbt": kbt, - "mass": mass, - "N": n_atoms, - "dof": dof, - "seed": seed, - "n_steps": n_steps, - "integrator": _INTEGRATOR_NAME, - } - if condition: - metadata.update(condition) - return TrajectoryArtifact( - pos=ref["pos"], - vel=ref["vel"], - energy=ref["energy"], - f_ref=f_ref, - f_quant=f_quant, - df=df, - metadata=metadata, - ) diff --git a/src/molix/md/forcefield.py b/src/molix/md/forcefield.py index aaf2146..077db87 100644 --- a/src/molix/md/forcefield.py +++ b/src/molix/md/forcefield.py @@ -8,22 +8,42 @@ * a **Potential** is a :class:`molpot.BasePotential` / ``PiNetPotential`` — energy from a batch ``TensorDict``; * a **ForceField** (this module) is an :class:`torch.nn.Module` that *binds* a - Potential (or an analytic form) to a system and maps positions ``(N, 3)`` to a - :class:`~molix.md.types.ForceOutput`. The ``Integrator`` consumes a - ``ForceField`` component — never a closure. - -Scope: open (non-periodic) systems, frozen neighbour list, small displacements -(see :class:`PotentialForceField`). PBC / minimum-image and neighbour-list -rebuild are out of scope. + Potential (or an analytic form, or any callable) to a system and maps + positions ``(N, 3)`` to a :class:`~molix.md.types.ForceOutput`. The + ``Integrator`` consumes a ``ForceField`` component — never a closure. + +Precision: a force field owns its own dtype, independent of the trajectory +state's (see :class:`molix.md.driver.MD` — ``MD(dtype=)`` governs the MD side +only; :meth:`MD.set_potential_dtype` casts the force field). Implementations +accept positions in any dtype and return their own; the integrator casts the +output back to the state dtype at the component boundary. + +Periodic systems and neighbour-list refresh are supported through +:attr:`ForceField.rebuilds_neighbors` (does this force field own a policy?), +:meth:`ForceField.rebuild_neighbors` (run it at these positions) and +:class:`~molix.md.neighbors.NeighborList` (the policy itself). The **list** +owns the cadence — ``skin`` / ``every`` / ``delay`` / ``check`` — and +:meth:`molix.md.integrators.Integrator.eval_force` asks it once per force +evaluation, at the positions being evaluated. :class:`PotentialForceField` +keeps its list frozen — valid for open systems and trajectories short enough +that no atom changes neighbours; periodic production runs use +:class:`PeriodicPotentialForceField` (TensorDict potentials consuming +``edges.shifts``), :class:`LennardJonesCutForceField` (analytic lj/cut over +the same rebuildable list), or :class:`CallableForceField` with a rebuildable +list. """ from __future__ import annotations +from collections.abc import Callable + import torch from tensordict import TensorDict from torch import nn +from molix.md.neighbors import NeighborStrategy from molix.md.types import ForceOutput +from molix.schema import ENERGY_KEY, FORCES_KEY, has_forces class ForceField(nn.Module): @@ -35,9 +55,60 @@ class ForceField(nn.Module): an energy-only path is cheaper than a full force evaluation. """ - def forward(self, pos: torch.Tensor) -> ForceOutput: # noqa: D102 + def forward(self, pos: torch.Tensor) -> ForceOutput: + """Energy + forces at ``pos`` ``(N, 3)``.""" raise NotImplementedError + @property + def rebuilds_neighbors(self) -> bool: + """Whether this force field owns a neighbour policy worth asking. + + ``False`` by default — a force field with no list (the analytic ones) + has nothing to refresh. A read-only capability property with a + documented default that subclasses override, mirroring + :attr:`molix.md.integrators.Integrator.removed_dof`; it exists so the + integrator can derive its static rebuild switch from a *declared* + capability instead of duck-reading ``getattr(force, "neighbors", None)``. + + Two different questions, hence two names: this one answers *can this + force field run a policy*, while + :attr:`molix.md.integrators.Integrator.rebuild` answers *does this + integrator ask*. + """ + return False + + def rebuild_neighbors(self, pos: torch.Tensor) -> None: + """Run this force field's neighbour policy at ``pos`` ``(N, 3)``. + + **Not** "rebuild now". The name is historical; the meaning is "the + positions are ``pos``, decide" — implementations delegate to + :meth:`molix.md.neighbors.NeighborList.update`, which applies the + list's own ``skin`` / ``every`` / ``delay`` / ``check`` gate and may + well decline. The list is the single owner of the cadence; this seam + only carries the positions to it. + :meth:`molix.md.integrators.Integrator.eval_force` is the caller, once + per force evaluation, at the positions ``F = -∇E`` is taken at. + + A **forced**, unconditional build is still one call away and is not + this method: ``force_field.neighbors.rebuild(pos)``, the primitive + ``update`` itself delegates to. Counts (``rebuild_count`` / + ``ndanger``) are read off the list, which is where they live — hence + the ``None`` return. + + A no-op by default: force fields with no neighbour list, and those that + deliberately freeze it, need do nothing. Implementations must keep + every tensor **shape** unchanged so a compiled / graph-captured force + path stays valid — see :class:`~molix.md.neighbors.NeighborList`. + + Note: + The policy costs one max-displacement reduction plus a ``float()`` + host sync per force evaluation. That is strictly cheaper than + rebuilding every step, but on GPU it is a per-step device→host sync + a fully captured loop would not have — the accepted price of a + correct, list-owned cadence. Freeze it with + ``Integrator(..., rebuild=False)`` when the list must not move. + """ + def calc_energy(self, pos: torch.Tensor) -> torch.Tensor: """Scalar energy ``()`` at ``pos``.""" return self(pos).energy @@ -51,23 +122,29 @@ class PotentialForceField(ForceField): """Bind a molpot Potential to a fixed system template. The collated template carries a *precomputed* ``edges.edge_diff`` / - ``edge_dist`` from its build-time positions; PiNet's ``_edge_bond_diff`` would + ``edge_dist`` from its build-time positions; PiNet's ``edge_bond_diff`` would use that as a straight-through *value*, freezing the PES (constant force) if left in place. The template is therefore stripped of those keys **once** so the Potential recomputes geometry from the live positions every call (correct for **open** systems). The neighbour list (``edge_index``) is **not** rebuilt: - valid for short, small-displacement, non-periodic trajectories only. + valid for short, small-displacement, non-periodic trajectories only — + periodic runs use :class:`PeriodicPotentialForceField`. + + ``.to(dtype)`` / ``.to(device)`` move the working batch together with the + module parameters (``_apply`` is overridden), so an explicit cast reaches + the whole bound system. Args: - potential: A molpot Potential / ``PiNetPotential``; its - ``forward(td, compute_forces=True)`` returns ``{"energy", "forces"}`` - and must not mutate ``td`` (PiNet clones internally — safe to reuse - the template). + potential: A molpot Potential / ``PiNetPotential``. Its ``forward(td)`` + writes ``graphs.energy`` and ``atoms.forces`` into the batch, per + :mod:`molix.schema`. Force derivation is fixed at the potential's + construction (``compute_forces=True``), not requested per call — + the monomorphic contract ``torch.compile`` needs. template: System ``TensorDict`` (``atoms.Z``, ``edges.edge_index``, ``atoms.batch``, ``graphs``); ``("atoms", "pos")`` is replaced per call. energy_scale: Unit-bridge multiplier applied to energy and forces, e.g. - ``1 / molix.md.integrators.EV_PER_AMU_A2_FS2`` to drive an eV/Å - potential in the integrator's (amu, Å, fs) system. Default ``1.0``. + ``1 / molix.units.EV_PER_AMU_A2_FS2`` to drive an eV/Å potential in + the integrator's (amu, Å, fs) system. Default ``1.0``. """ _STALE_EDGE_KEYS = ("edge_diff", "edge_dist") @@ -77,20 +154,35 @@ def __init__( ) -> None: super().__init__() self.potential = potential - batch = template.clone() - for key in self._STALE_EDGE_KEYS: - if ("edges", key) in batch.keys(include_nested=True): - del batch["edges", key] # Working batch: reuse structure and only replace pos each step # (full TensorDict.clone() every MD step was a measurable alloc cost). # Potential paths that need isolation (PiNet) clone internally. - self._template = batch - self._work = batch.clone() - ref = batch["atoms", "pos"] + work = template.clone() + for key in self._STALE_EDGE_KEYS: + if ("edges", key) in work.keys(include_nested=True): + del work["edges", key] + self._work = work + ref = work["atoms", "pos"] self._device = ref.device self._dtype = ref.dtype self.register_buffer("energy_scale", torch.as_tensor(float(energy_scale))) + def _apply( + self, fn: Callable[[torch.Tensor], torch.Tensor], recurse: bool = True + ) -> "nn.Module": + """Extend ``nn.Module._apply`` to the working batch. + + Without this, ``.to(dtype)`` walks parameters/buffers only and leaves + the bound system (``_work``, plain ``TensorDict``) at its construction + dtype — an explicit cast that silently does nothing. + """ + module = super()._apply(fn, recurse) + self._work = self._work.apply(fn) + ref = self._work["atoms", "pos"] + self._device = ref.device + self._dtype = ref.dtype + return module + def _batch_at(self, pos: torch.Tensor) -> TensorDict: """Bind live positions into the reusable working batch (in-place pos).""" self._work["atoms", "pos"] = pos.to(device=self._device, dtype=self._dtype) @@ -98,16 +190,159 @@ def _batch_at(self, pos: torch.Tensor) -> TensorDict: def forward(self, pos: torch.Tensor) -> ForceOutput: batch = self._batch_at(pos) - out = self.potential(batch, compute_forces=True) - energy = out["energy"].sum().detach() * self.energy_scale - forces = out["forces"].detach() * self.energy_scale + out = self.potential(batch) + if not has_forces(out): + raise RuntimeError( + f"{type(self.potential).__name__} wrote no {FORCES_KEY} — an MD force field " + "needs forces. Potentials fix this at construction now (e.g. " + "PiNetPotential(..., compute_forces=True)); it is no longer a per-call choice." + ) + energy = out[ENERGY_KEY].sum().detach() * self.energy_scale + forces = out[FORCES_KEY].detach() * self.energy_scale return ForceOutput(energy, forces) def calc_energy(self, pos: torch.Tensor) -> torch.Tensor: - """Energy only — skips the force derivation (cheaper than :meth:`forward`).""" - batch = self._batch_at(pos) - out = self.potential(batch, compute_forces=False) - return out["energy"].sum().detach() * self.energy_scale + """Scalar energy ``()`` at ``pos``. + + No longer cheaper than :meth:`forward`: whether a potential derives + forces is fixed when it is constructed, so an energy-only evaluation + means constructing an energy-only potential. + """ + out = self.potential(self._batch_at(pos)) + return out[ENERGY_KEY].sum().detach() * self.energy_scale + + +class PeriodicPotentialForceField(PotentialForceField): + """Bind a TensorDict potential to a periodic system with a rebuilding list. + + The component that joins the pieces the package already ships: the + ``rebuild_neighbors`` seam, :class:`~molix.md.neighbors.NeighborList`'s + ``skin`` / ``every`` / ``delay`` / ``check`` policy, and its + fixed-capacity buffers. The working batch's ``edges`` namespace holds the + list's live ``edge_index`` ``(capacity, 2)`` and ``shifts`` + ``(capacity, 3)`` **by reference**, so an in-place rebuild is visible to + the potential with every tensor shape unchanged (CUDA-graph safe). + + The **list** owns that namespace: this class calls + :meth:`~molix.md.neighbors.NeighborList.build` on the working batch at + construction and again after every cast, and never writes ``edges`` itself. + ``build`` also validates that the template describes the system the list + was constructed for (atom count, device/dtype, and its ``graphs.cell`` if + it carries one), so a mismatched pair is refused here rather than producing + a silently wrong PES. + + The potential's ``forward(td)`` must consume ``edges.shifts`` for periodic + correctness (e.g. :class:`molzoo.MACEMatpes`); dead padding edges + self-annihilate through the cutoff envelope. + + Args: + potential: TensorDict potential writing ``graphs.energy`` / + ``atoms.forces`` and reading ``edges.shifts``. + template: System ``TensorDict``; its ``edges`` namespace is replaced by + the neighbour list's buffers. + neighbors: The system's rebuilding neighbour list. + energy_scale: Unit-bridge multiplier applied to energy and forces. + """ + + def __init__( + self, + potential: nn.Module, + template: TensorDict, + *, + neighbors: NeighborStrategy, + energy_scale: float = 1.0, + ) -> None: + super().__init__(potential, template, energy_scale=energy_scale) + self.neighbors = neighbors + self.neighbors.build(self._work) + + def _apply( + self, fn: Callable[[torch.Tensor], torch.Tensor], recurse: bool = True + ) -> "nn.Module": + """Cast the neighbour list alongside the module, then let it re-bind. + + Both halves are still required. ``TensorDict.apply`` in the parent + produces new leaf tensors and ``NeighborList.to`` rebinds the list's + buffers, so a cast severs the by-reference tie twice over — a later + ``rebuild`` would update tensors the potential no longer sees, and the + PES would freeze silently. Only the *knowledge of how to bind* lives + elsewhere now: :meth:`~molix.md.neighbors.NeighborList.build` owns the + working batch's ``edges`` namespace, and this class merely says when. + """ + module = super()._apply(fn, recurse) + # Direct attribute access, deliberately: ``neighbors`` is a required + # constructor argument, and ``_apply`` cannot fire before construction + # completes — a subclass that violates that fails loud here rather + # than silently skipping the re-bind. + ref = self._work["atoms", "pos"] + self.neighbors.to(ref.device, ref.dtype) + self.neighbors.build(self._work) + return module + + @property + def rebuilds_neighbors(self) -> bool: + """``True`` — a periodic run is exactly the case with a live list.""" + return True + + def rebuild_neighbors(self, pos: torch.Tensor) -> None: + """Ask the list's policy at ``pos``; it rebuilds in place, or declines. + + The rebuild, when it happens, is in place (shapes unchanged), so the + ``edges`` buffers bound into the working batch stay valid. Force an + unconditional build with ``self.neighbors.rebuild(pos)``. + """ + self.neighbors.update(pos) + + +class CallableForceField(ForceField): + """Adapt any ``pos -> (energy, forces)`` callable to the ForceField contract. + + The escape hatch for force providers that are not TensorDict potentials — + an AOTI-exported ``.pt2``, a :class:`molix.engine.StaticForward`, a + compiled energy core with hand-rolled autograd, an external engine. The + callable may return a :class:`~molix.md.types.ForceOutput` or a plain + ``(energy, forces)`` tuple. + + Args: + fn: Maps positions ``(N, 3)`` to scalar energy ``()`` and forces + ``(N, 3)``. + neighbors: Optional rebuildable neighbour list. When one is bound, + :attr:`rebuilds_neighbors` reports ``True`` and + :meth:`rebuild_neighbors` runs its policy, so the integrator + derives its rebuild switch without being told. + energy_scale: Unit-bridge multiplier applied to energy and forces. + """ + + def __init__( + self, + fn: Callable[[torch.Tensor], tuple[torch.Tensor, torch.Tensor]], + *, + neighbors: NeighborStrategy | None = None, + energy_scale: float = 1.0, + ) -> None: + super().__init__() + self._fn = fn + self.neighbors = neighbors + self.register_buffer("energy_scale", torch.as_tensor(float(energy_scale))) + + @property + def rebuilds_neighbors(self) -> bool: + """Per instance: ``True`` iff a list was bound at construction.""" + return self.neighbors is not None + + def rebuild_neighbors(self, pos: torch.Tensor) -> None: + """Ask the bound list's policy at ``pos``; a no-op when none is bound. + + The list may decline (its ``skin`` / ``every`` / ``delay`` / ``check`` + gate). Force an unconditional build with + ``self.neighbors.rebuild(pos)``. + """ + if self.neighbors is not None: + self.neighbors.update(pos) + + def forward(self, pos: torch.Tensor) -> ForceOutput: + energy, forces = self._fn(pos) + return ForceOutput(energy * self.energy_scale, forces * self.energy_scale) class HarmonicForceField(ForceField): @@ -169,3 +404,149 @@ def forward(self, pos: torch.Tensor) -> ForceOutput: coef = (24.0 * self.epsilon * (2.0 * inv_r12 - inv_r6) / r2).masked_fill(eye, 0.0) force = (coef.unsqueeze(-1) * (-diff)).sum(1) # (N, 3) return ForceOutput(energy, force) + + +class LennardJonesCutForceField(ForceField): + """Truncated(-shifted) Lennard-Jones over a rebuildable neighbour list (``lj/cut``). + + The bulk counterpart of :class:`LennardJonesForceField` (which is all-pairs + and open): pair interactions are evaluated on the fixed-capacity buffers of + a :class:`~molix.md.neighbors.NeighborStrategy` and truncated at ``cutoff``, + LAMMPS ``pair_style lj/cut`` style. Energy and forces are the analytic + closed form (no autograd) over tensors whose **shapes never change** across + rebuilds, so ``forward`` stays ``torch.compile(fullgraph=True)`` / + CUDA-graph capturable while the list runs its policy eagerly between force + evaluations (:meth:`molix.md.integrators.Integrator.eval_force`). + + Pairs beyond ``cutoff`` — including the list's dead padding edges, whose + shift is ``DEAD_EDGE_CUTOFF_FACTOR × cutoff`` — contribute exactly zero + energy and force. With ``shift=True`` (default) the pair energy is shifted + by ``E_lj(cutoff)`` so it reaches zero *continuously* at the cutoff; + without the shift, every pair crossing r_cut steps the total energy by + ``E_lj(r_cut)``, which reads as noise/drift in an NVE total-energy trace. + The forces are identical under both conventions. + + Args: + epsilon: Well depth ε, in the run's energy unit (amu·Å²/fs² for the + stock integrator — convert eV via ``1/EV_PER_AMU_A2_FS2``). + sigma: Zero-crossing distance σ (Å). + neighbors: Rebuildable neighbour list with **full bidirectional** + edges (each pair present in both directions), e.g. + :class:`~molix.md.neighbors.NeighborList`. + cutoff: Truncation radius r_cut (Å). Defaults to the list's own + cutoff, and must not exceed it — pairs between the two radii would + simply be absent from the buffers, silently truncating the PES + harder than asked. + shift: Shift pair energies by ``E_lj(cutoff)`` (see above). + + Reference: + Lennard-Jones, "On the Determination of Molecular Fields", Proc. R. + Soc. Lond. A 106 (1924) 463. https://doi.org/10.1098/rspa.1924.0082 + Truncated-and-shifted convention: Allen & Tildesley, "Computer + Simulation of Liquids", 2nd ed. (2017), §5.2. + """ + + def __init__( + self, + *, + epsilon: float, + sigma: float, + neighbors: NeighborStrategy, + cutoff: float | None = None, + shift: bool = True, + ) -> None: + super().__init__() + # The interaction cutoff, not the list's r_build: a skinned list holds + # pairs further out, but only guarantees them *complete* to .cutoff + # between rebuilds, so anything beyond is the silent truncation this + # check exists to prevent. + list_cutoff = float(neighbors.cutoff) + if cutoff is None: + cutoff = list_cutoff + elif float(cutoff) > list_cutoff: + raise ValueError( + f"cutoff {cutoff} A exceeds the neighbour list's horizon {list_cutoff} A; " + "pairs between the two radii would be silently missing from the PES" + ) + self.neighbors = neighbors + self.shift = bool(shift) + eps, sig, r_cut = float(epsilon), float(sigma), float(cutoff) + sr6_cut = (sig / r_cut) ** 6 + self.register_buffer("epsilon", torch.as_tensor(eps)) + self.register_buffer("sigma", torch.as_tensor(sig)) + self.register_buffer("cutoff_sq", torch.as_tensor(r_cut * r_cut)) + self.register_buffer( + "energy_shift", + torch.as_tensor(4.0 * eps * (sr6_cut * sr6_cut - sr6_cut) if shift else 0.0), + ) + + def _apply( + self, fn: Callable[[torch.Tensor], torch.Tensor], recurse: bool = True + ) -> "nn.Module": + """Extend ``nn.Module._apply`` to the neighbour list's buffers. + + ``forward`` reads ``neighbors.edge_index`` / ``shifts`` live; a + ``.to(device/dtype)`` that skipped the list would leave the pair + geometry behind — a cross-device indexing error at best. + """ + module = super()._apply(fn, recurse) + # Direct access (see PeriodicPotentialForceField._apply): required + # attribute, post-construction call site, fail loud over silent skip. + ref = self.epsilon + self.neighbors.to(ref.device, ref.dtype) + return module + + @property + def rebuilds_neighbors(self) -> bool: + """``True`` — lj/cut is defined over a live, rebuildable list.""" + return True + + def rebuild_neighbors(self, pos: torch.Tensor) -> None: + """Ask the list's policy at ``pos``; it rebuilds in place, or declines. + + Inside the half-skin the list keeps the current (superset) buffers, + which the ``cutoff_sq`` mask already reduces to the same PES. Force an + unconditional build with ``self.neighbors.rebuild(pos)``. + + The list is a **geometry** object and tracks the MD-state dtype (see + :meth:`forward`); do not cast ``pos`` to the potential parameter dtype + here — that would pull the Verlet skin check onto the pot axis and make + a split-precision run no longer vary only the force arithmetic. + """ + self.neighbors.update(pos) + + def forward(self, pos: torch.Tensor) -> ForceOutput: + """Truncated-LJ energy + forces at ``pos`` ``(N, 3)`` from the live buffers. + + Positions (and the list's shifts) are cast into the **parameter** + dtype before the pair loop so a split-precision configuration + (``MD(dtype=fp64)`` + pot buffers in fp32) actually evaluates the PES + in pot precision — matching the MACE path, which keeps the neighbour + list on the MD axis and only casts into the model at force time. The + integrator then casts energy/forces back to the trajectory dtype. + """ + pot_dtype = self.epsilon.dtype + pos = pos.to(device=self.epsilon.device, dtype=pot_dtype) + edge_index = self.neighbors.edge_index # (capacity, 2) + source, target = edge_index[:, 0], edge_index[:, 1] + # Minimum-image displacement: live positions + the list's periodic + # remainder, both in pot precision for the arithmetic. + shifts = self.neighbors.shifts.to(dtype=pot_dtype) + diff = pos[target] - pos[source] + shifts # (capacity, 3) + r2 = (diff * diff).sum(-1) + inside = r2 < self.cutoff_sq # dead edges: |shift| ≫ cutoff → excluded here + safe_r2 = torch.where(inside, r2, torch.ones_like(r2)) + inv_r2 = (self.sigma * self.sigma) / safe_r2 + inv_r6 = inv_r2 * inv_r2 * inv_r2 + inv_r12 = inv_r6 * inv_r6 + pair_e = torch.where( + inside, 4.0 * self.epsilon * (inv_r12 - inv_r6) - self.energy_shift, 0.0 + ) + energy = 0.5 * pair_e.sum() # bidirectional edges visit each pair twice + # Per-edge force on the *source*: 24ε/r²·(2(σ/r)¹² − (σ/r)⁶)·(-diff). + # The reverse edge delivers the Newton pair to the other atom, so the + # forces need no ½ — only the energy double-counts. + coef = torch.where(inside, 24.0 * self.epsilon * (2.0 * inv_r12 - inv_r6) / safe_r2, 0.0) + force = torch.zeros_like(pos) + force.index_add_(0, source, coef.unsqueeze(-1) * (-diff)) + return ForceOutput(energy, force) diff --git a/src/molix/md/integrators.py b/src/molix/md/integrators.py index c888c83..aa4ea56 100644 --- a/src/molix/md/integrators.py +++ b/src/molix/md/integrators.py @@ -9,6 +9,15 @@ ``torch.compile(fullgraph=True)`` to a single graph *including* a traceable force field (PiNet's functorch force path is graph-break-free). +Neighbour-list cadence is **not** the integrator's: the list owns it (``skin`` / +``every`` / ``delay`` / ``check`` on :class:`~molix.md.neighbors.NeighborList`) +and :meth:`Integrator.eval_force` merely asks, once per force evaluation, at the +positions being evaluated. What the integrator keeps is a construction-time +on/off, :attr:`Integrator.rebuild`, derived from the force field. Over a live +list that question is an eager, host-syncing decision, so a compiled rollout is +built with ``rebuild=False`` (frozen list) while production compiles the force +field and keeps the loop eager. + BAOAB ordering (Leimkuhler & Matthews): B (half kick) → A (half drift) → O (Ornstein-Uhlenbeck) → A → B. The O step ``v ← c1·v + c2·σ·ξ`` with ``c1 = e^{-γΔt}``, ``c2 = √(1-c1²)``, ``σ = √(k_BT/m)`` satisfies the @@ -18,10 +27,16 @@ Units — one self-consistent system. The arithmetic ``v += Δt·F/m`` needs ``[F] = [m][length]/[time]²``, so eV/Å + amu + fs is **not** consistent. The -canonical system is (amu, Å, fs) with energy in amu·Å²/fs² (= ``EV_PER_AMU_A2_FS2`` -eV); drive an eV/Å potential by converting at the force field +canonical system is (amu, Å, fs) with energy in amu·Å²/fs² +(= :data:`molix.units.EV_PER_AMU_A2_FS2` eV); drive an eV/Å potential by +converting at the force field (``PotentialForceField(..., energy_scale=1/EV_PER_AMU_A2_FS2)``). +Precision boundary: the force field owns its own dtype, independently of the +trajectory state's (``MD(dtype=)`` governs the state; ``MD.set_potential_dtype`` +the model). :meth:`Integrator.eval_force` casts the force field's output back to +the state dtype so the two precisions never silently promote mid-step. + Reference: Leimkuhler & Matthews, "Rational Construction of Stochastic Numerical Methods for Molecular Sampling", Appl. Math. Res. Express 2013. @@ -36,13 +51,10 @@ from torch import nn from molix.md.forcefield import ForceField -from molix.md.types import MDState - -#: Energy-unit bridge: 1 amu·Å²/fs² = 103.6426965638 eV. -EV_PER_AMU_A2_FS2 = 103.6426965638 +from molix.md.types import ForceOutput, MDState -def as_mass_col(mass: float | torch.Tensor, ref: torch.Tensor) -> torch.Tensor: +def _as_mass_col(mass: float | torch.Tensor, ref: torch.Tensor) -> torch.Tensor: """Mass reshaped to broadcast against ``(N, 3)`` in ``ref``'s dtype/device. A per-atom ``(N,)`` tensor becomes ``(N, 1)``; a scalar stays scalar. Shared @@ -58,37 +70,128 @@ def as_mass_col(mass: float | torch.Tensor, ref: torch.Tensor) -> torch.Tensor: class Integrator(nn.Module): """Abstract integrator over a :class:`~molix.md.forcefield.ForceField` component. + The contract :class:`~molix.md.runner.MDRunner` drives — a conforming + subclass implements :meth:`advance` (one eager step) and :meth:`rollout` + (the compile-friendly fixed-length loop) and inherits the rest: + :meth:`initial` seeds the state, :meth:`advance_n` chunks eager steps + between observations, and :attr:`removed_dof` states the + temperature-estimator convention. + Args: force: The force-field component supplying ``forward(pos) -> ForceOutput``. + rebuild: Whether :meth:`eval_force` asks the force field's neighbour + policy. ``None`` (default) derives it from + :attr:`~molix.md.forcefield.ForceField.rebuilds_neighbors`, so a + list-backed force field drives its list and a listless one costs + nothing. Pass ``False`` to freeze the list for a compiled rollout, + ``True`` to force the question. """ - def __init__(self, force: ForceField) -> None: + def __init__(self, force: ForceField, *, rebuild: bool | None = None) -> None: super().__init__() self.force = force + #: Whether this integrator asks the force field's neighbour policy once + #: per force evaluation. A **plain Python bool**, fixed at construction + #: and never a buffer or a tensor: dynamo specialises the branch at + #: trace time, so ``rebuild=False`` leaves the body of the guard in + #: :meth:`eval_force` dead and the step still traces to one graph. + #: Answers *whether this integrator asks*, against + #: :attr:`~molix.md.forcefield.ForceField.rebuilds_neighbors`'s + #: *whether the force field can*. + self.rebuild: bool = bool(force.rebuilds_neighbors) if rebuild is None else bool(rebuild) + + @property + def removed_dof(self) -> int: + """Degrees of freedom the temperature estimator must not count. + + ``3`` by default — a deterministic integrator conserves the (removed) + centre-of-mass momentum, leaving ``3N - 3``. Thermostatted integrators + that agitate all ``3N`` DoF (Langevin's O step includes the COM) + override this to ``0``. + """ + return 3 + + def eval_force(self, pos: torch.Tensor) -> ForceOutput: + """Evaluate the force field, casting its output to the state dtype. + + The **only** neighbour-policy seam in the engine. With + :attr:`rebuild` set, the force field is asked once here, at ``pos``, + immediately before the force call — and the *list* decides, under its + own ``skin`` / ``every`` / ``delay`` / ``check`` gate. The integrator + keeps no cadence state: no counter, no modulo, no second owner. + + Why here and nowhere else: velocity-Verlet / BAOAB evaluates ``F`` at + the *end-of-step* positions, so a refresh wired to step-start leaves + the connectivity one displacement behind the positions entering + ``F = -∇E`` — a systematic, one-signed energy leak on long NVE runs + rather than a symmetric discretisation error. + + Compile invariant: :attr:`rebuild` is a Python bool, so dynamo + specialises this branch. ``rebuild=False`` leaves the call dead and + ``torch.compile(ig.rollout, fullgraph=True)`` still traces one graph; + ``rebuild=True`` is the eager production path, where the policy runs + *between* compiled force calls and the list's fixed-capacity buffers + keep every shape static, so a compiled / CUDA-graph-captured force + field survives each rebuild. + + The force field owns its own precision (deliberately independent of the + trajectory's — see :class:`molix.md.driver.MD`); the state must not + silently promote, so energy/forces come back in ``pos``'s dtype. A + same-dtype ``.to`` is the identity, so the matched case costs nothing. + """ + if self.rebuild: + self.force.rebuild_neighbors(pos) + out = self.force(pos) + return ForceOutput(out.energy.to(pos.dtype), out.forces.to(pos.dtype)) def initial(self, pos: torch.Tensor, vel: torch.Tensor) -> MDState: """Seed an :class:`~molix.md.types.MDState`, evaluating the entry force.""" - out = self.force(pos) + out = self.eval_force(pos) return MDState(pos, vel, out.forces, out.energy) - def step(self, state: MDState, noise: torch.Tensor) -> MDState: # noqa: D102 + def advance(self, state: MDState) -> MDState: + """One eager step.""" raise NotImplementedError - def rollout(self, state: MDState, n_steps: int) -> MDState: # noqa: D102 + def advance_n(self, state: MDState, n_steps: int) -> MDState: + """Advance ``n_steps`` eagerly; subclasses may specialise the loop.""" + for _ in range(n_steps): + state = self.advance(state) + return state + + def rollout(self, state: MDState, n_steps: int) -> MDState: + """Advance ``n_steps`` and return the final state (compile-friendly).""" raise NotImplementedError + def cast_state(self, dtype: torch.dtype) -> "Integrator": + """Cast this integrator's own step-constant buffers to ``dtype``. + + Unlike ``.to(dtype)`` this does **not** recurse into the force field: + the MD-side precision and the potential's precision are independent + concerns (:class:`molix.md.driver.MD` casts the two separately). + """ + for name, buf in self.named_buffers(recurse=False): + if buf.is_floating_point(): + self._buffers[name] = buf.to(dtype) + return self + class LangevinVerletIntegrator(Integrator): """Langevin velocity-Verlet (BAOAB) over a force-field component. Args: force: Force-field component (``forward(pos) -> ForceOutput``). - dt: Timestep Δt. - gamma: Langevin friction γ (``0`` → NVE; the O step becomes the identity). + dt: Timestep Δt in fs. + gamma: Langevin friction γ in fs⁻¹ (``0`` → NVE; the O step becomes the + identity). kbt: Thermal energy k_B·T (energy units). mass: Particle mass — scalar or per-atom ``(N,)`` tensor, strictly positive. - seed: Seed for the eager noise generator (reproducible :meth:`advance` / - :meth:`run`). :meth:`rollout` uses global RNG so it stays compilable. + seed: Seed for the eager noise generator (reproducible :meth:`advance`). + :meth:`rollout` uses global RNG so it stays compilable. + rebuild: Forwarded to :class:`Integrator` — whether + :meth:`Integrator.eval_force` asks the force field's neighbour + policy. ``None`` derives it from the force field; ``False`` is the + frozen-list configuration a ``fullgraph=True`` rollout needs. Scalar parameters are immutable after construction (baked into the compiled graph). Buffers ``dt``/``c1``/``c2``/``mass_col``/``inv_mass``/``sigma`` carry @@ -104,8 +207,9 @@ def __init__( kbt: float, mass: float | torch.Tensor, seed: int = 0, + rebuild: bool | None = None, ) -> None: - super().__init__(force) + super().__init__(force, rebuild=rebuild) if isinstance(mass, torch.Tensor): if not bool((mass > 0).all()): raise ValueError("mass must be strictly positive") @@ -116,7 +220,7 @@ def __init__( c1 = math.exp(-float(gamma) * float(dt)) c2 = math.sqrt(max(0.0, 1.0 - c1 * c1)) ref = torch.zeros(()) # CPU fp32 reference for buffer construction - mass_col = as_mass_col(mass, ref) + mass_col = _as_mass_col(mass, ref) self.register_buffer("dt", torch.as_tensor(float(dt))) self.register_buffer("c1", torch.as_tensor(c1)) self.register_buffer("c2", torch.as_tensor(c2)) @@ -125,6 +229,11 @@ def __init__( self.register_buffer("sigma", math.sqrt(float(kbt)) * mass_col.rsqrt()) self._generator: torch.Generator | None = None + @property + def removed_dof(self) -> int: + """``0`` under the thermostat (γ>0 agitates all 3N DoF, COM included); ``3`` NVE.""" + return 0 if self.gamma > 0.0 else 3 + def step(self, state: MDState, noise: torch.Tensor) -> MDState: """One BAOAB step from the cached entry force and a pre-drawn ``noise``. @@ -134,19 +243,55 @@ def step(self, state: MDState, noise: torch.Tensor) -> MDState: step (one force-field evaluation per step). ``torch.compile``-able. """ half_dt = 0.5 * self.dt - vel = state.vel + half_dt * state.force * self.inv_mass # B (cached force) + vel = state.vel + half_dt * state.forces * self.inv_mass # B (cached force) pos = state.pos + half_dt * vel # A vel = self.c1 * vel + self.c2 * self.sigma * noise # O (identity at γ=0) pos = pos + half_dt * vel # A - out = self.force(pos) + out = self.eval_force(pos) vel = vel + half_dt * out.forces * self.inv_mass # B return MDState(pos, vel, out.forces, out.energy) + def step_nve(self, state: MDState) -> MDState: + """One γ=0 step: BAOAB with the identity O step elided. + + Bit-identical to ``step(state, noise)`` at ``gamma=0`` — there + ``c1=1, c2=0``, so ``v ← 1.0·v + 0.0·σ·ξ`` is the identity in floating + point too (multiply by 1.0 and add of +0.0 are exact). The two half + drifts are kept as **separate adds** to preserve that bit identity; + fusing them into one full drift would reassociate. Used by + :meth:`advance_n` so long NVE runs skip the per-step noise draw; the + compiled :meth:`rollout` keeps the branch-free :meth:`step` per the + md-component-engine spec. + """ + half_dt = 0.5 * self.dt + vel = state.vel + half_dt * state.forces * self.inv_mass # B (cached force) + pos = state.pos + half_dt * vel # A + pos = pos + half_dt * vel # A (O step elided: identity at γ=0) + out = self.eval_force(pos) + vel = vel + half_dt * out.forces * self.inv_mass # B + return MDState(pos, vel, out.forces, out.energy) + + def advance_n(self, state: MDState, n_steps: int) -> MDState: + """Advance ``n_steps`` eagerly with no per-step host work. + + The γ selection is a construction-time Python branch out here in the + eager driver — :meth:`step` itself stays branch-free (spec invariant). + At γ=0 this also skips the per-step ``randn`` whose contribution the + O step would multiply by ``c2=0`` anyway. + """ + if self.gamma == 0.0: + for _ in range(n_steps): + state = self.step_nve(state) + return state + for _ in range(n_steps): + state = self.step(state, self.draw_noise(state.vel)) + return state + def draw_noise(self, ref: torch.Tensor) -> torch.Tensor: """Reproducible O-step noise ``(N, 3)`` from a seeded generator (eager). - Used by :meth:`advance` / :meth:`run`; kept out of :meth:`rollout` so the - compiled path has no ``Generator`` object in the graph. + Used by :meth:`advance` / :meth:`advance_n`; kept out of :meth:`rollout` + so the compiled path has no ``Generator`` object in the graph. """ if self._generator is None or self._generator.device != ref.device: self._generator = torch.Generator(device=ref.device).manual_seed(self._seed) @@ -166,34 +311,3 @@ def rollout(self, state: MDState, n_steps: int) -> MDState: for _ in range(n_steps): state = self.step(state, torch.randn_like(state.vel)) return state - - def run( - self, pos: torch.Tensor, vel: torch.Tensor, n_steps: int, *, stride: int = 1 - ) -> dict[str, torch.Tensor]: - """Eager trajectory: record every ``stride``-th frame's pos/vel/energy. - - History is detached and moved to CPU as recorded, so device memory does - not grow with ``n_steps`` (host memory grows O(T·N/stride); raise - ``stride`` or use :class:`molix.md.TrajectoryHook` for long runs). - Reproducible via the seeded generator. - - Returns: - ``pos`` / ``vel`` ``(⌈n_steps/stride⌉, N, 3)`` and ``energy`` - ``(⌈n_steps/stride⌉,)`` (detached, on CPU). - """ - stride = max(1, int(stride)) - pos_hist: list[torch.Tensor] = [] - vel_hist: list[torch.Tensor] = [] - energy_hist: list[torch.Tensor] = [] - state = self.initial(pos, vel) - for i in range(n_steps): - state = self.advance(state) - if (i + 1) % stride == 0: - pos_hist.append(state.pos.detach().to("cpu")) - vel_hist.append(state.vel.detach().to("cpu")) - energy_hist.append(state.energy.detach().reshape(()).to("cpu")) - return { - "pos": torch.stack(pos_hist), - "vel": torch.stack(vel_hist), - "energy": torch.stack(energy_hist), - } diff --git a/src/molix/md/neighbors.py b/src/molix/md/neighbors.py new file mode 100644 index 0000000..d4c0f3b --- /dev/null +++ b/src/molix/md/neighbors.py @@ -0,0 +1,1174 @@ +"""Rebuilding, fixed-capacity neighbour list for production MD. + +A frozen neighbour list is only valid while no atom moves far enough to change +its neighbour set — tens of steps for liquid water, not the millions a +production trajectory needs. This module rebuilds it under the LAMMPS +``neigh_modify every/delay/check`` policy (see below) while keeping **every +tensor shape constant**, which is what lets the force evaluation stay inside a +CUDA graph across the whole run. + +The trick is the *capacity*: edge tensors are allocated once at +``capacity = ceil(factor * E_initial)`` and only their contents change. Unused +rows carry a **dead edge** — source and target both atom 0, displaced by a shift +longer than the cutoff. Such an edge contributes exactly zero: + +* ``r = |shift| > r_cut`` so the polynomial cutoff envelope is 0, and the radial + features it multiplies are 0, so the message and the learned density are 0; +* pair-repulsion envelopes cut off at the pair's covalent radii, well inside + ``r_cut``, so they are 0 too; +* ``pos[0] - pos[0]`` cancels exactly, so no spurious force reaches atom 0. + +Periodicity is handled by the minimum-image convention of the compiled +neighbour kernel, whose sequential reduction is complete only up to half the +smallest **perpendicular width** of the cell: ``w_i = V / ||a_j x a_k||`` in +Angstrom, for cell vectors ``a_1, a_2, a_3`` (the rows of ``cell``) and volume +``V = |det(cell)|``. So the *build* radius must not exceed ``min_i w_i / 2`` — +:class:`NeighborList` refuses to construct otherwise rather than +silently dropping pairs that are inside the cutoff. For an orthorhombic cell +``w_i = ||a_i||``, i.e. the familiar half-shortest-cell-vector bound. + +The Verlet skin +--------------- + +``cutoff`` is the **interaction** cutoff ``r_cut`` every consumer means. The +list is built at the enlarged radius + +``r_build = cutoff + skin`` (Angstrom, :attr:`NeighborList.r_build`) + +so that pairs which walk *into* ``r_cut`` between rebuilds are already in the +buffers. The model's own cutoff envelope masks the skin-region pairs to zero +(``LennardJonesCutForceField`` compares against ``cutoff_sq``), so the skin +costs edges — the live count grows as ``(1 + skin/r_cut)^3`` — but changes no +energy. Both the capacity and the kernel's ``max_num_pairs`` are therefore +sized from ``r_build``, never from ``cutoff``. + +**Half-skin completeness criterion.** For atoms *i*, *j* let +``d_i = ||x_i(t) - x_i(t0)||`` be the displacement since the last build at +``t0``. The triangle inequality gives ``r_ij(t) >= r_ij(t0) - d_i - d_j``, so +any pair inside ``r_cut`` at time *t* was inside ``r_cut + d_i + d_j`` at +``t0``. A list built at ``r_build = r_cut + s`` is therefore still **complete** +while ``d_(1) + d_(2) <= s`` for the two largest displacements. Which two atoms +those are is unknown, so the conservative sufficient condition — and the one +LAMMPS tests — is the **half**-skin bound + +``max_i d_i <= s / 2`` (worst case ``d_(1) = d_(2) = s/2``; hence *half*) + +with a **strict** ``>`` triggering the rebuild: exactly ``s/2`` still satisfies +the proof, so rebuilding there would be wasted work. + +**Raw displacements, unwrapped positions (load-bearing invariant).** The +displacement test uses the *raw* difference ``x - x_hold``, never a minimum +image — min-imaging it would clamp a genuine ``> L/2`` excursion and so +*suppress* the very rebuild it is meant to force. That is correct only while +positions drift unwrapped, which is what this repo's MD does (nothing in +``molix.md.integrators`` / ``molix.md.runner`` wraps; LAMMPS wraps only on +reneighbour steps, before ``xhold`` is stored). :meth:`NeighborList.update` +promotes it to a checked invariant: a displacement at or beyond +``min_i w_i / 2`` raises :class:`RuntimeError`, which catches mid-run wrapping, +a changed cell and a blown-up trajectory alike. Do not wrap positions mid-run. + +**``ndanger`` — the correctness alarm.** A rebuild that fires at the *first* +permitted opportunity, ``ago == max(every, delay)``, may already have been +overdue on an earlier, non-permitted step, so pairs may have been missed; +:attr:`NeighborList.ndanger` counts those (LAMMPS "Dangerous builds"). A run +that reports ``ndanger > 0`` should be rerun with a larger ``skin`` or a +tighter gate. **Caveat at the default gate** ``every=1, delay=0``: +``max(every, delay) == 1``, so *every* rebuild lands on the first permitted +opportunity and ``ndanger`` simply counts rebuilds. There the counter carries +no information — read it only for a coarser gate or a nonzero skin. + +**What ``check=False`` and a coarse ``delay`` cost.** A pair inside ``r_cut`` +but absent from the list contributes exactly zero, and the next build inserts +it discontinuously: the total energy takes an ``O(1)``, one-signed injection +per missed event — not the ``O(dt^2)`` error of a discretisation artefact. +Repeated events integrate into a systematic NVE energy leak. ``check=False`` +(and any ``delay`` long enough to skip a needed rebuild) buys speed by +accepting that leak; it is never a free optimisation. ``check=False`` +additionally disables the unwrapped-positions guard, which lives inside the +displacement branch. + +**LAMMPS parity.** The gate, the strict comparison, the raw difference and the +``ndanger`` threshold are ``Neighbor::decide`` / ``Neighbor::check_distance`` +verbatim, and the constructor mirrors ``Neighbor::init`` in rejecting a +``delay`` that is not a multiple of ``every`` (which would make the danger +threshold an ``ago`` the gate never permits, silently killing the alarm). + +The binned build +---------------- + +``bin=None`` (the default) hands the whole system to the compiled O(N^2) pair +kernel, which enumerates ``N(N-1)/2`` candidates per rebuild. Passing a float +switches to a pure-torch **cell list**: atoms are sorted into a periodic grid +of bins derived once from the cell and ``r_build``, and only a fixed stencil of +neighbouring bins is searched, so the build is linear in ``N`` at fixed density +with no per-atom Python loop and no device-specific code +(:meth:`NeighborList._configure_bins` derives the grid, +:meth:`NeighborList._build_binned` runs the search). Both backends return the +same triple into the same ``_write``, so the capacity, the dead-edge padding +and every consumer are unaffected by which one ran. ``bin`` is a **cost** knob: +the edge set is identical, and the kernel path is the equivalence oracle the +binned one is tested against. + +**Stencil completeness — why a bounded search is exact.** Bins are cubes in +*fractional* space: atom A sits in bin ``p_i = floor(s_i n_i)`` along axis +``i``. If two atoms are ``m`` bins apart along axis ``i`` (minimal modular +difference), then ``s_B >= (p + m)/n_i`` while ``s_A < (p + 1)/n_i``, so +``|Ds_i| > (m - 1)/n_i`` **strictly**. The displacement's component along the +axis normal is ``|Ds_i| w_i``, and ``||d|| >= |d . n_i| = |Ds_i| w_i``, so with +the *effective* perpendicular bin thickness ``b_i = w_i / n_i`` (Angstrom) + +``||d|| > (m - 1) * b_i`` + +A stencil half-width ``k_i`` with ``k_i * b_i >= r_build`` is therefore +**complete**: any pair more than ``k_i`` bins apart on some axis has +``||d|| > k_i b_i >= r_build`` and is excluded by the ``r <= r_build`` filter +anyway. The strict inequality is what makes the textbook statement exact at the +boundary — ``b_i = r_build`` gives ``k_i = 1``, i.e. the 27-cell 3x3x3 stencil, +with no epsilon fudge. ``_configure_bins`` re-checks the relation per axis +instead of trusting the arithmetic that produced it: an incomplete stencil is a +silently short neighbour list, which is a wrong energy that never raises. + +**Bin size.** The candidate volume for ``b = r_build / k`` is +``(2 + 1/k)^3 r_build^3`` — ``27 r^3`` at ``k = 1``, ``15.6 r^3`` at ``k = 2``, +against the ``4.19 r^3`` sphere actually needed — decreasing in ``k`` while the +sorting and gather cost grows with the bin count. LAMMPS settles this at +``k = 2`` (``src/nbin_standard.cpp``, ``binsize_optimal = 0.5 * cutneighmax``), +which is what ``bin=0.0`` requests here. + +**Minimum image by fractional rounding is exact inside the half-width guard.** +Write a periodic displacement as ``d = sum_i f_i a_i``. Its component along the +axis-``i`` normal is ``|f_i| w_i = |d . n_i| <= ||d||``. So whenever +``||d|| <= r_build <= min_i w_i / 2 <= w_i / 2`` — exactly the guard above — +``|f_i| <= 1/2`` on every axis: **any** in-range image is already the one +``f <- f - round(f)`` selects. Fractional rounding and the kernel's sequential +diagonal reduction therefore return the same, unique minimum image everywhere +the guard admits, which is what makes edge-set equality between the two +backends a theorem rather than a coincidence. (Ties at ``|f_i| = 1/2`` are +reachable only when ``r_build = min_i w_i / 2`` exactly *and* a pair sits +exactly on the boundary, where ``torch.round``'s half-to-even and C's +half-away-from-zero can differ: measure zero, and outside every fixture.) + +**Filter parity.** The binned path accepts a pair iff ``0 < r <= r_build``, +which is the compiled backends' filter verbatim (``(distances <= cutoff) & +(distances > 0)`` in the C++ kernel; ``distance2 > cutoff2 || distance2 == 0`` +dropped in the CUDA one) — the same *closed* upper bound and the same ``r > 0`` +rejection, so coincident atoms and pairs separated by exactly one lattice +vector (whose minimum image is the zero vector) are dropped identically. + +**Edge order is not part of the contract.** The binned path emits edges in +bin-sorted order, the kernel in upper-triangle index order. What binds is the +*set* of ``(source, target, shift)`` triples plus the edge count: consumers +reduce with order-independent scatter / ``index_add_`` and read the shift +buffer positionally alongside ``edge_index``, never by index-order assumption. + +**Small cells degenerate gracefully.** When ``2 k_i + 1 >= n_i`` on every axis +the stencil's residue sets cover the whole grid and the search enumerates all +pairs — correct, just not faster (and a single bin per axis, the ``bin`` wider +than the cell case, is exactly that). The O(N) win appears once +``n_i > 2 k_i + 1``, i.e. cells wider than about ``5 r_build / 2`` per axis at +the automatic bin size, so a small-cell timing is not a regression. + +References: + LAMMPS ``neigh_modify`` documentation — + https://docs.lammps.org/neigh_modify.html — and ``lammps/lammps`` develop + ``src/neighbor.cpp`` (``Neighbor::decide``, ``Neighbor::check_distance``, + ``Neighbor::init``) / ``src/verlet.cpp``. The binning policy behind + ``bin=0.0`` is ``src/nbin_standard.cpp`` + (``binsize_optimal = 0.5 * cutneighmax``) and + https://docs.lammps.org/neighbor.html; the paper of record is + A. P. Thompson et al., *Comput. Phys. Commun.* **271**, 108171 (2022), + https://doi.org/10.1016/j.cpc.2021.108171. + + K. Nordlund, *Introduction to molecular dynamics simulations*, lecture 3, + https://www.mv.helsinki.fi/home/knordlun/moldyn/lecture03.pdf — the open, + directly re-verified source for the two-atom criterion above. + + M. P. Allen & D. J. Tildesley, *Computer Simulation of Liquids*, 2nd ed., + Oxford University Press (2017), + https://doi.org/10.1093/oso/9780198803195.001.0001 — cell (link-cell) + lists, the 27-cell stencil and minimum-image validity. + + L. Verlet, *Phys. Rev.* **159**, 98 (1967), + https://doi.org/10.1103/PhysRev.159.98 — the original neighbour list. + + B. Quentrec & C. Brot, *J. Comput. Phys.* **13**, 430 (1973), + https://doi.org/10.1016/0021-9991(73)90046-6 — the original cell method, + and the skin refinement. + + Caveat: the Verlet 1967 and Quentrec & Brot 1973 texts are paywalled and + were **not** re-verified here; they are cited for attribution only. Every + equation above is verified against the Nordlund notes and the LAMMPS + source, and the stencil-completeness and ``|f_i| <= 1/2`` relations are + derived in line above rather than taken from a text. + +Owning ``edges``: the bind surface +---------------------------------- + +Two entry points, two audiences, one owner. The MD hot path drives the list +with raw ``(N, 3)`` position tensors in Angstrom (:meth:`NeighborList.rebuild` +forces a build, :meth:`NeighborList.update` applies the policy); a TensorDict +caller drives it with the batch itself (:meth:`NeighborList.build` builds *and* +binds, :meth:`NeighborList.update` again — one method with an ``isinstance`` +dispatch at the top, never an ``update_td`` / ``update_pos`` pair of twins). + +**The list owns ``edges`` once bound.** :meth:`NeighborList.build` writes +``batch["edges"]`` as a ``TensorDict`` holding the live ``edge_index`` / +``shifts`` buffers **by reference**, replacing whatever was there — including a +precomputed ``edge_diff`` / ``edge_dist`` pair, which a potential would +otherwise consume straight through as a value and so freeze the PES. Because +every rebuild is in place, that tie survives all later ``rebuild`` / ``update`` +calls with no re-binding: the potential simply sees the current neighbour set. +It does **not** survive :meth:`NeighborList.to`, which rebinds the buffers to +new tensors — the owner re-binds (see ``PeriodicPotentialForceField._apply``). + +Two ``NeighborList`` classes, deliberately +----------------------------------------- + +The repository holds two classes named ``NeighborList``, in different layers, +and the collision is intentional rather than an accident awaiting cleanup: + +* :class:`molix.md.neighbors.NeighborList` (this one) — the **MD engine** + symbol: a *stateful, fixed-capacity buffer owner* holding ``edge_index + (capacity, 2)``, ``shifts (capacity, 3)``, ``num_edges`` and + ``rebuild_count``, rebuilt in place so shapes never change and the force path + stays CUDA-graph capturable. One instance per run, held by a + :class:`~molix.md.forcefield.ForceField`. +* :class:`molix.data.tasks.neighbor.NeighborList` — the **data-pipeline** + symbol: a *stateless* ``SampleTask`` mapping one flat sample dict to + ``edge_index`` / ``edge_diff`` / ``edge_dist`` and contributing to ``task_id`` + for cache keying. Constructed once per pipeline definition, no per-call state. + +They are not two variants of one concept — a per-run mutable buffer with an +overflow policy versus a pure pipeline transform — and each is the shortest +natural name in its own layer, so neither gives up the bare name. Because this +module *consumes* the pipeline task (see the import below), and a bare +``from ... import NeighborList`` here would be rebound by the class definition +further down — making the constructor call itself recursively — the import is +aliased to ``NeighborListTask``. +""" + +from __future__ import annotations + +import math +from typing import Protocol, runtime_checkable + +import torch +from tensordict import TensorDict, TensorDictBase + +# The one owner of kernel-output normalisation (pbc handling, NaN-padding +# strip, symmetry expansion, edge-sign convention); reimplementing that here +# against the raw ``molix.F.locality`` kernel would fork it. Aliased because +# the MD buffer owner defined below is *also* named ``NeighborList`` and would +# otherwise shadow its own dependency — the class would then call itself. +from molix.data.tasks.neighbor import NeighborList as NeighborListTask +from molix.units import DEAD_EDGE_CUTOFF_FACTOR + +#: Largest stencil half-width ``k_i`` (in bins) the binned build accepts. At the +#: cap the Python loop runs over ``17^3 = 4913`` bin offsets; beyond it an +#: absurdly small explicit ``bin`` stops being a fine grid and becomes a hang. +_MAX_STENCIL_HALF_WIDTH = 8 + + +def _perpendicular_widths(cell: torch.Tensor) -> torch.Tensor: + """Distances between the three pairs of opposite faces of a periodic cell. + + For cell vectors ``a_1, a_2, a_3`` — the **rows** of ``cell`` — the width + perpendicular to the face spanned by ``a_j`` and ``a_k`` is + ``w_i = V / ||a_j x a_k||`` with ``V = |det(cell)|``. For an orthorhombic + cell ``||a_j x a_k|| = ||a_j||*||a_k||`` and ``V = ||a_1||*||a_2||*||a_3||``, + hence ``w_i = ||a_i||``. + + Two callers need different reductions of the same three numbers: the + minimum-image guard bounds ``r_build`` by ``min_i w_i / 2``, while the + binned build sizes its grid **per axis** so that a requested bin thickness + means the same thing on a sheared cell as on a cube. + + Evaluated in ``float64`` so a ``float32`` cell cannot jitter an + accept/reject decision taken right at the bound. + + Args: + cell: Cell vectors ``(3, 3)`` in Angstrom, one vector per row. + + Returns: + The perpendicular widths ``(w_1, w_2, w_3)`` as a ``(3,)`` ``float64`` + tensor in Angstrom, in cell-row order. + + Raises: + ValueError: If ``cell`` is not ``(3, 3)``, or is degenerate (volume + zero or non-finite, or two rows collinear). A degenerate cell has + no finite width, and returning ``nan`` would make every ``>`` + comparison against the bound silently succeed. + """ + if tuple(cell.shape) != (3, 3): + raise ValueError(f"cell must have shape (3, 3), got {tuple(cell.shape)}") + vectors = cell.detach().to(torch.float64) + volume = float(torch.linalg.det(vectors).abs()) + areas = torch.linalg.norm( + torch.stack( + ( + torch.linalg.cross(vectors[1], vectors[2]), + torch.linalg.cross(vectors[2], vectors[0]), + torch.linalg.cross(vectors[0], vectors[1]), + ) + ), + dim=-1, + ) + max_area = float(areas.max()) + if not math.isfinite(volume) or volume <= 0.0 or not math.isfinite(max_area) or max_area <= 0.0: + raise ValueError( + f"degenerate cell: volume {volume} A^3, largest face area {max_area} A^2; " + "a cell with no interior has no perpendicular width to bound the cutoff by." + ) + # A positive volume forces every face area positive (a zero area means two + # rows are collinear, which collapses the determinant), so this cannot divide + # by zero once the guard above has passed. + return volume / areas + + +@runtime_checkable +class NeighborStrategy(Protocol): + """Contract a force field expects from a rebuildable neighbour list. + + Any strategy (minimum-image, Verlet-skin, cell list, …) is usable by + :class:`~molix.md.forcefield.PeriodicPotentialForceField` / + :class:`~molix.md.forcefield.CallableForceField` as long as it exposes + these members with fixed-shape, in-place-rebuilt buffers. + + Attributes: + edge_index: Edge buffer ``(capacity, 2)`` — ``[:, 0]`` source, + ``[:, 1]`` target, per the repo edge convention. + shifts: Periodic shift vectors ``(capacity, 3)``. + num_edges: Live edges occupy ``[0, num_edges)``; the rest are dead. + capacity: Fixed buffer length. + cutoff: Interaction cutoff ``r_cut`` in Angstrom — the radius the list + guarantees *complete* between rebuilds, and the horizon a force + field's own cutoff must not exceed. A strategy may build at a + larger radius (see ``skin``); that is its business, not the + consumer's. + skin: Verlet skin in Angstrom, ``0.0`` for a strategy that rebuilds at + the bare cutoff. + """ + + edge_index: torch.Tensor + shifts: torch.Tensor + num_edges: int + capacity: int + cutoff: float + skin: float + + def rebuild(self, positions: torch.Tensor) -> None: + """Recompute the neighbour list at ``positions``, in place.""" + ... + + def build(self, batch: TensorDict) -> TensorDict: + """Rebuild at ``batch["atoms", "pos"]`` and bind the buffers into ``batch``. + + Returns the same batch object, with ``batch["edges"]`` holding the + strategy's live buffers by reference — the strategy owns that namespace + from here on. ``PeriodicPotentialForceField`` calls this through this + annotation, at construction and after every ``.to()``. + """ + ... + + def update(self, positions: TensorDict | torch.Tensor) -> bool: + """Rebuild at ``positions`` if the strategy's policy says so. + + Called once per force evaluation, with either the raw ``(N, 3)`` + positions of the MD hot path or the batch carrying them. Returns + whether a rebuild happened. + """ + ... + + def to( + self, + device: torch.device | str | torch.dtype | None = None, + dtype: torch.dtype | None = None, + ) -> "NeighborStrategy": + """Move / cast the buffers, mirroring ``Tensor.to`` semantics.""" + ... + + +class NeighborList: + """Minimum-image Verlet-skin neighbour list in fixed-capacity buffers. + + Not an ``nn.Module``: it owns plain buffers and a rebuild policy, and is held + by a :class:`~molix.md.forcefield.ForceField`. Keeping it out of the module + tree also keeps it out of ``state_dict``, where a per-run neighbour list has + no business — the owning force field forwards device/dtype changes through + :meth:`to` instead (see ``PeriodicPotentialForceField._apply``). + + Two entry points drive it: :meth:`rebuild` forces a build unconditionally, + :meth:`update` applies the ``every`` / ``delay`` / ``check`` policy. A + TensorDict caller uses :meth:`build` (force a build *and* bind the buffers + into the batch) and the same :meth:`update`, which takes either input type. + See the module docstring for the half-skin criterion, the + unwrapped-positions invariant and what ``check=False`` costs. + + Args: + cell: Cell vectors ``(3, 3)`` in Angstrom, one vector per row. + cutoff: **Interaction** cutoff ``r_cut`` in Angstrom — what every + consumer means by "cutoff", and the radius the list stays complete + to between rebuilds. The build radius is :attr:`r_build`. + positions: Initial positions ``(N, 3)`` in Angstrom; the constructor + builds the list at them (not counted in :attr:`rebuild_count`) and + sizes the capacity from the result. + skin: Verlet skin ``s`` in Angstrom. The list is built at + ``r_build = cutoff + skin`` and is provably complete out to + ``cutoff`` while no atom has moved more than ``s/2`` since the last + build. ``0.0`` reproduces the pre-skin rebuild-on-any-motion + behaviour. + every: Attempt a rebuild only when the number of steps since the last + build is a multiple of this (steps). + delay: Attempt no rebuild until at least this many steps have passed + since the last build (steps). Must be a multiple of ``every`` + (LAMMPS ``Neighbor::init`` parity); ``0`` is always legal. + check: Rebuild only when the maximum displacement since the last build + exceeds ``skin/2``. ``False`` rebuilds on cadence alone — cheaper, + but it drops both the completeness criterion and the + unwrapped-positions guard (module docstring). + capacity_factor: Buffer capacity as a multiple of the initial edge + count at ``r_build``. Density fluctuations grow the edge count + during a run; overflow raises rather than truncating. **Caveat for + crystalline starts:** a lattice initial configuration systematically + *under*-estimates the live-edge count a warm run reaches, because a + coordination shell sitting just outside ``r_build`` at ``t = 0`` is + pulled in by thermal motion (the LJ-lattice test harness needed + ``2.5`` where the default would have allocated 519 rows against 600 + live edges). Raise it when starting from a lattice. + bin: Selects the **build backend** (never the physics — see the module + docstring). ``None`` (default) keeps the compiled O(N^2) pair + kernel. A float switches to the pure-torch binned (cell-list) + build and is the *requested* perpendicular bin thickness in + Angstrom: ``0.0`` asks for the automatic ``r_build / 2`` (LAMMPS + ``nbin_standard`` ``binsize_optimal``), a positive value asks for + that thickness. The effective thickness is ``w_i / n_i`` with + ``n_i = max(1, floor(w_i / bin))``, so it is never *smaller* than + requested and is reported per axis through :attr:`n_bins`. The + binned path wins on cells wider than about ``5 * r_build / 2`` per + axis and merely degenerates to an all-pairs search below that. + Measured 2026-08-09 (torch 2.12.1+cpu, x86_64, N=4096, 48 A cube, + r_build 5.0, n_bins (19,19,19)): binned 0.037 s vs kernel 0.450 s + per rebuild at ``OMP_NUM_THREADS=4`` — but 7.2 s vs 1.0 s at the + node default of 48 threads, where per-op OpenMP region overhead + dominates the 125-offset loop. The thread count is part of any + such number; no timing threshold is asserted anywhere. + device: Device for the buffers; defaults to ``positions``'. + + Raises: + ValueError: If ``cell`` is not ``(3, 3)`` or is degenerate; if ``skin``, + ``every`` or ``delay`` is out of domain or ``delay`` is not a + multiple of ``every``; if ``r_build`` exceeds half the smallest + perpendicular cell width ``w_i = V / ||a_j x a_k||`` (Angstrom; + ``w_i = ||a_i||`` for an orthorhombic cell), beyond which the + minimum-image reduction silently drops pairs that lie inside the + cutoff; if ``skin`` reaches the dead-edge padding radius + ``(DEAD_EDGE_CUTOFF_FACTOR - 1) * cutoff``; or if ``bin`` is + negative or so small that the stencil half-width exceeds + ``_MAX_STENCIL_HALF_WIDTH`` bins. + """ + + def __init__( + self, + *, + cell: torch.Tensor, + cutoff: float, + positions: torch.Tensor, + skin: float = 0.0, + every: int = 1, + delay: int = 0, + check: bool = True, + capacity_factor: float = 1.35, + bin: float | None = None, + device: torch.device | None = None, + ) -> None: + cutoff_f, skin_f = float(cutoff), float(skin) + if skin_f < 0.0: + raise ValueError(f"skin must be >= 0 A, got {skin_f} A") + every_i, delay_i = int(every), int(delay) + if every_i != every or every_i < 1: + raise ValueError(f"every must be an integer >= 1 step, got {every!r}") + if delay_i != delay or delay_i < 0: + raise ValueError(f"delay must be an integer >= 0 steps, got {delay!r}") + if delay_i % every_i: + raise ValueError( + f"delay {delay_i} steps must be a multiple of every {every_i} steps " + f"({delay_i} % {every_i} = {delay_i % every_i}); LAMMPS Neighbor::init " + f"rejects the same pair, because the danger threshold max(every, delay) = " + f"{max(every_i, delay_i)} would be an ago the gate never permits, leaving " + "ndanger silently dead." + ) + + min_width = float(_perpendicular_widths(cell).min()) + half_width = 0.5 * min_width + r_build = cutoff_f + skin_f + if r_build > half_width: + raise ValueError( + f"r_build {r_build} A (cutoff {cutoff_f} A + skin {skin_f} A) exceeds half the " + f"minimum perpendicular cell width " + f"({half_width:.3f} A; widths from V/||a_j x a_k||, minimum {min_width:.3f} A); " + "the kernel's sequential minimum-image reduction would silently drop pairs " + "inside the cutoff. Use a larger cell, a shorter cutoff, or a thinner skin." + ) + dead_edge_headroom = (DEAD_EDGE_CUTOFF_FACTOR - 1.0) * cutoff_f + if skin_f >= dead_edge_headroom: + raise ValueError( + f"skin {skin_f} A reaches the dead padding edges: they sit at " + f"DEAD_EDGE_CUTOFF_FACTOR * cutoff = {DEAD_EDGE_CUTOFF_FACTOR} * {cutoff_f} = " + f"{DEAD_EDGE_CUTOFF_FACTOR * cutoff_f} A, so the skin must stay below " + f"(DEAD_EDGE_CUTOFF_FACTOR - 1) * cutoff = {dead_edge_headroom} A or a dead " + "edge falls inside r_build and is built as a real pair." + ) + + self.cutoff = cutoff_f + self.skin = skin_f + self.every = every_i + self.delay = delay_i + self.check = bool(check) + self._r_build = r_build + self.capacity_factor = float(capacity_factor) + self._device = device if device is not None else positions.device + self._dtype = positions.dtype + self.cell = cell.to(device=self._device, dtype=self._dtype) + n_atoms = int(positions.shape[0]) + # Built at r_build, not cutoff — and with an explicit half-pair bound, so + # the task's stale 512 default can never truncate the enlarged radius. + self._nl = NeighborListTask( + cutoff=self._r_build, + max_num_pairs=max(1, n_atoms * (n_atoms - 1) // 2), + pbc=True, + symmetry=True, + ) + #: Requested perpendicular bin thickness in Angstrom, or ``None`` for + #: the compiled O(N^2) backend. A cost knob, never a physics knob. + self.bin = None if bin is None else float(bin) + #: Bins per cell axis ``(n_1, n_2, n_3)``, ``None`` when ``bin is None`` + #: and no grid was derived. The only public window onto the stencil. + self.n_bins: tuple[int, int, int] | None = None + if self.bin is not None: + self._configure_bins(self.bin) + + # Through the dispatch, so the capacity is sized by the backend that + # will keep refilling the buffers (the two agree, by the equivalence). + source, target, shifts = self._build_pairs(positions) + self.capacity = max(1, int(math.ceil(self.capacity_factor * source.numel()))) + self.edge_index = torch.zeros(self.capacity, 2, dtype=torch.long, device=self._device) + self.shifts = torch.zeros(self.capacity, 3, dtype=self._dtype, device=self._device) + #: Edges live in ``[0, num_edges)``; the rest are dead. Kept for + #: diagnostics — the model needs no mask, dead edges self-annihilate. + self.num_edges = 0 + #: Rebuilds since construction; the initial build is not one of them. + self.rebuild_count = 0 + #: Steps since the last build (LAMMPS ``ago``); 0 on a fresh list. + self.ago = 0 + #: Rebuilds that fired at the first permitted opportunity and may + #: therefore have come too late (LAMMPS "Dangerous builds"). + self.ndanger = 0 + self._half_skin_sq = (0.5 * skin_f) ** 2 + self._wrap_guard_sq = half_width**2 + self._danger_ago = max(every_i, delay_i) # LAMMPS neighbor.cpp:2488, verbatim + # Reference configuration of the last build, differenced raw (never + # minimum-imaged) by ``update``. In place from here on: no per-rebuild + # allocation, no shape churn. + self._x_hold = torch.zeros(n_atoms, 3, dtype=self._dtype, device=self._device) + self._write(source, target, shifts) + self._hold(positions) + + @property + def r_build(self) -> float: + """Build radius ``cutoff + skin`` in Angstrom (derived, never settable). + + The capacity, the kernel's radius and the half-width guard were all + sized from this at construction, so a writable ``r_build`` could only + drift out of step with them. + """ + return self._r_build + + def _compute(self, positions: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Run the compiled kernel; returns ``(source, target, shifts)``.""" + pos = positions.detach() + graph = self._nl.execute({"pos": pos, "cell": self.cell.to(pos.dtype)}) + edge_index = graph["edge_index"] # (E, 2) — canonical layout throughout + source, target = edge_index[:, 0], edge_index[:, 1] + # The kernel returns minimum-image displacements; the model recomputes + # pos[target]-pos[source] itself, so hand it the periodic remainder. + shifts = graph["edge_diff"] - (pos[target] - pos[source]) + return source, target, shifts + + def _configure_bins(self, requested: float) -> None: + """Derive the bin grid and the search stencil from the cell and ``r_build``. + + Called once from ``__init__`` when ``bin`` is not ``None``. The grid + depends only on the (fixed) cell and the (derived) build radius, so + nothing here runs again per rebuild; :meth:`to` only carries the two + tensors it produces to their new device / dtype. + + Per axis ``i``, with the perpendicular widths ``w_i`` in Angstrom + (:func:`_perpendicular_widths`) and all counts dimensionless: + + 1. requested thickness ``b = requested`` in Angstrom, or ``r_build / 2`` + when ``requested == 0.0`` (LAMMPS ``nbin_standard`` + ``binsize_optimal``); + 2. ``n_i = max(1, floor(w_i / b))`` — sized on the **perpendicular** + width, not the row norm, so a requested thickness means the same + thing on a sheared cell as on a cube; + 3. effective thickness ``b_i = w_i / n_i`` in Angstrom, never below the + request except where the ``max(1, ...)`` clamp caught a cell thinner + than one requested bin; + 4. half-width ``k_i = ceil(r_build / b_i)`` bins, bumped while + ``k_i * b_i < r_build`` — a float-exactness guard that fires at most + once, and the completeness relation the whole search rests on (module + docstring); + 5. offsets ``o_i = unique(arange(-k_i, k_i + 1) mod n_i)``. The + **distinct residues** are load-bearing: as soon as ``2 k_i + 1 > n_i`` + the raw stencil wraps onto the same bin twice, and every candidate + pair in it would be emitted twice. + + Sets :attr:`n_bins` (the public diagnostic), ``self._stencil`` — the + ``(S, 3)`` Cartesian product of the three residue sets, ``S <= (2 * + _MAX_STENCIL_HALF_WIDTH + 1) ** 3`` — and ``self._inv_cell``, the cached + ``cell^-1`` the build maps positions to fractional coordinates with. + + Args: + requested: Requested perpendicular bin thickness in Angstrom; + ``0.0`` asks for the automatic ``r_build / 2``. + + Raises: + ValueError: If ``requested`` is negative (``0.0`` is how one asks + for the automatic size); if any ``k_i`` exceeds + ``_MAX_STENCIL_HALF_WIDTH``, i.e. the bin is so far below + ``r_build`` that the stencil loop stops being a search and + becomes a hang; or — a fail-loud tripwire on step 4 rather than + a user knob — if the completeness relation ``k_i * b_i >= + r_build`` fails on any axis, which would be a silently short + neighbour list. + """ + if requested < 0.0: + raise ValueError( + f"bin must be >= 0 A, got {requested} A: a negative bin thickness has no " + f"meaning. bin=0.0 selects the automatic r_build / 2 = {0.5 * self._r_build} A " + "size (LAMMPS nbin_standard), any positive value is an explicit requested " + "perpendicular bin thickness, and bin=None keeps the compiled O(N^2) backend." + ) + thickness = requested if requested > 0.0 else 0.5 * self._r_build + widths = [float(width) for width in _perpendicular_widths(self.cell)] + counts = [max(1, int(math.floor(width / thickness))) for width in widths] + effective = [width / count for width, count in zip(widths, counts, strict=True)] + halves: list[int] = [] + for size in effective: + half = int(math.ceil(self._r_build / size)) + while half * size < self._r_build: # float-exactness bump; fires at most once + half += 1 + halves.append(half) + + if max(halves) > _MAX_STENCIL_HALF_WIDTH: + raise ValueError( + f"bin {requested} A gives effective bin thicknesses " + f"{[round(size, 6) for size in effective]} A, so the stencil half-widths are " + f"{halves} bins — above the cap of {_MAX_STENCIL_HALF_WIDTH}. The build would " + f"loop over {math.prod(2 * half + 1 for half in halves)} bin offsets per " + f"rebuild, against the {(2 * _MAX_STENCIL_HALF_WIDTH + 1) ** 3} the cap allows. " + f"Pass a larger bin, or bin=0.0 for the automatic " + f"r_build / 2 = {0.5 * self._r_build} A size." + ) + for count, size, half in zip(counts, effective, halves, strict=True): + if half * size < self._r_build: + raise ValueError( + f"incomplete stencil from bin {requested} A: an axis with n_i = {count} " + f"bins of b_i = {size} A reaches only k_i * b_i = {half * size} A at " + f"half-width k_i = {half}, short of r_build {self._r_build} A. Pairs " + "further apart than the stencil are dropped without being measured, so " + "this is a tripwire on the derivation above, not a user knob." + ) + + axes = [ + torch.unique( + torch.arange(-half, half + 1, dtype=torch.long, device=self._device) % count + ) + for count, half in zip(counts, halves, strict=True) + ] + self._stencil = torch.stack( + [axis.reshape(-1) for axis in torch.meshgrid(*axes, indexing="ij")], dim=-1 + ) + # Inverted in float64 and cast down: a float32 cell inverted in float32 + # loses digits the bin index is then floored from. + self._inv_cell = torch.linalg.inv(self.cell.to(torch.float64)).to( + device=self._device, dtype=self.cell.dtype + ) + self.n_bins = (counts[0], counts[1], counts[2]) + + def _build_binned( + self, positions: torch.Tensor, n_bins: tuple[int, int, int] + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Build the pair list from the bin grid; returns ``(source, target, shifts)``. + + The pure-torch O(N) backend behind ``bin=``, with the same return + contract as :meth:`_compute` so :meth:`_write` — capacity, overflow and + dead-edge padding — is shared verbatim. Atoms are sorted into the + :meth:`_configure_bins` grid once, then each of the ``S`` stencil offsets + gathers one neighbour bin per atom; the Python loop runs over those + offsets (``S <= 4913``, independent of ``N``) and everything inside it is + vectorised over all atoms at once. Every tensor is created with an + explicit ``device=`` / ``dtype=`` taken from the positions, so the path + runs on CUDA unchanged. Eager, like every other build here. + + Fractional coordinates are wrapped into ``[0, 1)`` **for indexing only**: + displacements are taken from the unwrapped ``frac`` and the stored + ``positions``, so a trajectory that has drifted out of the box keeps + both its coordinates and its shifts (module docstring, "Raw + displacements, unwrapped positions"). + + Args: + positions: Positions ``(N, 3)`` in Angstrom, wrapped or not. Cast to + the cell's dtype — the precision the buffers already hold. + n_bins: Bins per axis, i.e. :attr:`n_bins` narrowed to non-``None`` + by :meth:`_build_pairs`, which is the only caller. + + Returns: + ``(source, target, shifts)``: atom indices ``(E,)`` per the repo + edge convention and periodic remainders ``(E, 3)`` in Angstrom, as a + **full bidirectional** list — each pair appears as ``(s, t, D)`` and + ``(t, s, -D)``. Edge *order* is not part of the contract (module + docstring). + """ + cell = self.cell + pos = positions.detach().to(cell.dtype) + device, dtype = pos.device, pos.dtype + n_atoms = int(pos.shape[0]) + counts_per_axis = torch.tensor(n_bins, dtype=torch.long, device=device) + strides = torch.tensor( + [n_bins[1] * n_bins[2], n_bins[2], 1], dtype=torch.long, device=device + ) + + frac = pos @ self._inv_cell + # Indexing device only — never differenced, never stored. + wrapped = frac - frac.floor() + bin_ijk = (wrapped * counts_per_axis.to(dtype)).floor().long() + # clamp: frac_w can round to exactly 1.0, and a negative index would wrap + # silently through Python's semantics instead of landing in bin 0. + bin_ijk = bin_ijk.clamp_(min=0).minimum(counts_per_axis - 1) + bin_id = (bin_ijk * strides).sum(-1) + + order = torch.argsort(bin_id, stable=True) + occupancy = torch.bincount(bin_id, minlength=n_bins[0] * n_bins[1] * n_bins[2]) + starts = torch.cumsum(occupancy, 0) - occupancy + atoms = torch.arange(n_atoms, dtype=torch.long, device=device) + r_build_sq = self._r_build**2 + + sources: list[torch.Tensor] = [] + targets: list[torch.Tensor] = [] + remainders: list[torch.Tensor] = [] + for offset in self._stencil: + neighbour_bin = (((bin_ijk + offset) % counts_per_axis) * strides).sum(-1) + occupied = occupancy[neighbour_bin] + total = int(occupied.sum()) + if total == 0: + continue + # Ragged gather, the molix.data.collate._gather_indices idiom (that + # one is CPU-pinned for DataLoader workers, so it is followed, not + # imported): counts -> segment ids -> exclusive cumsum -> row index. + segment = torch.repeat_interleave(atoms, occupied) + exclusive = torch.cumsum(occupied, 0) - occupied + candidate = order[ + starts[neighbour_bin][segment] + + (torch.arange(total, dtype=torch.long, device=device) - exclusive[segment]) + ] + + half_pair = segment < candidate # each unordered pair once, no self pairs + source, target = segment[half_pair], candidate[half_pair] + if source.numel() == 0: + continue + # Minimum image by fractional rounding — exact here because + # r_build <= min_i w_i / 2 bounds every in-range image by |f_i| <= 1/2. + fractional = frac[target] - frac[source] + displacement = (fractional - fractional.round()) @ cell + distance_sq = (displacement * displacement).sum(-1) + # 0 < r <= r_build, the compiled kernels' filter verbatim: coincident + # atoms and pairs separated by exactly one lattice vector are dropped. + keep = (distance_sq <= r_build_sq) & (distance_sq > 0) + source, target, displacement = source[keep], target[keep], displacement[keep] + sources.append(source) + targets.append(target) + # Same definition as _compute's edge_diff - (pos[target] - pos[source]), + # against the stored (unwrapped) positions. + remainders.append(displacement - (pos[target] - pos[source])) + + if sources: + half_source, half_target = torch.cat(sources), torch.cat(targets) + half_shifts = torch.cat(remainders) + else: # an isolated system at this radius: no half pairs to expand + half_source = torch.zeros(0, dtype=torch.long, device=device) + half_target = torch.zeros(0, dtype=torch.long, device=device) + half_shifts = torch.zeros(0, 3, dtype=dtype, device=device) + # Symmetry expansion to the full bidirectional list: the shift flips + # sign wholesale with the displacement it is the remainder of. + return ( + torch.cat((half_source, half_target)), + torch.cat((half_target, half_source)), + torch.cat((half_shifts, -half_shifts)), + ) + + def _build_pairs( + self, positions: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Route the build to the backend ``bin`` selected; ``(source, target, shifts)``. + + The single seam between the two backends: the constructor's initial + (capacity-sizing) build and every :meth:`_build_at` go through it, so a + binned list is binned from the first edge on and its capacity is sized + by the backend that will keep refilling it. + + Args: + positions: Positions ``(N, 3)`` in Angstrom to build at. + """ + grid = self.n_bins + if grid is None: + return self._compute(positions) + return self._build_binned(positions, grid) + + def _dead_shift(self) -> torch.Tensor: + """A displacement long enough that every cutoff envelope evaluates to 0.""" + return torch.tensor( + [DEAD_EDGE_CUTOFF_FACTOR * self.cutoff, 0.0, 0.0], + dtype=self._dtype, + device=self._device, + ) + + def _write(self, source: torch.Tensor, target: torch.Tensor, shifts: torch.Tensor) -> None: + """Fill the buffers in place; pad the tail with dead edges.""" + n = int(source.numel()) + if n > self.capacity: + raise RuntimeError( + f"neighbour-list overflow: {n} edges > capacity {self.capacity}. Raise " + f"capacity_factor (currently {self.capacity_factor})." + ) + # In-place so the compiled graph keeps seeing the same tensors. + self.edge_index[:n, 0] = source + self.edge_index[:n, 1] = target + self.shifts[:n] = shifts.to(self._dtype) + self.edge_index[n:] = 0 + self.shifts[n:] = self._dead_shift() + self.num_edges = n + + def _hold(self, positions: torch.Tensor) -> None: + """Adopt ``positions`` as the reference of the current build. + + Restarts the ``ago`` clock and refreshes ``_x_hold`` in place (cast to + the buffer dtype, so a list that was moved with :meth:`to` keeps + differencing in one precision). + """ + self.ago = 0 + self._x_hold.copy_(positions.detach().to(self._x_hold.dtype)) + + def _positions_from(self, batch: TensorDictBase) -> torch.Tensor: + """Validate that ``batch`` describes *this* system; return its positions. + + Shared by :meth:`build` and the batch arm of :meth:`update`, so the + per-step path cannot drift from the bind-time path. The checks are + metadata compares — key presence, shape, ``device``, ``dtype`` — and so + cost the hot loop nothing measurable; the single value read is the + optional cell comparison, and an MD working batch carries no cell. + + Args: + batch: Batch ``TensorDict`` with positions at ``("atoms", "pos")`` + in Angstrom, optionally carrying ``("graphs", "cell")``. + + Returns: + The batch's positions ``(N, 3)`` in Angstrom — the *same* tensor, + never a cast copy. + + Raises: + ValueError: If ``("atoms", "pos")`` is missing; if its shape is not + ``(N, 3)`` for the ``N`` this list was constructed with; if its + device or dtype differ from the list's buffers (the owner + casts, via :meth:`to` — no silent cast here); or if + ``("graphs", "cell")`` is present and is neither ``(3, 3)`` nor + ``(1, 3, 3)`` equal to the constructor cell within ``1e-8`` A + (``rtol=0``) — loose enough to survive a float32 template + round-trip, far below any physically meaningful difference. The + constructor cell stays the **owner**: the batch's copy is + checked, never adopted, because every frozen shift in the + buffers is a lattice vector of *that* cell. + """ + if ("atoms", "pos") not in batch.keys(include_nested=True): + present = sorted( + "/".join(key) if isinstance(key, tuple) else str(key) + for key in batch.keys(include_nested=True) + ) + raise ValueError( + "the batch carries no ('atoms', 'pos'): a neighbour list builds at atom " + f"positions, and this batch holds {present}. Post-collate batches nest " + "positions under the 'atoms' namespace (two-tier data contract); a flat " + "sample dict is the other tier and is not what this binds into." + ) + pos = batch["atoms", "pos"] + if not isinstance(pos, torch.Tensor): + raise ValueError( + "('atoms', 'pos') must be a positions tensor (N, 3) in Angstrom, but the " + f"batch holds a {type(pos).__name__} there — a nested namespace, not the " + "leaf this list builds at." + ) + n_atoms = int(self._x_hold.shape[0]) + if tuple(pos.shape) != (n_atoms, 3): + raise ValueError( + f"positions must have shape ({n_atoms}, 3) in Angstrom — this list was " + f"constructed for {n_atoms} atoms — but the batch's ('atoms', 'pos') has " + f"shape {tuple(pos.shape)}. A different atom count is a different system: " + "the capacity was sized for the original and the displacement reference " + "_x_hold has its shape." + ) + # _x_hold, not _device/_dtype: it is the tensor these positions are + # actually differenced against, and it tracks every to() exactly. + if pos.device != self._x_hold.device or pos.dtype != self._x_hold.dtype: + raise ValueError( + f"positions are on {pos.device} in {pos.dtype}, but this list's buffers are " + f"on {self._x_hold.device} in {self._x_hold.dtype}. Nothing is cast here: a " + "silently promoted difference against _x_hold is a mixed-precision " + "comparison nobody asked for, and a cross-device index is worse. The owner " + "casts — call NeighborList.to(device, dtype) first." + ) + if ("graphs", "cell") not in batch.keys(include_nested=True): + return pos + cell = batch["graphs", "cell"] + if not isinstance(cell, torch.Tensor): + raise ValueError( + "('graphs', 'cell') must be a cell tensor (3, 3) — or (1, 3, 3) for a " + f"single-system batch — in Angstrom, but the batch holds a " + f"{type(cell).__name__} there." + ) + vectors = cell + if vectors.dim() == 3: + if vectors.shape[0] != 1: + raise ValueError( + f"('graphs', 'cell') has a leading batch dimension of {vectors.shape[0]} " + f"(shape {tuple(vectors.shape)}), but this neighbour list is " + "single-system: one cell, one displacement reference, one capacity. " + "Multi-system batched MD is not supported here rather than silently " + "reduced to the first cell." + ) + vectors = vectors[0] + if tuple(vectors.shape) != (3, 3): + raise ValueError( + "('graphs', 'cell') must be (3, 3), or (1, 3, 3) for a single-system batch, " + f"in Angstrom; got shape {tuple(cell.shape)}." + ) + reference = self.cell.detach().to(torch.float64) + candidate = vectors.detach().to(device=reference.device, dtype=torch.float64) + if not torch.allclose(candidate, reference, rtol=0.0, atol=1e-8): + raise ValueError( + f"the batch's ('graphs', 'cell')\n{candidate.tolist()}\ndisagrees with the " + f"cell this list was built against\n{reference.tolist()}\n(Angstrom, tolerance " + "atol=1e-8, rtol=0). The stored shifts are lattice vectors of the " + "constructor's cell, so adopting a different one would leave every periodic " + "remainder silently wrong; construct a new list instead." + ) + return pos + + def _build_at(self, positions: torch.Tensor) -> None: + """Recompute and rewrite the buffers at ``positions``, restarting ``ago``. + + The shared body of :meth:`rebuild` (which adds the counter increment) + and :meth:`build` (which adds the batch validation and the bind). + Deliberately does *not* touch :attr:`rebuild_count`: only the caller + knows whether this build is a rebuild driven by the run. + + Args: + positions: Positions ``(N, 3)`` in Angstrom to build at. + """ + source, target, shifts = self._build_pairs(positions) + self._write(source, target, shifts) + self._hold(positions) + + def rebuild(self, positions: torch.Tensor) -> None: + """Recompute the neighbour list at ``positions`` (eager, outside any graph). + + The unconditional primitive: it builds whatever the policy would have + said. ``ago`` restarts here, so a forced rebuild also re-phases the + ``every`` / ``delay`` schedule :meth:`update` runs on. + + Args: + positions: Positions ``(N, 3)`` in Angstrom to build at. + """ + self._build_at(positions) + self.rebuild_count += 1 + + def build(self, batch: TensorDict) -> TensorDict: + """Rebuild at the batch's positions and bind the live buffers into it. + + The TensorDict-side forced build: it validates that ``batch`` describes + the system this list was constructed for, rebuilds at + ``batch["atoms", "pos"]`` (Angstrom), then writes ``batch["edges"]`` as + a ``TensorDict`` of :attr:`edge_index` ``(capacity, 2)`` and + :attr:`shifts` ``(capacity, 3)`` held **by reference**, at + ``batch_size=[capacity]``. Every later in-place :meth:`rebuild` / + :meth:`update` is therefore visible to whatever reads that batch, with + no shape change and no re-binding — which is what keeps a compiled or + graph-captured force path valid across a rebuild. + + **The list owns ``edges`` once bound.** The namespace is replaced + wholesale, not merged: a pre-existing ``edge_diff`` / ``edge_dist`` + pair would be consumed straight through as a *value* by a potential and + freeze the PES, and a surviving shorter ``edge_index`` would disagree + with the capacity. + + This is *not* counted as a :attr:`rebuild_count` rebuild — it is a + binding operation, and it runs again on every ``.to()`` re-sync, so + counting it would make a dtype cast look like physics. It does restart + the policy clock (:attr:`ago` back to 0, ``_x_hold`` refreshed): the + buffers are fresh here. + + Warning: + ``build`` returns the batch so it composes with the repo's + ``forward(td) -> td`` convention (``potential(nl.build(batch))``) — + **not** so the policy call can be chained. + ``nl.build(batch).update(batch)`` parses, but that ``.update`` is + ``TensorDict.update``: it merges the batch into itself and never + touches this list. The idiom is two statements:: + + nl.build(batch) # once, and after every .to() + ... + nl.update(batch) # per step — NeighborList.update + + Args: + batch: Batch ``TensorDict`` carrying ``("atoms", "pos")`` + ``(N, 3)`` in Angstrom and optionally ``("graphs", "cell")`` + ``(3, 3)`` or ``(1, 3, 3)`` in Angstrom, which is validated + against the constructor cell and never adopted. + + Returns: + The **same** ``batch`` object, with ``batch["edges"]`` bound. + + Raises: + ValueError: Per :meth:`_positions_from` — missing positions, a + different atom count, a device/dtype the owner has not cast, or + a cell that is not this list's. + RuntimeError: If the rebuilt edge count overflows the capacity. + """ + positions = self._positions_from(batch) + self._build_at(positions) + batch["edges"] = TensorDict( + {"edge_index": self.edge_index, "shifts": self.shifts}, + batch_size=[self.capacity], + ) + return batch + + def update(self, positions: TensorDict | torch.Tensor) -> bool: + """Rebuild at ``positions`` if the ``every``/``delay``/``check`` gate says so. + + Call once per force evaluation, at the positions being evaluated. The + gate is ``Neighbor::decide`` verbatim: ``ago`` is incremented, a rebuild + is *permitted* only when ``ago >= delay`` **and** ``ago % every == 0`` + (conjunctive), and — with ``check`` — happens only when the largest raw + displacement since the last build exceeds ``skin/2``. Crossing that + bound at the first permitted opportunity ``ago == max(every, delay)`` + increments :attr:`ndanger`. + + **One method, two input types.** A raw ``(N, 3)`` tensor is the MD hot + path; a batch ``TensorDict`` is the bound path, dispatched by an + ``isinstance`` test against ``TensorDictBase`` (so lazy / stacked + batches dispatch too) and reduced to its positions by the same + validation :meth:`build` runs — a batch whose ``pos`` was silently + re-cast therefore fails loud here instead of promoting against + ``_x_hold`` for the rest of the run. Everything after the dispatch is + identical: same decisions, same :attr:`ago` / :attr:`rebuild_count` / + :attr:`ndanger` bookkeeping, same buffers. There is no ``update_td`` / + ``update_pos`` pair — the policy state lives behind one door. + + Note: + ``nl.update(batch)`` is :class:`NeighborList`'s ``update``; + ``batch.update(...)`` is ``TensorDict.update``, a merge that never + touches this list. Keep the bind and the step as two statements + (see :meth:`build`). + + Args: + positions: Positions ``(N, 3)`` in Angstrom, **unwrapped** (see + Raises and the module docstring), or the batch + ``TensorDict`` carrying them at ``("atoms", "pos")``. + + Returns: + ``True`` if the list was rebuilt, ``False`` if the frozen list is + still valid (or the gate simply did not permit a build this step). + + Raises: + ValueError: If a batch was passed and it does not describe this + system — see :meth:`_positions_from`. + RuntimeError: If the largest displacement since the last build + reaches half the smallest perpendicular cell width — mid-run + wrapping, a changed cell, or a blown-up trajectory. Not raised + under ``check=False``, which skips the displacement branch + entirely. + """ + if isinstance(positions, TensorDictBase): + positions = self._positions_from(positions) + self.ago += 1 + if self.ago < self.delay or self.ago % self.every: + return False # not a permitted opportunity + if not self.check: + self.rebuild(positions) + return True + # Raw difference, never a minimum image: min-imaging would clamp a + # genuine > L/2 excursion and suppress the rebuild it should force. + max_d2 = float(((positions.detach() - self._x_hold) ** 2).sum(-1).max()) + if max_d2 >= self._wrap_guard_sq: + raise RuntimeError( + f"positions are no longer unwrapped: the largest displacement since the last " + f"build is {math.sqrt(max_d2):.3f} A, at or beyond half the minimum " + f"perpendicular cell width ({math.sqrt(self._wrap_guard_sq):.3f} A). The frozen " + "periodic shifts and this raw displacement test hold only for continuously " + "drifting coordinates, so a jump this large means the positions were wrapped " + "mid-run, the cell changed, or the trajectory blew up. All three are fatal; " + "none is recoverable by rebuilding." + ) + if max_d2 > self._half_skin_sq: # strict: exactly skin/2 is still complete + if self.ago == self._danger_ago: + self.ndanger += 1 + self.rebuild(positions) + return True + return False + + def to( + self, + device: torch.device | str | torch.dtype | None = None, + dtype: torch.dtype | None = None, + ) -> "NeighborList": + """Move / cast the buffers, mirroring ``Tensor.to`` semantics. + + Accepts ``nl.to("cuda")``, ``nl.to(torch.float64)`` and + ``nl.to(device, dtype)`` alike. ``_x_hold`` travels with the rest: a + displacement reference left behind in the old dtype would silently + promote the next :meth:`update` comparison instead of failing. + + Warning: + **This severs any :meth:`build` binding.** The move rebinds + :attr:`edge_index` / :attr:`shifts` to *new* tensors, so a batch + bound beforehand keeps pointing at the old ones and would freeze at + the pre-cast neighbour set. Re-binding from inside ``to`` is + deliberately not done — it would make the list hold a reference to + a batch it does not own — so the **owner** re-binds: call + ``nl.build(batch)`` after the cast (this is exactly what + ``PeriodicPotentialForceField._apply`` does). + """ + if isinstance(device, torch.dtype): + if dtype is not None: + raise TypeError("dtype given twice") + device, dtype = None, device + if device is not None: + self._device = torch.device(device) + self.edge_index = self.edge_index.to(self._device) + self.shifts = self.shifts.to(self._device) + self.cell = self.cell.to(self._device) + self._x_hold = self._x_hold.to(self._device) + if dtype is not None: + self._dtype = dtype + self.shifts = self.shifts.to(dtype) + self.cell = self.cell.to(dtype) + self._x_hold = self._x_hold.to(dtype) + if self.n_bins is not None: + # The grid itself is a property of the cell and r_build, so only its + # two tensors move: the stencil is integer offsets, and the inverse + # cell is re-derived (in float64, then cast) rather than converted, + # so a float32 hop does not compound its own rounding. + self._stencil = self._stencil.to(self._device) + self._inv_cell = torch.linalg.inv(self.cell.to(torch.float64)).to( + device=self._device, dtype=self.cell.dtype + ) + return self diff --git a/src/molix/md/runner.py b/src/molix/md/runner.py index 0c6e65d..6c0392b 100644 --- a/src/molix/md/runner.py +++ b/src/molix/md/runner.py @@ -1,47 +1,84 @@ -"""Hook-driven MD runner: drive an :class:`Integrator` through the hook lifecycle. - -:class:`MDRunner` plays the role :class:`molix.core.trainer.Trainer` plays for -training — it owns a :class:`~molix.core.state.TrainState` and a priority-sorted -hook list — but its loop is a molecular-dynamics integration. It fires -``on_train_start`` once, ``on_train_batch_end`` per step, ``on_train_end`` at the -close, advancing ``state["global_step"]`` each step, so the existing hook -ecosystem observes an MD run unchanged. - -Hook dispatch is **static**: hooks are :class:`~molix.core.hook.BaseHook` -instances (all lifecycle methods have no-op defaults), so the runner calls the -typed methods directly — no ``getattr`` name lookup. The integrator advances a -typed :class:`~molix.md.types.MDState`; per-step physics is unpacked into the -``outputs`` dict (the same channel the Trainer uses) so the runner never writes -physics into the reserved ``TrainState`` namespaces. +"""Hook-driven MD runner: drive an :class:`Integrator` through an MD lifecycle. + +:class:`MDRunner` owns the observation loop the way +:class:`molix.core.trainer.Trainer` owns the training loop, but it speaks its +own, deliberately narrow protocol: :class:`MDHook`. MD hooks receive the step +count and typed physics (:class:`~molix.md.types.MDState` / +:class:`~molix.md.types.MDObservables`) — they are **not** Trainer hooks, and +Trainer hooks (which dereference ``trainer.model`` / ``trainer.optimizer``) +are not accepted. One runner, one honest contract. + +Hook dispatch is **static**: hooks subclass :class:`MDHook` (all lifecycle +methods have no-op defaults), so the runner calls the typed methods directly — +no ``getattr`` name lookup. A hook that acts on a step cadence declares it via +:attr:`MDHook.cadence` so :meth:`MDRunner.run` can refuse a ``chunk`` that +would silently skip firings. :class:`TrajectoryHook` captures strided frames to host buffers, spilling to on-disk shards every ``flush_every`` frames so host memory stays bounded, and -writes one ``.pt`` (+ optional extended-XYZ) at ``on_train_end``. +writes one ``.pt`` (+ optional extended-XYZ via +:func:`molix.datasets._extxyz.write_extxyz_frames`) at :meth:`MDHook.on_run_end`. """ from __future__ import annotations from collections.abc import Sequence from pathlib import Path -from typing import Any import torch -from molix.core.hook import BaseHook -from molix.core.state import Stage, TrainState -from molix.md.integrators import EV_PER_AMU_A2_FS2, LangevinVerletIntegrator, as_mass_col +from molix.md.integrators import Integrator, _as_mass_col +from molix.md.types import MDObservables, MDState +from molix.units import KB_AMU_A_FS + + +class MDHook: + """Lifecycle observer for an MD run — the MD-specific hook contract. + + Subclass and override only what you need; every method is a no-op by + default. Hooks that act on a step cadence (every N-th step) must declare + it in :attr:`cadence` so the runner can validate ``chunk`` against it. + """ + + #: Steps between the firings this hook acts on (``None``: every + #: observation). :meth:`MDRunner.run` rejects a ``chunk`` that is not a + #: divisor of a declared cadence — chunking must never silently skip a + #: hook's step. + cadence: int | None = None + + def on_run_start(self, runner: "MDRunner") -> None: + """Called once before the first step (the entry force is already cached).""" + + def on_step_start(self, runner: "MDRunner", step: int, state: MDState) -> None: + """Called before each hook-visible advance. + + Not the place to refresh a neighbour list: velocity-Verlet evaluates + ``F`` at the *end-of-step* positions, so a step-start rebuild leaves the + connectivity one displacement behind. The list's policy runs inside + :meth:`~molix.md.integrators.Integrator.eval_force` instead. + + Args: + runner: The driving runner. + step: Steps completed so far (``0`` on the first call). + state: The live state the upcoming advance will consume. + """ + + def on_step_end(self, runner: "MDRunner", step: int, obs: MDObservables) -> None: + """Called after each hook-visible advance with the step's physics. + + Args: + runner: The driving runner. + step: Steps completed including this advance. + obs: Typed thermodynamic snapshot at ``step``. + """ -#: Boltzmann constant in eV/K. (``molix.quant`` keeps its own copy for the -#: quantization subsystem; this is the MD package's single named source.) -KB_EV_PER_K = 8.617333262e-5 -#: k_B in the integrator's (amu, Å, fs) energy unit (amu·Å²/fs²), so temperature -#: comes out in kelvin: k_B[eV/K] / (1 amu·Å²/fs² in eV). -KB_AMU_A_FS = KB_EV_PER_K / EV_PER_AMU_A2_FS2 + def on_run_end(self, runner: "MDRunner") -> None: + """Called once after the last step (persist buffered results here).""" def _normalize_hooks( - hooks: Sequence[BaseHook | tuple[BaseHook, int]] | None, -) -> list[BaseHook]: + hooks: Sequence[MDHook | tuple[MDHook, int]] | None, +) -> list[MDHook]: """Priority-sort hooks (lower priority first; ties keep registration order).""" if not hooks: return [] @@ -57,98 +94,112 @@ def _normalize_hooks( class MDRunner: - """Drive a :class:`LangevinVerletIntegrator` through the hook lifecycle. + """Drive an :class:`~molix.md.integrators.Integrator` through the MD hook lifecycle. Args: integrator: The integrator advancing the typed ``MDState``. mass: Per-atom mass ``(N,)`` or scalar (integrator's mass unit). Used to report kinetic energy / temperature. - hooks: :class:`~molix.core.hook.BaseHook` instances or ``(hook, priority)`` - tuples; same protocol as the Trainer. + hooks: :class:`MDHook` instances or ``(hook, priority)`` tuples; lower + priority fires earlier, ties keep registration order. kb: Boltzmann constant in the integrator's energy unit (default: the (amu, Å, fs) value, so temperature comes out in kelvin). - dof: Degrees of freedom for the temperature estimator; defaults to ``3 N`` - under Langevin (γ>0, the O step thermostats the COM too) and ``3 N - 3`` - under NVE (centre-of-mass momentum removed). + dof: Degrees of freedom for the temperature estimator; defaults to + ``3 N - integrator.removed_dof`` (``3 N`` under Langevin — the O + step thermostats the COM too — and ``3 N - 3`` under NVE with + centre-of-mass momentum removed). """ def __init__( self, - integrator: LangevinVerletIntegrator, + integrator: Integrator, *, mass: float | torch.Tensor, - hooks: Sequence[BaseHook | tuple[BaseHook, int]] | None = None, + hooks: Sequence[MDHook | tuple[MDHook, int]] | None = None, kb: float = KB_AMU_A_FS, dof: int | None = None, ) -> None: self.integrator = integrator - self.hooks: list[BaseHook] = _normalize_hooks(hooks) + self.hooks: list[MDHook] = _normalize_hooks(hooks) self._mass = mass self._kb = float(kb) self._dof = dof - self.state = TrainState() - def run(self, pos: torch.Tensor, vel: torch.Tensor, n_steps: int) -> dict[str, Any]: - """Integrate ``n_steps`` steps, firing the hook lifecycle each step. + def run(self, pos: torch.Tensor, vel: torch.Tensor, n_steps: int, *, chunk: int = 1) -> MDState: + """Integrate ``n_steps`` steps, firing the hook lifecycle per chunk. + + ``chunk > 1`` advances the integrator ``chunk`` steps between hook + firings (``Integrator.advance_n`` — no per-step Python, no per-step + thermodynamics). The dynamics are bit-identical to ``chunk=1``; only + the observation cadence changes, so every declared hook cadence + (``TrajectoryHook.stride``, ``MDCheckpointHook.every``) must be a + multiple of ``chunk`` — enforced via :attr:`MDHook.cadence`. The + neighbour policy is **not** among them: it runs inside + :meth:`~molix.md.integrators.Integrator.eval_force`, once per force + evaluation, whatever ``chunk`` is. Args: pos: Initial positions ``(N, 3)``. vel: Initial velocities ``(N, 3)``. n_steps: Number of MD steps. + chunk: Steps advanced between hook firings. Returns: - Dict with the final ``pos`` / ``vel`` / ``force`` tensors and the - terminal :class:`~molix.core.state.TrainState`. Trajectory capture is - the job of hooks (see :class:`TrajectoryHook`). + The final typed :class:`~molix.md.types.MDState`. Trajectory + capture is the job of hooks (see :class:`TrajectoryHook`). """ - state = self.state - state["stage"] = Stage.TRAIN - state["global_step"] = 0 - mass = as_mass_col(self._mass, pos) + mass = _as_mass_col(self._mass, pos) if self._dof is not None: dof = self._dof else: - # Langevin (γ>0) thermostats all 3N DoF including the COM; NVE with - # COM momentum removed leaves 3N-3. 3N-3 under Langevin would - # over-report T by 3N/(3N-3). - n = int(pos.shape[0]) - dof = max(1, 3 * n - (0 if self.integrator.gamma > 0.0 else 3)) + dof = max(1, 3 * int(pos.shape[0]) - self.integrator.removed_dof) + chunk = max(1, int(chunk)) for hook in self.hooks: - hook.on_train_start(self, state) + if hook.cadence is not None and hook.cadence % chunk: + raise ValueError( + f"{type(hook).__name__} fires every {hook.cadence} steps, which chunk=" + f"{chunk} would silently skip; make it a multiple of chunk" + ) md = self.integrator.initial(pos, vel) - for i in range(n_steps): - md = self.integrator.advance(md) + for hook in self.hooks: + hook.on_run_start(self) + done = 0 + while done < n_steps: + n = min(chunk, n_steps - done) + # Fired *before* the advance so a hook can refresh position-derived + # state (the neighbour list) while it still precedes the force + # evaluation inside. + for hook in self.hooks: + hook.on_step_start(self, done, md) + md = self.integrator.advance_n(md, n) + done += n kinetic = 0.5 * (mass * md.vel * md.vel).sum() potential = md.energy.reshape(()) temperature = 2.0 * kinetic / (dof * self._kb) - state["global_step"] = i + 1 - # Physics rides the ``outputs`` channel (like Trainer step outputs), - # NOT the reserved state namespaces — keeps the state contract clean. - outputs = { - "pos": md.pos, - "vel": md.vel, - "forces": md.force, - "potential": potential, - "kinetic": kinetic, - "total": potential + kinetic, - "temperature": temperature, - } - batch = {"pos": md.pos, "vel": md.vel} + obs = MDObservables( + pos=md.pos, + vel=md.vel, + forces=md.forces, + potential=potential, + kinetic=kinetic, + total=potential + kinetic, + temperature=temperature, + ) for hook in self.hooks: - hook.on_train_batch_end(self, state, batch, outputs) + hook.on_step_end(self, done, obs) for hook in self.hooks: - hook.on_train_end(self, state) - return {"pos": md.pos, "vel": md.vel, "force": md.force, "state": state} + hook.on_run_end(self) + return md -class TrajectoryHook(BaseHook): +class TrajectoryHook(MDHook): """Capture an MD trajectory to host buffers; persist at run end. Strided frames are copied to CPU and accumulated on the host; every ``flush_every`` kept frames the buffer is spilled to an on-disk shard and cleared, so host memory stays O(``flush_every`` · N) regardless of run - length. At ``on_train_end`` the shards (if any) are concatenated into one + length. At :meth:`on_run_end` the shards (if any) are concatenated into one ``.pt`` (+ optional extended-XYZ) and removed. Runs whose kept frames fit one buffer skip sharding entirely (identical output to a single write). @@ -166,7 +217,6 @@ class TrajectoryHook(BaseHook): flush_every: Kept-frame budget before spilling a shard to disk. Default 10000. """ - _SYMBOLS = {1: "H", 6: "C", 7: "N", 8: "O", 9: "F", 15: "P", 16: "S", 17: "Cl"} _FIELDS = ("pos", "vel", "f", "pe", "ke", "etot", "temp") def __init__( @@ -185,22 +235,23 @@ def __init__( self._write_xyz = write_xyz self._with_forces = with_forces self._flush_every = max(1, int(flush_every)) + self.cadence = self._stride self._buf: dict[str, list[torch.Tensor]] = {k: [] for k in self._FIELDS} self._n_buffered = 0 self._shards: list[Path] = [] - def on_train_batch_end(self, trainer: Any, state: TrainState, batch: Any, outputs: Any) -> None: - if state["global_step"] % self._stride: + def on_step_end(self, runner: MDRunner, step: int, obs: MDObservables) -> None: + if step % self._stride: return b = self._buf - b["pos"].append(outputs["pos"].detach().to("cpu")) - b["vel"].append(outputs["vel"].detach().to("cpu")) - if self._with_forces and outputs.get("forces") is not None: - b["f"].append(outputs["forces"].detach().to("cpu")) - b["pe"].append(outputs["potential"].detach().to("cpu")) - b["ke"].append(outputs["kinetic"].detach().to("cpu")) - b["etot"].append(outputs["total"].detach().to("cpu")) - b["temp"].append(outputs["temperature"].detach().to("cpu")) + b["pos"].append(obs.pos.detach().to("cpu")) + b["vel"].append(obs.vel.detach().to("cpu")) + if self._with_forces: + b["f"].append(obs.forces.detach().to("cpu")) + b["pe"].append(obs.potential.detach().to("cpu")) + b["ke"].append(obs.kinetic.detach().to("cpu")) + b["etot"].append(obs.total.detach().to("cpu")) + b["temp"].append(obs.temperature.detach().to("cpu")) self._n_buffered += 1 if self._n_buffered >= self._flush_every: self._flush_shard() @@ -234,7 +285,7 @@ def _combine_shards(self) -> dict[str, torch.Tensor] | None: self._shards.clear() return fields - def on_train_end(self, trainer: Any, state: TrainState) -> None: + def on_run_end(self, runner: MDRunner) -> None: if self._shards: self._flush_shard() # spill the trailing partial buffer fields = self._combine_shards() @@ -242,7 +293,7 @@ def on_train_end(self, trainer: Any, state: TrainState) -> None: fields = self._stack_buffer() if fields is None: return - payload: dict[str, Any] = { + payload: dict[str, torch.Tensor | int] = { "pos": fields["pos"].to(torch.float32), "vel": fields["vel"].to(torch.float32), "pe": fields["pe"], @@ -261,13 +312,72 @@ def on_train_end(self, trainer: Any, state: TrainState) -> None: self._dump_xyz(payload["pos"], payload["etot"], payload["temp"]) def _dump_xyz(self, pos: torch.Tensor, etot: torch.Tensor, temp: torch.Tensor) -> None: - zs = [int(z) for z in self._numbers.detach().cpu()] # type: ignore[union-attr] - syms = [self._SYMBOLS.get(z, "X") for z in zs] - n = len(zs) - with self._out.with_suffix(".xyz").open("w") as fh: - for t in range(pos.shape[0]): - coords = pos[t].to(torch.float64).numpy() - fh.write(f"{n}\n") - fh.write(f"Etot={float(etot[t]):.6f} T={float(temp[t]):.2f} frame={t}\n") - for sy, (x, y, z) in zip(syms, coords): - fh.write(f"{sy} {x:.6f} {y:.6f} {z:.6f}\n") + from molpy import Element + + from molix.datasets._extxyz import write_extxyz_frames + + species = [Element(int(z)).symbol for z in self._numbers.detach().cpu()] + write_extxyz_frames( + self._out.with_suffix(".xyz"), + species=species, + positions=pos.to(torch.float64).numpy(), + energies=etot.to(torch.float64).numpy(), + tags=[f"temperature={float(t):.2f}" for t in temp], + ) + + +class MDCheckpointHook(MDHook): + """Persist a restartable NVE state (pos, vel, absolute step) every N steps. + + Named ``MDCheckpointHook`` — :class:`molix.hooks.CheckpointHook` is the + training-side checkpointer with an unrelated constructor; the two must not + collide in a ``from molix... import *`` namespace. + + Multi-hour trajectories die to walltime and node failures; without this, + everything after the last :class:`TrajectoryHook` shard is gone. The write + is atomic (temp file + ``rename``) so a kill mid-write leaves the previous + checkpoint intact. Restart is exact for γ=0 — an NVE state is fully + determined by ``(pos, vel)`` — and approximate for γ>0 (the Langevin noise + stream restarts, which changes the realisation but not the ensemble). + + Doubles as the run's heartbeat: each checkpoint prints one line, so a + day-long job's log shows progress instead of silence. + + Args: + path: Checkpoint file, overwritten in place. + every: Step interval between checkpoints. + step_offset: Absolute step count this run resumed from, added to the + in-run step so a chain of resumed segments keeps one monotonic + step axis. + """ + + def __init__(self, path: str | Path, *, every: int, step_offset: int = 0) -> None: + if every < 1: + raise ValueError(f"every must be >= 1, got {every}") + self._path = Path(path) + self._every = int(every) + self._offset = int(step_offset) + self.cadence = self._every + + def on_step_end(self, runner: MDRunner, step: int, obs: MDObservables) -> None: + if step % self._every == 0: + self._save(step, obs) + + def _save(self, step: int, obs: MDObservables) -> None: + absolute = self._offset + int(step) + self._path.parent.mkdir(parents=True, exist_ok=True) + tmp = self._path.with_suffix(self._path.suffix + ".tmp") + torch.save( + { + "pos": obs.pos.detach().cpu(), + "vel": obs.vel.detach().cpu(), + "step": absolute, + }, + tmp, + ) + tmp.replace(self._path) + print( + f"[checkpoint] step {absolute} E_tot={float(obs.total):.6f} " + f"T={float(obs.temperature):.1f} K", + flush=True, + ) diff --git a/src/molix/md/types.py b/src/molix/md/types.py index f572e59..1c9736c 100644 --- a/src/molix/md/types.py +++ b/src/molix/md/types.py @@ -1,13 +1,14 @@ """Typed tensor containers for the MD engine — compile-friendly pytrees. -:class:`ForceOutput` and :class:`MDState` are ``NamedTuple``s. PyTorch registers -namedtuples as pytree nodes, so they flatten / unflatten transparently under -:func:`torch.compile`, ``torch.func`` and ``torch.utils._pytree`` with no custom -registration. They are the **only** data contract crossing component boundaries -(``ForceField`` → ``Integrator`` → ``MDRunner``); the heterogeneous topology -``TensorDict`` stays inside :class:`~molix.md.forcefield.ForceField` (model I/O). -This is the deliberate ``TensorDict`` ↔ typed-state split: rich dict where the -model needs it, static tensors where the hot loop and type-checking need it. +:class:`ForceOutput`, :class:`MDState` and :class:`MDObservables` are +``NamedTuple``s. PyTorch registers namedtuples as pytree nodes, so they +flatten / unflatten transparently under :func:`torch.compile`, ``torch.func`` +and ``torch.utils._pytree`` with no custom registration. They are the **only** +data contract crossing component boundaries (``ForceField`` → ``Integrator`` → +``MDRunner`` → ``MDHook``); the heterogeneous topology ``TensorDict`` stays +inside :class:`~molix.md.forcefield.ForceField` (model I/O). This is the +deliberate ``TensorDict`` ↔ typed-state split: rich dict where the model needs +it, static tensors where the hot loop and type-checking need it. """ from __future__ import annotations @@ -32,17 +33,44 @@ class ForceOutput(NamedTuple): class MDState(NamedTuple): """Dynamical state advanced one BAOAB step by an ``Integrator``. - Carries the force-cache (``force`` / ``energy`` at ``pos``, both from the + Carries the force-cache (``forces`` / ``energy`` at ``pos``, both from the same force-field evaluation) so the loop does one evaluation per step. Attributes: pos: Positions ``(N, 3)``. vel: Velocities ``(N, 3)``. - force: Cached force ``(N, 3)`` at ``pos``. + forces: Cached forces ``(N, 3)`` at ``pos``. energy: Cached scalar energy ``()`` at ``pos``. """ pos: torch.Tensor vel: torch.Tensor - force: torch.Tensor + forces: torch.Tensor energy: torch.Tensor + + +class MDObservables(NamedTuple): + """Per-observation thermodynamic snapshot handed to MD hooks. + + Produced by :class:`~molix.md.runner.MDRunner` after each hook-visible + chunk of steps; consumed by :class:`~molix.md.runner.MDHook.on_step_end` + implementations (trajectory capture, checkpointing, logging). + + Attributes: + pos: Positions ``(N, 3)``. + vel: Velocities ``(N, 3)``. + forces: Forces ``(N, 3)`` at ``pos``. + potential: Scalar potential energy ``()``. + kinetic: Scalar kinetic energy ``()``. + total: Scalar total energy ``()`` (``potential + kinetic``). + temperature: Scalar instantaneous temperature ``()`` in kelvin (for the + runner's default ``kb``). + """ + + pos: torch.Tensor + vel: torch.Tensor + forces: torch.Tensor + potential: torch.Tensor + kinetic: torch.Tensor + total: torch.Tensor + temperature: torch.Tensor diff --git a/src/molix/nn/__init__.py b/src/molix/nn/__init__.py index 326e1e9..e06fd84 100644 --- a/src/molix/nn/__init__.py +++ b/src/molix/nn/__init__.py @@ -1,13 +1,11 @@ """Neural network utilities for molix.""" -from .locality import NeighborList from .mlp import KeyedMLP, KeyedMLPSpec from .scatter import BatchAggregation, ScatterSum __all__ = [ + "BatchAggregation", "KeyedMLP", "KeyedMLPSpec", - "NeighborList", "ScatterSum", - "BatchAggregation", ] diff --git a/src/molix/nn/locality.py b/src/molix/nn/locality.py deleted file mode 100644 index 149fc39..0000000 --- a/src/molix/nn/locality.py +++ /dev/null @@ -1,53 +0,0 @@ -""" -Module wrappers for locality operations (molix) -""" - -import torch.nn as nn - -from ..F import locality as F - - -class NeighborList(nn.Module): - """Build neighbor pairs within a cutoff radius. - - A :class:`~torch.nn.Module` wrapper over - :func:`molix.F.locality.get_neighbor_pairs`. - - Args: - cutoff: Neighbor cutoff radius. - pbc: Whether periodic boundary conditions apply (``box_vectors`` used - only when ``True``). - max_num_pairs: Buffer size for the C++ kernel (``-1`` = all pairs). - """ - - def __init__(self, cutoff, pbc=True, max_num_pairs: int = -1): - super().__init__() - self.cutoff = cutoff - self.pbc = pbc - self.max_num_pairs = max_num_pairs - - def forward(self, positions, cell): - """Return neighbor pairs for ``positions`` under the given cell. - - Args: - positions: Atom positions ``(N, 3)``. - cell: Box vectors ``(3, 3)`` defining the simulation cell. - - Returns: - The neighbor-pair output of - :func:`molix.F.locality.get_neighbor_pairs`. - """ - box = cell if self.pbc else None - return F.get_neighbor_pairs( - positions, - self.cutoff, - max_num_pairs=self.max_num_pairs, - box_vectors=box, - ) - - def extra_repr(self): - """Render constructor args for ``repr(module)``.""" - return f"cutoff={self.cutoff}, pbc={self.pbc}, max_num_pairs={self.max_num_pairs}" - - -__all__ = ["NeighborList"] diff --git a/src/molix/nn/mlp.py b/src/molix/nn/mlp.py index 12063f1..3d1e667 100644 --- a/src/molix/nn/mlp.py +++ b/src/molix/nn/mlp.py @@ -5,6 +5,8 @@ import torch.nn as nn from pydantic import BaseModel, Field +from molix.config import config + Key = str | tuple[str, ...] @@ -62,9 +64,16 @@ def __init__( } act_fn = activation_map[self.config.activation.lower()] + ftype = config.ftype + layers: list[nn.Module] = [] layers.append( - nn.Linear(self.config.in_dim, self.config.hidden_dims[0], bias=self.config.use_bias) + nn.Linear( + self.config.in_dim, + self.config.hidden_dims[0], + bias=self.config.use_bias, + dtype=ftype, + ) ) layers.append(act_fn) @@ -74,12 +83,18 @@ def __init__( self.config.hidden_dims[idx], self.config.hidden_dims[idx + 1], bias=self.config.use_bias, + dtype=ftype, ) ) layers.append(act_fn) layers.append( - nn.Linear(self.config.hidden_dims[-1], self.config.out_dim, bias=self.config.use_bias) + nn.Linear( + self.config.hidden_dims[-1], + self.config.out_dim, + bias=self.config.use_bias, + dtype=ftype, + ) ) self.mlp = nn.Sequential(*layers) diff --git a/src/molix/profiler/__init__.py b/src/molix/profiler/__init__.py index 7420c38..89b710f 100644 --- a/src/molix/profiler/__init__.py +++ b/src/molix/profiler/__init__.py @@ -3,7 +3,7 @@ Standalone, OOP profiling tools for identifying performance bottlenecks **before and outside** of training. No Trainer, no hooks. -Three profilers, each targeting a single component: +Five profilers, each targeting a single component: - :class:`TaskProfiler` — wall-clock timing of a single pipeline task (:class:`~molix.data.task.SampleTask`, :class:`~molix.data.task.DatasetTask`, @@ -14,6 +14,10 @@ - :class:`DataLoaderProfiler` — DataLoader stall-time measurement. +- :class:`DatasetProfiler` — dataset characterisation: exact packed-pointer + size distributions, per-sample access latency and byte footprint, field + layout and per-target value statistics. + - :class:`TrainerProfiler` — Trainer-loop per-step framework overhead. Drives a real Trainer over a near-zero-compute :class:`~molix.profiler.mock.MockModel` so a ``cProfile`` window @@ -64,25 +68,27 @@ """ from molix.profiler.dataloader import DataLoaderProfiler, DataLoaderResult +from molix.profiler.dataset import DatasetProfiler, DatasetResult from molix.profiler.mock import MockBatch, MockModel, MockSource, mock_node_feature_loss from molix.profiler.module import ModuleProfiler, ModuleResult from molix.profiler.task import TaskProfiler, TaskResult from molix.profiler.trainer import TrainerProfiler, TrainerResult +# Profilers, their result types and the mock data generators share one +# alphabetically sorted list (notes.md 2026-08-09: `__all__` stays alphabetized). __all__ = [ - # Profilers - "TaskProfiler", - "ModuleProfiler", "DataLoaderProfiler", - "TrainerProfiler", - # Results - "TaskResult", - "ModuleResult", "DataLoaderResult", - "TrainerResult", - # Data generators + "DatasetProfiler", + "DatasetResult", "MockBatch", - "MockSource", "MockModel", + "MockSource", + "ModuleProfiler", + "ModuleResult", + "TaskProfiler", + "TaskResult", + "TrainerProfiler", + "TrainerResult", "mock_node_feature_loss", ] diff --git a/src/molix/profiler/_utils.py b/src/molix/profiler/_utils.py index f5493af..1283ebc 100644 --- a/src/molix/profiler/_utils.py +++ b/src/molix/profiler/_utils.py @@ -129,6 +129,62 @@ def from_list(cls, values: list[float | int]) -> "ValueStat": ) +# --------------------------------------------------------------------------- +# Count extraction — one function per tier of the two-tier data contract +# --------------------------------------------------------------------------- + + +def batch_counts(batch: object) -> tuple[int, int]: + """Extract ``(n_atoms, n_graphs)`` from a **post-collate** nested batch. + + Reads the collated tier of the two-tier data contract: a nested + ``TensorDict`` addressed by namespace, ``batch["atoms"]["Z"]`` ``(N,)`` + and ``batch["graphs"]["num_atoms"]`` ``(B,)``. For the raw flat sample + tier use :func:`sample_counts`. + + Args: + batch: Post-collate batch, normally a nested ``TensorDict``. + + Returns: + ``(n_atoms, n_graphs)``, or ``(0, 0)`` if the batch does not carry + those namespaces — this is a diagnostic helper and never raises. + """ + try: + n_atoms = int(batch["atoms"]["Z"].shape[0]) + n_graphs = int(batch["graphs"]["num_atoms"].shape[0]) + return n_atoms, n_graphs + except (KeyError, AttributeError, TypeError): + return 0, 0 + + +def sample_counts(sample: object) -> tuple[int, int]: + """Extract ``(n_atoms, n_edges)`` from one **raw** flat sample dict. + + Reads the pre-collate tier of the two-tier data contract: a flat + ``dict`` with ``sample["Z"]`` ``(N,)`` and ``sample["edge_index"]`` + ``(E, 2)``. The two keys are probed independently because + ``edge_index`` is genuinely optional — a cache packed before + :class:`~molix.data.tasks.NeighborList` ran has atoms but no edges. + For the collated nested tier use :func:`batch_counts`. + + Args: + sample: Raw sample, normally a flat ``dict`` of tensors. + + Returns: + ``(n_atoms, n_edges)``; a missing or malformed key contributes + ``0`` — this is a diagnostic helper and never raises. + """ + try: + n_atoms = int(sample["Z"].shape[0]) + except (KeyError, AttributeError, TypeError, IndexError): + n_atoms = 0 + try: + n_edges = int(sample["edge_index"].shape[0]) + except (KeyError, AttributeError, TypeError, IndexError): + n_edges = 0 + return n_atoms, n_edges + + # --------------------------------------------------------------------------- # ASCII table formatter # --------------------------------------------------------------------------- diff --git a/src/molix/profiler/dataloader.py b/src/molix/profiler/dataloader.py index 6fece0c..256d87e 100644 --- a/src/molix/profiler/dataloader.py +++ b/src/molix/profiler/dataloader.py @@ -32,7 +32,7 @@ from molix.data.collate import DEFAULT_TARGET_SCHEMA, TargetSchema, collate_molecules from molix.data.dataset import CachedDataset from molix.data.pipeline import PipelineSpec -from molix.profiler._utils import TimingStat, ValueStat, _fmt_table +from molix.profiler._utils import TimingStat, ValueStat, _fmt_table, batch_counts # --------------------------------------------------------------------------- # Result @@ -121,7 +121,7 @@ def print_report(self) -> None: # --------------------------------------------------------------------------- -# Batch stats extraction +# Collate # --------------------------------------------------------------------------- @@ -148,16 +148,6 @@ def __call__(self, batch_samples: list[dict]) -> object: return batch -def _extract_batch_counts(batch: object) -> tuple[int, int]: - """Extract (n_atoms, n_graphs) from a TensorDict batch.""" - try: - n_atoms = int(batch["atoms"]["Z"].shape[0]) # type: ignore[index] - n_graphs = int(batch["graphs"]["num_atoms"].shape[0]) # type: ignore[index] - return n_atoms, n_graphs - except (KeyError, AttributeError, TypeError): - return 0, 0 - - # --------------------------------------------------------------------------- # Profiler # --------------------------------------------------------------------------- @@ -256,7 +246,7 @@ def run( load_ms = (time.perf_counter() - t0) * 1000 if i >= n_warmup: load_times_ms.append(load_ms) - n_a, n_g = _extract_batch_counts(batch) + n_a, n_g = batch_counts(batch) atom_counts.append(n_a) graph_counts.append(n_g) if i + 1 >= total: @@ -301,8 +291,8 @@ def _source_to_dataset(self, source: object) -> Dataset: from molix.data.cache import PackedCache - n = len(source) # type: ignore[arg-type] - samples = [source[i] for i in range(n)] # type: ignore[index] + n = len(source) + samples = [source[i] for i in range(n)] tmp_file = Path(tempfile.mkdtemp(prefix="molix_profiler_")) / "samples.pt" PackedCache(tmp_file).save(samples) return CachedDataset(tmp_file) diff --git a/src/molix/profiler/dataset.py b/src/molix/profiler/dataset.py new file mode 100644 index 0000000..f6a3931 --- /dev/null +++ b/src/molix/profiler/dataset.py @@ -0,0 +1,790 @@ +"""Dataset characterisation profiler. + +Answers *"what is in my data, and what does one sample cost?"* for any object +yielding **flat sample dicts** (the raw-sample tier of the two-tier data +contract): :class:`~molix.data.dataset.CachedDataset` / +:class:`~molix.data.dataset.MmapDataset` / +:class:`~molix.data.dataset.SubsetDataset`, a +:class:`~molix.profiler.mock.MockSource`, or a plain ``Sequence[dict]``. +Complements :class:`~molix.profiler.dataloader.DataLoaderProfiler`, which +times batch *throughput* rather than the data itself. + +:meth:`DatasetProfiler.run` merges two paths: + +1. **Exact fast path** — when the object exposes the packed-pointer + properties (``atom_counts`` / ``edge_counts`` / ``avg_num_neighbors`` / + ``max_atoms`` / ``max_edges``), size statistics are read straight off the + ``atom_ptr`` / ``edge_ptr`` cumsums: **every** record, zero unpacking. + The result then carries ``counts_exact=True``. +2. **Sampled slow path** — only for what genuinely needs per-sample reads: + cold / steady-state ``__getitem__`` latency, per-sample byte footprint, + and target-value statistics. ``n_samples`` records are visited on a + ``stride``. Objects without packed pointers get their size statistics + here too, and report ``counts_exact=False``. + +Field layout comes from the packed ``payload["schema"]`` when available +(``fields_exact=True``), else it is inferred from the sampled records. + +Diagnostics — non-finite targets, atom-count skew, a missing ``edge_ptr``, +any inexact section — are emitted as ``[WARN]`` lines and **never** raise; +only degenerate *inputs* raise :class:`ValueError`. No unit conversion is +performed and no units are guessed: positions, energies and every target +are printed exactly as the dataset stores them. + +Example:: + + from molix.profiler import DatasetProfiler + + result = DatasetProfiler(n_samples=500).run(train_dataset) + result.print_report() + result.avg_num_neighbors # exact, straight off the packed pointers +""" + +from __future__ import annotations + +import math +from collections.abc import Mapping +from dataclasses import dataclass, field + +import torch + +from molix.profiler._utils import Timer, TimingStat, ValueStat, _fmt_table, sample_counts + +#: Atom-count skew — ``p95 / p50`` over the records — above which the report +#: warns that padded batches will waste compute. Documented as the "3×" +#: threshold in ``docs/molix/user-guide/profiling.md``. +_ATOM_SKEW_WARN_RATIO = 3.0 + +# --------------------------------------------------------------------------- +# Field / target descriptions +# --------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class FieldSpec: + """Packed layout of one sample key, as shown in the report. + + This is the **report view** of one entry of a packed cache's + ``payload["schema"]`` — not to be confused with + :class:`molix.data.collate.TargetSchema`, which is a collate-time + routing rule. ``FieldSpec`` describes storage, ``TargetSchema`` + describes destination. + + Attributes: + key: Dotted key path into the flat sample, e.g. ``"targets.U0"``. + axis: Packing axis — ``"atom"``, ``"edge"``, ``"graph"`` or + ``"scalar"``, matching :func:`molix.data.cache._infer_schema_across`. + dtype: ``torch.dtype`` of the tensor, or the Python type name for a + non-tensor scalar. + extra_shape: Trailing shape after the packing axis — ``(3,)`` for + ``pos`` ``(N, 3)``, ``()`` for ``Z`` ``(N,)``. ``"graph"`` fields + keep their full per-sample shape; ``"scalar"`` fields use ``()``. + """ + + key: str + axis: str + dtype: torch.dtype | str + extra_shape: tuple[int, ...] + + @classmethod + def from_packed(cls, key: str, spec: tuple) -> "FieldSpec": + """Build the report view of one packed ``payload["schema"]`` entry. + + Args: + key: Dotted key the entry describes. + spec: Schema entry — ``(axis, dtype, extra_shape)`` for packed + tensors, or the two-element ``("scalar", type_name)`` form. + + Returns: + The corresponding :class:`FieldSpec`. + """ + return cls( + key=key, + axis=spec[0], + dtype=spec[1], + extra_shape=tuple(spec[2]) if len(spec) > 2 else (), + ) + + @classmethod + def from_sample(cls, key: str, value: object, n_atoms: int, n_edges: int) -> "FieldSpec": + """Infer the layout of one leaf from a single sampled record. + + Best-effort fallback for datasets with no packed view; classification + mirrors :func:`molix.data.cache._infer_schema_across` (leading + dimension tracks atoms, else edges, else the field is per-graph) but + sees one record instead of all of them — the caller reports + ``fields_exact=False``. + + Args: + key: Dotted key of the leaf. + value: The leaf value, normally a tensor. + n_atoms: Atom count of the record this leaf came from. + n_edges: Edge count of the record this leaf came from. + + Returns: + The inferred :class:`FieldSpec`. + """ + if not isinstance(value, torch.Tensor): + return cls(key=key, axis="scalar", dtype=type(value).__name__, extra_shape=()) + if value.ndim and int(value.shape[0]) == n_atoms: + return cls(key=key, axis="atom", dtype=value.dtype, extra_shape=tuple(value.shape[1:])) + if value.ndim and int(value.shape[0]) == n_edges: + return cls(key=key, axis="edge", dtype=value.dtype, extra_shape=tuple(value.shape[1:])) + return cls(key=key, axis="graph", dtype=value.dtype, extra_shape=tuple(value.shape)) + + +@dataclass(frozen=True) +class TargetStat: + """Numeric distribution of one label column across the sampled records. + + Unrelated to :class:`molix.data.collate.TargetSchema`: this is a + *measurement* of the label values, whereas ``TargetSchema`` is the rule + that routes labels into batch namespaces at collate time. + + Attributes: + key: Dotted key of the label, e.g. ``"targets.U0"``. + stat: Mean / std / p50 / p95 over the **finite** sampled values. + min: Smallest finite value seen (``nan`` if none was). + max: Largest finite value seen (``nan`` if none was). + n_nonfinite: How many sampled values were ``nan`` or ``inf``. + """ + + key: str + stat: ValueStat + min: float + max: float + n_nonfinite: int + + @classmethod + def from_values(cls, key: str, values: list[float]) -> "TargetStat": + """Summarise one label column, tolerating non-finite entries. + + Args: + key: Dotted key of the label column. + values: Sampled scalar values, possibly containing ``nan`` / + ``inf`` — those are counted, then excluded from the moments + so a single bad row cannot poison the whole column. + + Returns: + The corresponding :class:`TargetStat`. + """ + finite = [v for v in values if math.isfinite(v)] + empty = ValueStat(mean=math.nan, std=math.nan, p50=math.nan, p95=math.nan) + return cls( + key=key, + stat=ValueStat.from_list(finite) if finite else empty, + min=min(finite) if finite else math.nan, + max=max(finite) if finite else math.nan, + n_nonfinite=len(values) - len(finite), + ) + + +# --------------------------------------------------------------------------- +# Result +# --------------------------------------------------------------------------- + + +@dataclass +class DatasetResult: + """Profiling results for a dataset. + + Attributes: + n_total: Records in the dataset (``len(data)``). + n_sampled: Records actually read by the slow path. + counts_exact: Size statistics came from the packed pointers (all + records) rather than the sampled subset. + fields_exact: Field layout came from the packed ``payload["schema"]`` + rather than sampled inference. + atom_stats: Atoms per record. + edge_stats: Edges per record, or ``None`` when the dataset carries no + edges (see :attr:`warnings`). + max_atoms: Largest single-record atom count. + max_edges: Largest single-record edge count (``0`` without edges). + avg_num_neighbors: ``total_edges / total_atoms``, the MACE/Allegro + normalisation constant. Exact iff :attr:`counts_exact`. + access_ms: Steady-state ``__getitem__`` latency. + cold_access_ms: Latency of the very first access (page-in / mmap + fault included). + sample_bytes: Leaf-tensor footprint of one record, in bytes. + est_total_mb: ``sample_bytes.mean * n_total`` in MB — what a full + in-RAM materialisation would cost. + fields: Per-key packed layout. + targets: Per-label value statistics. + task_states: ``{task_name: state}`` restored from the cache, i.e. + which fitted pipeline tasks (e.g. ``AtomicDress``) baked it. + warnings: Diagnostic lines rendered as ``[WARN]`` by + :meth:`print_report`. + data_description: Human-readable description of the input. + """ + + n_total: int + n_sampled: int + counts_exact: bool + fields_exact: bool + atom_stats: ValueStat + edge_stats: ValueStat | None + max_atoms: int + max_edges: int + avg_num_neighbors: float + access_ms: TimingStat + cold_access_ms: float + sample_bytes: ValueStat + est_total_mb: float + fields: list[FieldSpec] = field(default_factory=list) + targets: list[TargetStat] = field(default_factory=list) + task_states: dict[str, object] = field(default_factory=dict) + warnings: list[str] = field(default_factory=list) + data_description: str = "" + + def print_report(self) -> None: + """Print a sectioned dataset characterisation report to stdout.""" + print(f"\nDataset Profile (n_total={self.n_total:,}, n_sampled={self.n_sampled:,})") + print(f"Data: {self.data_description}") + print("─" * 72) + self._print_size() + self._print_access() + self._print_footprint() + self._print_fields() + self._print_targets() + self._print_task_states() + print("─" * 72) + self._print_warnings() + + def _print_size(self) -> None: + """Print the ``Size`` section: atom / edge counts per record.""" + print(" Size") + size_rows = [ + { + "Quantity": "atoms / sample", + "mean": f"{self.atom_stats.mean:.2f}", + "std": f"{self.atom_stats.std:.2f}", + "p50": f"{self.atom_stats.p50:.0f}", + "p95": f"{self.atom_stats.p95:.0f}", + "max": str(self.max_atoms), + } + ] + if self.edge_stats is not None: + size_rows.append( + { + "Quantity": "edges / sample", + "mean": f"{self.edge_stats.mean:.2f}", + "std": f"{self.edge_stats.std:.2f}", + "p50": f"{self.edge_stats.p50:.0f}", + "p95": f"{self.edge_stats.p95:.0f}", + "max": str(self.max_edges), + } + ) + print(_fmt_table(size_rows, ["Quantity", "mean", "std", "p50", "p95", "max"], col_width=8)) + print( + f" avg_num_neighbors (E/N): {self.avg_num_neighbors:.4f}" + f" (exact={self.counts_exact})" + ) + print() + + def _print_access(self) -> None: + """Print the ``Access`` section: steady-state and cold ``__getitem__`` latency.""" + print(" Access") + s = self.access_ms + access_rows = [ + { + "Metric": "__getitem__", + "mean(ms)": f"{s.mean_ms:.4f}", + "std(ms)": f"{s.std_ms:.4f}", + "p50(ms)": f"{s.p50_ms:.4f}", + "p95(ms)": f"{s.p95_ms:.4f}", + "cold(ms)": f"{self.cold_access_ms:.4f}", + } + ] + cols = ["Metric", "mean(ms)", "std(ms)", "p50(ms)", "p95(ms)", "cold(ms)"] + print(_fmt_table(access_rows, cols, col_width=8)) + print() + + def _print_footprint(self) -> None: + """Print the ``Footprint`` section: per-record bytes and the full-set estimate.""" + print(" Footprint") + print( + f" {self.sample_bytes.mean / 1e6:.6f} MB / sample" + f" (~{self.est_total_mb:,.1f} MB for all {self.n_total:,} records)" + ) + print() + + def _print_fields(self) -> None: + """Print the ``Fields`` section: packed layout of every sample key.""" + print(f" Fields (exact={self.fields_exact})") + field_rows = [ + { + "Key": f.key, + "axis": f.axis, + "dtype": str(f.dtype).removeprefix("torch."), + "extra_shape": str(tuple(f.extra_shape)), + } + for f in self.fields + ] + print(_fmt_table(field_rows, ["Key", "axis", "dtype", "extra_shape"], col_width=10)) + print() + + def _print_targets(self) -> None: + """Print the ``Targets`` section: value distribution of every label column.""" + print(" Targets") + if self.targets: + target_rows = [ + { + "Target": t.key, + "mean": f"{t.stat.mean:.4g}", + "std": f"{t.stat.std:.4g}", + "min": f"{t.min:.4g}", + "max": f"{t.max:.4g}", + "nonfinite": str(t.n_nonfinite), + } + for t in self.targets + ] + cols = ["Target", "mean", "std", "min", "max", "nonfinite"] + print(_fmt_table(target_rows, cols, col_width=8)) + else: + print(" (none)") + print() + + def _print_task_states(self) -> None: + """Print the fitted ``DatasetTask`` names, when the cache restored any.""" + if self.task_states: + print(f" Fitted task states: {', '.join(sorted(self.task_states))}") + print() + + def _print_warnings(self) -> None: + """Print the diagnostics collected during the run as ``[WARN]`` lines.""" + for message in self.warnings: + print(f" [WARN] {message}") + print() + + +# --------------------------------------------------------------------------- +# Sample helpers +# --------------------------------------------------------------------------- + + +def _flatten_leaves(sample: Mapping[str, object], prefix: str = "") -> dict[str, object]: + """Flatten a nested sample dict to dotted keys, read-only and never raising. + + Deliberately diverges from :func:`molix.data.cache._flatten`, the + packing-time flattener: that one is a **validator** and raises + :class:`ValueError` when a sample key collides with a reserved + packed-cache key (``schema``, ``atom_ptr``, …). This walker performs + **no** reserved-key validation and raises nothing, because the profiler + must survive any malformed sample — diagnostics are ``[WARN]`` lines, + never exceptions. + + Args: + sample: Sample dict, possibly with nested sub-dicts (``targets``). + prefix: Dotted prefix accumulated during recursion. + + Returns: + ``{dotted_key: leaf}``; leaf values are returned by reference, not + copied. + """ + leaves: dict[str, object] = {} + for name, value in sample.items(): + key = f"{prefix}{name}" + if isinstance(value, Mapping): + leaves.update(_flatten_leaves(value, prefix=f"{key}.")) + else: + leaves[key] = value + return leaves + + +def _sample_bytes(leaves: Mapping[str, object]) -> int: + """Total leaf-tensor footprint of one sample, in bytes. + + Args: + leaves: Flattened sample, as returned by :func:`_flatten_leaves`. + + Returns: + ``Σ numel() * element_size()`` over tensor leaves; non-tensor leaves + contribute nothing (their Python overhead is not the dataset's cost). + """ + return sum( + value.numel() * value.element_size() + for value in leaves.values() + if isinstance(value, torch.Tensor) + ) + + +# --------------------------------------------------------------------------- +# Profiler +# --------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class _SampledPass: + """What one pass over the sampled records measured. + + Everything here comes from the slow path — ``len(access_ms)`` records were + actually read. The size counts are only *used* when the packed pointers of + the exact fast path are unavailable. + + Attributes: + cold_access_ms: Latency of the very first access, page-in included. + access_ms: Steady-state per-record ``__getitem__`` latency, in ms. + byte_counts: Leaf-tensor footprint of each visited record, in bytes. + atom_counts: Atoms per visited record. + edge_counts: Edges per visited record. + fields: ``{dotted_key: FieldSpec}`` inferred from the first record that + carried each key. + target_values: ``{dotted_key: values}`` for every scalar ``targets.*`` + leaf seen. + """ + + cold_access_ms: float + access_ms: list[float] + byte_counts: list[int] + atom_counts: list[int] + edge_counts: list[int] + fields: dict[str, FieldSpec] + target_values: dict[str, list[float]] + + +@dataclass(frozen=True) +class _SizeStats: + """Record-size statistics, from the packed pointers or the sampled records. + + Attributes: + counts_exact: Statistics cover every record (packed pointers) rather + than the sampled subset. + atom_stats: Atoms per record. + edge_stats: Edges per record, or ``None`` when there are no edges. + max_atoms: Largest single-record atom count. + max_edges: Largest single-record edge count (``0`` without edges). + avg_num_neighbors: ``total_edges / total_atoms``. + """ + + counts_exact: bool + atom_stats: ValueStat + edge_stats: ValueStat | None + max_atoms: int + max_edges: int + avg_num_neighbors: float + + +class DatasetProfiler: + """Profile a dataset's sizes, access cost, field layout and labels. + + Args: + n_samples: Upper bound on records read by the sampled slow path. + stride: Step between inspected indices — ``stride > 1`` spreads the + sample across an ordered dataset instead of reading a prefix. + n_warmup: Accesses discarded before the timed ones, so page-in and + allocator warm-up do not land in :attr:`DatasetResult.access_ms` + (they show up in ``cold_access_ms`` instead). + + Raises: + ValueError: ``n_samples`` or ``stride`` is not positive. + + Example:: + + profiler = DatasetProfiler(n_samples=500, stride=4) + profiler.run(cached_dataset).print_report() + profiler.run([sample_a, sample_b]).print_report() # plain Sequence + """ + + def __init__(self, n_samples: int = 200, stride: int = 1, n_warmup: int = 3) -> None: + if n_samples <= 0: + raise ValueError( + f"n_samples must be > 0, got {n_samples}. Pass the number of records to " + "inspect (default 200); the sampled path needs at least one read." + ) + if stride <= 0: + raise ValueError( + f"stride must be > 0, got {stride}. Pass stride=1 (default) to walk " + "consecutive records, or a larger step to spread the sample out." + ) + self.n_samples = n_samples + self.stride = stride + self.n_warmup = max(0, n_warmup) + + def run(self, data: object) -> DatasetResult: + """Profile *data* and return the grouped result. + + Args: + data: Any object with ``__len__`` and ``__getitem__`` returning + flat sample dicts. Packed-cache-backed datasets additionally + get exact, whole-dataset size statistics and an exact field + layout for free. + + Returns: + :class:`DatasetResult`. Everything diagnostic — non-finite + labels, size skew, missing ``edge_ptr``, inexact sections — lands + in :attr:`DatasetResult.warnings`, not in an exception. + + Raises: + ValueError: *data* is empty, or exposes no ``__len__`` / + ``__getitem__``. + """ + n_total = self._require_indexable(data) + indices = list(range(0, n_total, self.stride))[: self.n_samples] + + sampled = self._sampled_pass(data, indices) + sizes, size_warnings = self._size_stats( + data, sampled.atom_counts, sampled.edge_counts, len(indices) + ) + fields, fields_exact, field_warnings = self._field_layout(data, sampled.fields) + targets, target_warnings = self._target_diagnostics(sampled.target_values, sizes.atom_stats) + + bytes_stat = ValueStat.from_list(sampled.byte_counts) + return DatasetResult( + n_total=n_total, + n_sampled=len(indices), + counts_exact=sizes.counts_exact, + fields_exact=fields_exact, + atom_stats=sizes.atom_stats, + edge_stats=sizes.edge_stats, + max_atoms=sizes.max_atoms, + max_edges=sizes.max_edges, + avg_num_neighbors=sizes.avg_num_neighbors, + access_ms=TimingStat.from_list(sampled.access_ms), + cold_access_ms=sampled.cold_access_ms, + sample_bytes=bytes_stat, + est_total_mb=bytes_stat.mean * n_total / 1e6, + fields=fields, + targets=targets, + task_states=dict(data.stats()) if hasattr(data, "stats") else {}, + warnings=[*size_warnings, *field_warnings, *target_warnings], + data_description=getattr(data, "describe", lambda: type(data).__name__)(), + ) + + # -- Stages ------------------------------------------------------------------- + + def _sampled_pass(self, data: object, indices: list[int]) -> _SampledPass: + """Read every index in *indices* once, measuring the sampled slow path. + + Args: + data: The dataset being profiled. + indices: Record indices to visit, in order. ``indices[0]`` is read + once beforehand for the cold measurement, then + :attr:`n_warmup` further reads are discarded, so page-in does + not land in the steady-state latency. + + Returns: + The measurements of the pass — latency, footprint, sampled sizes, + inferred field layout and label values. + """ + with Timer() as timer: + data[indices[0]] + cold_access_ms = timer.elapsed * 1000.0 + + for i in range(self.n_warmup): + data[indices[i % len(indices)]] + + access_ms: list[float] = [] + byte_counts: list[int] = [] + sampled_atoms: list[int] = [] + sampled_edges: list[int] = [] + sampled_fields: dict[str, FieldSpec] = {} + target_values: dict[str, list[float]] = {} + for idx in indices: + with Timer() as timer: + sample = data[idx] + access_ms.append(timer.elapsed * 1000.0) + + n_atoms, n_edges = sample_counts(sample) + sampled_atoms.append(n_atoms) + sampled_edges.append(n_edges) + + leaves = _flatten_leaves(sample) + byte_counts.append(_sample_bytes(leaves)) + for key, value in leaves.items(): + if key not in sampled_fields: + sampled_fields[key] = FieldSpec.from_sample(key, value, n_atoms, n_edges) + if not key.startswith("targets."): + continue + if isinstance(value, torch.Tensor) and value.numel() == 1: + target_values.setdefault(key, []).append(float(value.item())) + elif isinstance(value, (int, float)) and not isinstance(value, bool): + target_values.setdefault(key, []).append(float(value)) + + return _SampledPass( + cold_access_ms=cold_access_ms, + access_ms=access_ms, + byte_counts=byte_counts, + atom_counts=sampled_atoms, + edge_counts=sampled_edges, + fields=sampled_fields, + target_values=target_values, + ) + + def _size_stats( + self, + data: object, + sampled_atoms: list[int], + sampled_edges: list[int], + n_sampled: int, + ) -> tuple[_SizeStats, list[str]]: + """Summarise record sizes, preferring the exact packed-pointer fast path. + + Args: + data: The dataset being profiled. + sampled_atoms: Atoms per sampled record, used only as the fallback. + sampled_edges: Edges per sampled record, used only as the fallback. + n_sampled: How many records the sampled pass read, for the message. + + Returns: + ``(stats, warnings)``. ``warnings`` carries any missing-pointer + reason plus, on the fallback, the ``counts_exact=False`` notice. + """ + warnings: list[str] = [] + atom_counts, atom_warning = self._packed_counts(data, "atom_counts") + edge_counts, edge_warning = self._packed_counts(data, "edge_counts") + warnings.extend(w for w in (atom_warning, edge_warning) if w is not None) + counts_exact = atom_counts is not None + + if atom_counts is not None: + atom_stats = ValueStat.from_list(atom_counts.tolist()) + edge_stats = None if edge_counts is None else ValueStat.from_list(edge_counts.tolist()) + max_atoms = int(getattr(data, "max_atoms", 0)) + max_edges = int(getattr(data, "max_edges", 0)) if edge_counts is not None else 0 + avg_num_neighbors = float(getattr(data, "avg_num_neighbors", 0.0)) + else: + atom_stats = ValueStat.from_list(sampled_atoms) + edge_stats = ValueStat.from_list(sampled_edges) if any(sampled_edges) else None + max_atoms = max(sampled_atoms, default=0) + max_edges = max(sampled_edges, default=0) + total_atoms = sum(sampled_atoms) + avg_num_neighbors = sum(sampled_edges) / total_atoms if total_atoms else 0.0 + warnings.append( + f"size stats estimated from {n_sampled} sampled records (counts_exact=False)" + " — a packed-cache-backed dataset would give exact whole-dataset counts." + ) + + stats = _SizeStats( + counts_exact=counts_exact, + atom_stats=atom_stats, + edge_stats=edge_stats, + max_atoms=max_atoms, + max_edges=max_edges, + avg_num_neighbors=avg_num_neighbors, + ) + return stats, warnings + + def _field_layout( + self, data: object, sampled_fields: Mapping[str, FieldSpec] + ) -> tuple[list[FieldSpec], bool, list[str]]: + """Describe the packed layout of every key, exactly where possible. + + Args: + data: The dataset being profiled. + sampled_fields: Layout inferred during the sampled pass, used only + when *data* exposes no packed ``payload["schema"]``. + + Returns: + ``(fields, fields_exact, warnings)``. ``warnings`` carries the + inference notice when the packed schema was unreachable. + """ + schema = self._packed_schema(data) + if schema is not None: + fields = [FieldSpec.from_packed(key, spec) for key, spec in schema.items()] + return fields, True, [] + + fields = [sampled_fields[key] for key in sorted(sampled_fields)] + warning = ( + "field layout inferred from the sampled records (fields_exact=False)" + " — trailing shapes and axes may differ on unsampled records." + ) + return fields, False, [warning] + + def _target_diagnostics( + self, target_values: Mapping[str, list[float]], atom_stats: ValueStat + ) -> tuple[list[TargetStat], list[str]]: + """Summarise the label columns and flag distribution problems. + + Args: + target_values: ``{dotted_key: values}`` collected by the sampled pass. + atom_stats: Record-size statistics, checked for padding-wasting skew. + + Returns: + ``(targets, warnings)``. ``warnings`` covers non-finite label + values and an atom-count skew above :data:`_ATOM_SKEW_WARN_RATIO`. + """ + targets = [ + TargetStat.from_values(key, values) for key, values in sorted(target_values.items()) + ] + warnings = [ + f"target {t.key!r} has {t.n_nonfinite} non-finite value(s) in the sample" + for t in targets + if t.n_nonfinite + ] + if atom_stats.p50 > 0: + skew = atom_stats.p95 / atom_stats.p50 + if skew > _ATOM_SKEW_WARN_RATIO: + warnings.append( + f"atom-count skew p95/p50 = {skew:.1f}x" + " — padded batches will waste compute; consider a token-budget sampler." + ) + return targets, warnings + + # -- Helpers ---------------------------------------------------------------- + + def _require_indexable(self, data: object) -> int: + """Return ``len(data)``, rejecting inputs the sampled path cannot read. + + Args: + data: Candidate dataset. + + Returns: + Number of records. + + Raises: + ValueError: *data* is not a non-empty indexable sequence. + """ + remedy = ( + "Point it at a prepared cache (PipelineSpec.run(...) → CachedDataset) " + "or a non-empty Sequence[dict]." + ) + if not hasattr(data, "__len__") or not hasattr(data, "__getitem__"): + raise ValueError( + f"{type(data).__name__} exposes no __len__ / __getitem__, so no sample " + f"can be read. {remedy}" + ) + n_total = len(data) + if n_total == 0: + raise ValueError(f"Cannot profile an empty dataset ({type(data).__name__}). {remedy}") + return n_total + + def _packed_counts(self, data: object, attr: str) -> tuple[torch.Tensor | None, str | None]: + """Read an exact per-record count vector off a packed-pointer property. + + Args: + data: The dataset being profiled. + attr: Property name — ``"atom_counts"`` or ``"edge_counts"``. + + Returns: + ``(counts, warning)``. ``counts`` is ``None`` when the dataset is + not cache-backed (silent — the sampled path covers it) or when the + cache was packed without that pointer (a legitimate case, e.g. no + :class:`~molix.data.tasks.NeighborList` in the pipeline), in which + case ``warning`` carries the reason for the report. + """ + try: + counts = getattr(data, attr) + except AttributeError: + return None, None + except ValueError as exc: # cache built without the matching pointer + return None, f"exact {attr} unavailable — {exc}" + if not isinstance(counts, torch.Tensor): + return None, None + return counts, None + + def _packed_schema(self, data: object) -> Mapping[str, tuple] | None: + """Return the packed ``payload["schema"]`` mapping, or ``None``. + + Args: + data: The dataset being profiled. + + Returns: + The exact per-key ``(axis, dtype, extra_shape)`` mapping inferred + when the cache was packed, or ``None`` when *data* has no packed + view (``SubsetDataset`` over a non-cache dataset raises + :class:`AttributeError` here, which is the documented "no fast + path" signal). + """ + try: + return data.packed_view().payload["schema"] + except (AttributeError, KeyError, TypeError): + return None diff --git a/src/molix/profiler/mock.py b/src/molix/profiler/mock.py index cff38a0..8e2d88f 100644 --- a/src/molix/profiler/mock.py +++ b/src/molix/profiler/mock.py @@ -42,11 +42,14 @@ _IntOrRange = Union[int, tuple[int, int]] -def _resolve(value: _IntOrRange) -> int: +def _resolve(value: _IntOrRange, rng: random.Random) -> int: """Sample a concrete integer from a fixed value or (lo, hi) range. Args: value: Either a fixed ``int`` or a ``(lo, hi)`` inclusive range. + rng: Caller-owned generator to draw from. Passed explicitly rather + than read off the ``random`` module so a seeded owner + (:class:`MockBatch`, :class:`MockSource`) really is reproducible. Returns: A concrete integer. @@ -54,7 +57,7 @@ def _resolve(value: _IntOrRange) -> int: if isinstance(value, int): return value lo, hi = value - return random.randint(lo, hi) + return rng.randint(lo, hi) # --------------------------------------------------------------------------- @@ -113,9 +116,9 @@ def __call__(self) -> TensorDict: Returns: A ``TensorDict`` with random tensor values and the configured shape. """ - n_a = _resolve(self.n_atoms) - n_e = _resolve(self.n_edges) - n_g = _resolve(self.n_graphs) + n_a = _resolve(self.n_atoms, self._rng) + n_e = _resolve(self.n_edges, self._rng) + n_g = _resolve(self.n_graphs, self._rng) dev = self.device gen = self._torch_gen @@ -213,9 +216,7 @@ def __init__( self.atomic_numbers = atomic_numbers # Pre-generate atom counts for each sample so source_id is stable rng = random.Random(seed) - self._atom_counts: list[int] = [ - _resolve(n_atoms) if not isinstance(n_atoms, int) else n_atoms for _ in range(n_samples) - ] + self._atom_counts: list[int] = [_resolve(n_atoms, rng) for _ in range(n_samples)] # Per-sample generator seeds for reproducible, independent samples self._seeds: list[int] = [rng.randint(0, 2**31) for _ in range(n_samples)] diff --git a/src/molix/profiler/module.py b/src/molix/profiler/module.py index 829ba66..74498a3 100644 --- a/src/molix/profiler/module.py +++ b/src/molix/profiler/module.py @@ -29,7 +29,13 @@ import torch import torch.nn as nn -from molix.profiler._utils import TimingStat, ValueStat, _fmt_table, reset_peak_memory +from molix.profiler._utils import ( + TimingStat, + ValueStat, + _fmt_table, + batch_counts, + reset_peak_memory, +) # --------------------------------------------------------------------------- # Result @@ -205,16 +211,6 @@ def _move_to_device(batch: object, device: torch.device) -> object: return batch -def _extract_counts(batch: object) -> tuple[int, int]: - """Extract (n_atoms, n_graphs) from a TensorDict batch; returns (0, 0) on failure.""" - try: - n_atoms = int(batch["atoms"]["Z"].shape[0]) # type: ignore[index] - n_graphs = int(batch["graphs"]["num_atoms"].shape[0]) # type: ignore[index] - return n_atoms, n_graphs - except (KeyError, AttributeError, TypeError): - return 0, 0 - - def _make_batch_iter( data: object, n_total: int, @@ -234,7 +230,7 @@ def _make_batch_iter( return (data[i % len(data)] for i in range(n_total)) # Assume DataLoader or other iterable — wrap with cycling - def _cycle(iterable: Iterable, n: int): # type: ignore[return] + def _cycle(iterable: Iterable, n: int): buf: list = [] it = iter(iterable) count = 0 @@ -536,7 +532,7 @@ def run( bwd_times_ms.append(bwd_ms) opt_times_ms.append(opt_ms) peak_mem_mb.append(mem_mb) - n_a, n_g = _extract_counts(batch) + n_a, n_g = batch_counts(batch) atom_counts.append(n_a) graph_counts.append(n_g) @@ -676,9 +672,9 @@ def hook(_mod, _inp, _out): if not samples: continue if use_cuda: - ms = [s.elapsed_time(e) for _, s, e in samples] # type: ignore[not-iterable] + ms = [s.elapsed_time(e) for _, s, e in samples] else: - ms = [float(v) * 1000.0 for v in samples] # type: ignore[not-iterable] + ms = [float(v) * 1000.0 for v in samples] means[name] = sum(ms) / len(ms) if not means: return None diff --git a/src/molix/profiler/task.py b/src/molix/profiler/task.py index 487262a..5c55754 100644 --- a/src/molix/profiler/task.py +++ b/src/molix/profiler/task.py @@ -124,11 +124,11 @@ def run( Returns: :class:`TaskResult` with timing statistics. """ - n_source = len(source) # type: ignore[arg-type] + n_source = len(source) # For DatasetTask: fit on the full source first if isinstance(self.task, DatasetTask): - all_samples = [source[i] for i in range(n_source)] # type: ignore[index] + all_samples = [source[i] for i in range(n_source)] self.task.fit(all_samples) task_name = type(self.task).__name__ @@ -140,7 +140,7 @@ def run( for i in range(total): idx = i % n_source - sample = source[idx] # type: ignore[index] + sample = source[idx] with Timer() as t: entry.apply(sample) if i >= n_warmup: diff --git a/src/molix/profiler/trainer.py b/src/molix/profiler/trainer.py index fa93cbc..5fc4678 100644 --- a/src/molix/profiler/trainer.py +++ b/src/molix/profiler/trainer.py @@ -51,7 +51,9 @@ class TrainerResult: (Trainer self-time minus raw-loop baseline) / ``gross_us`` (raw Trainer self-time) / ``calls``. model_name: ``type(model).__name__``. - data_description: Human-readable batch description. + data_description: Human-readable batch description — the + :class:`~molix.profiler.mock.MockBatch` config when the default + batch was built, else the supplied batch's namespace shape. baseline_ms_per_step: Raw-loop (no Trainer) wall per step. overhead_ms_per_step: ``wall_ms_per_step - baseline_ms_per_step``. """ @@ -168,8 +170,18 @@ def run( A :class:`TrainerResult`. """ if batch is None: - batch = MockBatch(n_atoms=32, n_edges=128, n_graphs=4, device=str(self.device))() - desc = "MockBatch(n_atoms=32, n_edges=128, n_graphs=4)" + factory = MockBatch(n_atoms=32, n_edges=128, n_graphs=4, device=str(self.device)) + batch = factory() + desc = factory.describe() + else: + # Caller-supplied batch: label it from the per-namespace batch_size + # of the post-collate contract (``graphs`` is optional there). + shape = [ + f"{tag}={batch[ns].batch_size[0]}" + for tag, ns in (("N", "atoms"), ("E", "edges"), ("B", "graphs")) + if ns in batch + ] + desc = f"TensorDict({', '.join(shape)})" # Warm up both the raw-loop baseline and the Trainer path. if n_warmup > 0: diff --git a/src/molix/quant.py b/src/molix/quant.py index 2d74545..9a37f4e 100644 --- a/src/molix/quant.py +++ b/src/molix/quant.py @@ -31,6 +31,9 @@ import torch.nn as nn import torch.nn.utils.parametrize as parametrize +from molix.schema import FORCES_KEY +from molix.units import KB_EV_PER_K + class QuantScheme(ABC): """Abstract fake-quantization strategy (quantize-then-dequantize to float). @@ -252,8 +255,10 @@ def __init__(self, delta_f: torch.Tensor): @classmethod def between(cls, model_ref: nn.Module, model_quant: nn.Module, batch: object) -> ForceDelta: """ΔF from a reference vs quantized model on a fresh clone of ``batch`` each.""" - f_ref = model_ref(batch.clone(), compute_forces=True)["forces"].detach() # type: ignore[union-attr] - f_quant = model_quant(batch.clone(), compute_forces=True)["forces"].detach() # type: ignore[union-attr] + # Potentials fix force derivation at construction (monomorphic forward, + # since b85d12f); both models must already be built with it. + f_ref = model_ref(batch.clone())[FORCES_KEY].detach() + f_quant = model_quant(batch.clone())[FORCES_KEY].detach() return cls(f_quant - f_ref) def summary(self) -> dict[str, float]: @@ -289,8 +294,8 @@ class EffectiveTemperature: for the dimensionless reported quantity. """ - #: Boltzmann constant in eV/K - KB_EV_PER_K: float = 8.617333262e-5 + #: Boltzmann constant in eV/K (single source: :mod:`molix.units`). + KB_EV_PER_K: float = KB_EV_PER_K def __init__(self, *, dt: float, gamma: float, mass: float, dof: int): self.dt = dt diff --git a/src/molix/schema.py b/src/molix/schema.py new file mode 100644 index 0000000..34f1e97 --- /dev/null +++ b/src/molix/schema.py @@ -0,0 +1,36 @@ +"""Post-collate batch-schema keys shared across the molix stack. + +The nested ``atoms / edges / graphs`` TensorDict layout produced by +:func:`molix.data.collate.collate_molecules` is a molix contract (see +CLAUDE.md, "Post-collate batch schema"). The tuple keys under which models +write energy/force outputs live here so every layer — molix execution +utilities and the higher molpot/molzoo packages alike — addresses one schema +without molix ever importing upward. :mod:`molpot.derivation.protocol` +re-exports these names for potential-side code; molix-internal consumers +(:mod:`molix.md`, :mod:`molix.quant`, :mod:`molix.engine`) import them from +here, keeping the one-way package dependency ``molix ← molrep ← molzoo/molpot`` +intact. +""" + +from __future__ import annotations + +from tensordict import TensorDict + +#: Graph-level total energy written by a potential (``batch["graphs", "energy"]``). +ENERGY_KEY: tuple[str, str] = ("graphs", "energy") +#: Per-atom energy contributions (``batch["atoms", "energy"]``). +ATOMIC_ENERGY_KEY: tuple[str, str] = ("atoms", "energy") +#: Per-atom forces ``= -∂E/∂pos`` (``batch["atoms", "forces"]``). +FORCES_KEY: tuple[str, str] = ("atoms", "forces") +#: Per-atom positions (``batch["atoms", "pos"]``). +POS_KEY: tuple[str, str] = ("atoms", "pos") + + +def has_energy(batch: TensorDict) -> bool: + """Whether ``batch`` carries a graph-level energy at :data:`ENERGY_KEY`.""" + return "graphs" in batch.keys() and "energy" in batch["graphs"].keys() + + +def has_forces(batch: TensorDict) -> bool: + """Whether ``batch`` carries per-atom forces at :data:`FORCES_KEY`.""" + return "atoms" in batch.keys() and "forces" in batch["atoms"].keys() diff --git a/src/molix/units.py b/src/molix/units.py new file mode 100644 index 0000000..84586f4 --- /dev/null +++ b/src/molix/units.py @@ -0,0 +1,26 @@ +"""Physical constants and shared numerical conventions — the single source. + +Every module needing k_B or the (amu, Å, fs) energy bridge imports from here. +Restating a constant locally is exactly the drift this module exists to +prevent: ``KB_EV_PER_K`` used to live in two verbatim copies (``molix.md`` and +``molix.quant``) that could only diverge silently. +""" + +from __future__ import annotations + +#: Boltzmann constant in eV/K (CODATA 2018). +KB_EV_PER_K: float = 8.617333262e-5 + +#: Energy-unit bridge: 1 amu·Å²/fs² = 103.6426965638 eV. +EV_PER_AMU_A2_FS2: float = 103.6426965638 + +#: k_B in the MD integrator's (amu, Å, fs) energy unit (amu·Å²/fs²), so +#: ``T = 2·KE / (dof · KB_AMU_A_FS)`` comes out in kelvin. +KB_AMU_A_FS: float = KB_EV_PER_K / EV_PER_AMU_A2_FS2 + +#: Padded ("dead") neighbour edges are displaced to ``factor · r_cut`` so every +#: cutoff envelope evaluates to exactly zero. Shared by +#: :class:`molix.md.neighbors.NeighborList` and +#: :class:`molix.engine.static.StaticForward` so both padding paths silence +#: dead edges identically. +DEAD_EDGE_CUTOFF_FACTOR: float = 10.0 diff --git a/src/molpot/__init__.py b/src/molpot/__init__.py index 1ced9f1..1940281 100644 --- a/src/molpot/__init__.py +++ b/src/molpot/__init__.py @@ -6,16 +6,23 @@ # Potentials # Composition from molpot.composition import ( + KCAL_MOL_TO_EV, + AngleParamHead, + BondParamHead, ChargeHead, ChargeTransferParameterHead, + ClassicalMMComposer, + ClassicalMMParameterizer, + ImproperParamHead, LJParameterHead, MultiHead, PotentialComposer, + ProperTorsionParamHead, RepulsionParameterHead, Sonata, SonataSpec, TSScalingHead, - build_sonata, + energy_kcal_to_ev, ) # Physical derivation @@ -24,6 +31,21 @@ # Prediction heads from molpot.heads import AtomicEnergyMLP, EnergyHead, TypeHead +# Class-I Potential IR +from molpot.ir import ( + CLASS_I_CANONICAL, + AngleBag, + BondBag, + ChargeBag, + ImproperHarmonicBag, + ImproperPeriodicBag, + LJBag, + NonbondedScaling, + PotentialIR, + ProperTorsionBag, + UnitTag, +) + # Pooling from molpot.pooling import ( EdgeToNodePooling, @@ -40,7 +62,10 @@ ChargeTransfer, DihedralHarmonic, DispersionC6, + ImproperHarmonic, + ImproperPeriodic, Polarization, + ProperTorsionPeriodic, RepulsionExp6, geometric_arithmetic_mixing, lorentz_berthelot, @@ -54,11 +79,26 @@ "BondHarmonic", "AngleHarmonic", "DihedralHarmonic", + "ProperTorsionPeriodic", + "ImproperPeriodic", + "ImproperHarmonic", "RepulsionExp6", "DispersionC6", "ChargeTransfer", "Polarization", "geometric_arithmetic_mixing", + # Potential IR + "UnitTag", + "CLASS_I_CANONICAL", + "BondBag", + "AngleBag", + "ProperTorsionBag", + "ImproperPeriodicBag", + "ImproperHarmonicBag", + "LJBag", + "ChargeBag", + "NonbondedScaling", + "PotentialIR", # Heads "AtomicEnergyMLP", "EnergyHead", @@ -79,9 +119,16 @@ "ChargeTransferParameterHead", "ChargeHead", "TSScalingHead", + "BondParamHead", + "AngleParamHead", + "ProperTorsionParamHead", + "ImproperParamHead", "MultiHead", "PotentialComposer", + "ClassicalMMComposer", + "ClassicalMMParameterizer", + "KCAL_MOL_TO_EV", + "energy_kcal_to_ev", "Sonata", "SonataSpec", - "build_sonata", ] diff --git a/src/molpot/composition/__init__.py b/src/molpot/composition/__init__.py index a9b7121..b5b486f 100644 --- a/src/molpot/composition/__init__.py +++ b/src/molpot/composition/__init__.py @@ -11,8 +11,8 @@ outputs = composer(node_features=node_features, data=data) """ +from molpot.composition.classical_mm import ClassicalMMComposer from molpot.composition.composer import PotentialComposer -from molpot.composition.energy_force import EnergyForceModel from molpot.composition.heads import ( ChargeHead, ChargeTransferParameterHead, @@ -20,8 +20,21 @@ RepulsionParameterHead, TSScalingHead, ) +from molpot.composition.mm_heads import ( + AngleParamHead, + BondParamHead, + ImproperParamHead, + ProperTorsionParamHead, +) from molpot.composition.multihead import MultiHead -from molpot.composition.sonata import Sonata, SonataSpec, build_sonata +from molpot.composition.parameterizer import ( + KCAL_MOL_TO_EV, + ChemEmbeddingsLike, + ChemEncoderProtocol, + ClassicalMMParameterizer, + energy_kcal_to_ev, +) +from molpot.composition.sonata import Sonata, SonataSpec __all__ = [ "LJParameterHead", @@ -29,10 +42,18 @@ "ChargeTransferParameterHead", "ChargeHead", "TSScalingHead", + "BondParamHead", + "AngleParamHead", + "ProperTorsionParamHead", + "ImproperParamHead", "MultiHead", "PotentialComposer", - "EnergyForceModel", + "ClassicalMMComposer", + "ClassicalMMParameterizer", + "ChemEmbeddingsLike", + "ChemEncoderProtocol", + "KCAL_MOL_TO_EV", + "energy_kcal_to_ev", "Sonata", "SonataSpec", - "build_sonata", ] diff --git a/src/molpot/composition/classical_mm.py b/src/molpot/composition/classical_mm.py new file mode 100644 index 0000000..e463f80 --- /dev/null +++ b/src/molpot/composition/classical_mm.py @@ -0,0 +1,488 @@ +"""ClassicalMMComposer: features → MM heads → PotentialIR → Class-I energy. + +Does **not** own a chemical encoder. The caller injects feature tensors keyed +by interaction class (and optional atom features for LJ/charge). Sibling of +:class:`~molpot.composition.composer.PotentialComposer` specialized for +Class-I IR bags; does not subclass Sonata. + +Pipeline primitives (caller-composed):: + + ir = composer.parameterize(features, batch) + energy = composer.energy(ir, batch, pos=...) + # thin convenience: + out = composer(batch, features) # {"energy", "ir", "term_energies"?} + +Forces are **not** hand-rolled. Differentiate the energy path with +:class:`~molpot.derivation.ForceDerivation` / ``BasePotential.calc_forces``. + +Units: CLASS_I_CANONICAL (kcal/mol, Å, e, rad). + +References: + Spec: learnable-classical-ff-03-mm-heads + OpenMM User Guide §19 "Forces" +""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping +from typing import Any + +import torch +import torch.nn as nn +from tensordict import TensorDict + +from molpot.ir import ( + AngleBag, + BondBag, + ChargeBag, + ImproperHarmonicBag, + ImproperPeriodicBag, + LJBag, + NonbondedScaling, + PotentialIR, + ProperTorsionBag, +) +from molpot.potentials import ( + AngleHarmonic, + BondHarmonic, + ImproperHarmonic, + ImproperPeriodic, + ProperTorsionPeriodic, +) + +__all__ = ["ClassicalMMComposer"] + +# Evaluator: (ir, batch, *, pos) -> scalar energy +EnergyEvaluator = Callable[..., torch.Tensor] + + +def _get_namespace(batch: Any, name: str) -> Any | None: + """Fetch a top-level namespace from TensorDict / Mapping.""" + if batch is None: + return None + if isinstance(batch, Mapping) or isinstance(batch, TensorDict): + try: + if name in batch: + return batch[name] + except Exception: # noqa: BLE001 — TensorDict key miss variants + return None + return None + return None + + +def _ns_tensor(ns: Any, *keys: str) -> torch.Tensor | None: + """Read a tensor from a namespace under any of ``keys``.""" + if ns is None: + return None + for key in keys: + try: + if isinstance(ns, (Mapping, TensorDict)) and key in ns: + return ns[key] + except Exception: # noqa: BLE001 + continue + return None + + +def _bond_index_from_batch(batch: Any) -> torch.Tensor | None: + """Build COO ``bond_index`` ``[2, N]`` from column namespaces or packed keys.""" + bonds = _get_namespace(batch, "bonds") + if bonds is not None: + packed = _ns_tensor(bonds, "bond_index") + if packed is not None: + return packed + atomi = _ns_tensor(bonds, "atomi") + atomj = _ns_tensor(bonds, "atomj") + if atomi is not None and atomj is not None: + return torch.stack([atomi.long(), atomj.long()], dim=0) + # Flat legacy + if isinstance(batch, Mapping) and "bond_index" in batch: + return batch["bond_index"] + return None + + +def _angle_index_from_batch(batch: Any) -> torch.Tensor | None: + angles = _get_namespace(batch, "angles") + if angles is None: + return None + atomi = _ns_tensor(angles, "atomi") + atomj = _ns_tensor(angles, "atomj") + atomk = _ns_tensor(angles, "atomk") + if atomi is None or atomj is None or atomk is None: + return None + return torch.stack([atomi.long(), atomj.long(), atomk.long()], dim=0) + + +def _proper_index_from_batch(batch: Any) -> torch.Tensor | None: + propers = _get_namespace(batch, "propers") + if propers is None: + return None + cols = [_ns_tensor(propers, k) for k in ("atomi", "atomj", "atomk", "atoml")] + if any(c is None for c in cols): + return None + return torch.stack([c.long() for c in cols], dim=0) # type: ignore[union-attr] + + +def _improper_index_from_batch(batch: Any) -> torch.Tensor | None: + impropers = _get_namespace(batch, "impropers") + if impropers is None: + return None + cols = [_ns_tensor(impropers, k) for k in ("atomi", "atomj", "atomk", "atoml")] + if any(c is None for c in cols): + return None + # molrs center-first: atomi = center + return torch.stack([c.long() for c in cols], dim=0) # type: ignore[union-attr] + + +def _pos_from_batch(batch: Any, pos: torch.Tensor | None) -> torch.Tensor: + if pos is not None: + return pos + atoms = _get_namespace(batch, "atoms") + if atoms is not None: + p = _ns_tensor(atoms, "pos") + if p is not None: + return p + if isinstance(batch, Mapping) and "pos" in batch: + return batch["pos"] + raise ValueError("ClassicalMMComposer.energy requires pos or batch atoms.pos") + + +def _identity_types(n: int, device: torch.device) -> torch.Tensor: + return torch.arange(n, device=device, dtype=torch.long) + + +class ClassicalMMComposer(nn.Module): + """Wire continuous MM heads into Class-I :class:`~molpot.ir.PotentialIR`. + + Args: + bond_head: Optional :class:`~molpot.composition.mm_heads.BondParamHead` + (or compatible module) mapping bond features → ``k``, ``r0``. + angle_head: Optional angle parameter head. + proper_head: Optional proper-torsion parameter head. + improper_head: Optional improper parameter head (harmonic and/or + periodic outputs). + atom_head: Optional atom-level head (typically + :class:`~molpot.composition.multihead.MultiHead` of + :class:`~molpot.composition.heads.LJParameterHead` and + :class:`~molpot.composition.heads.ChargeHead`). + scaling: Optional nonbonded scaling; defaults to Class-I AMBER/GAFF. + evaluator: Optional callable ``(ir, batch, *, pos) -> energy``. When + set, :meth:`energy` delegates to it instead of built-in kernels. + + Notes: + Prefer :meth:`parameterize` and :meth:`energy` as separate primitives. + :meth:`forward` is a thin chain for ``nn.Module`` convenience only. + """ + + def __init__( + self, + *, + bond_head: nn.Module | None = None, + angle_head: nn.Module | None = None, + proper_head: nn.Module | None = None, + improper_head: nn.Module | None = None, + atom_head: nn.Module | None = None, + scaling: NonbondedScaling | None = None, + evaluator: EnergyEvaluator | None = None, + ) -> None: + super().__init__() + self.bond_head = bond_head + self.angle_head = angle_head + self.proper_head = proper_head + self.improper_head = improper_head + self.atom_head = atom_head + self.scaling = scaling if scaling is not None else NonbondedScaling() + self.evaluator = evaluator + + # ------------------------------------------------------------------ + # parameterize + # ------------------------------------------------------------------ + + def parameterize( + self, + features: Mapping[str, torch.Tensor], + batch: Any, + ) -> PotentialIR: + """Run enabled heads and assemble a Class-I :class:`PotentialIR`. + + Args: + features: Feature tensors keyed by interaction class. Supported + keys: ``"bonds"``, ``"angles"``, ``"propers"``, + ``"impropers"``, ``"atoms"``. Missing keys skip that bag. + batch: Topology batch (TensorDict or Mapping) with valence + namespaces. Used for atom ``batch`` index when charge + neutrality is required; topology counts are taken from + feature row counts. + + Returns: + :class:`PotentialIR` with ``unit_system="class_i_canonical"`` and + bags populated for enabled terms. Default + :class:`~molpot.ir.NonbondedScaling` is always attached. + """ + bonds = self._run_bond_head(features) + angles = self._run_angle_head(features) + propers = self._run_proper_head(features) + imp_h, imp_p = self._run_improper_head(features) + lj, charges = self._run_atom_head(features, batch) + + return PotentialIR( + bonds=bonds, + angles=angles, + propers=propers, + impropers_harmonic=imp_h, + impropers_periodic=imp_p, + lj=lj, + charges=charges, + scaling=self.scaling, + unit_system="class_i_canonical", + ) + + def _run_bond_head(self, features: Mapping[str, torch.Tensor]) -> BondBag | None: + if self.bond_head is None or "bonds" not in features: + return None + out = self.bond_head(features["bonds"]) + return BondBag(k=out["k"], r0=out["r0"]) + + def _run_angle_head(self, features: Mapping[str, torch.Tensor]) -> AngleBag | None: + if self.angle_head is None or "angles" not in features: + return None + out = self.angle_head(features["angles"]) + return AngleBag(k=out["k"], theta0=out["theta0"]) + + def _run_proper_head(self, features: Mapping[str, torch.Tensor]) -> ProperTorsionBag | None: + if self.proper_head is None or "propers" not in features: + return None + out = self.proper_head(features["propers"]) + return ProperTorsionBag( + k=out["k"], + periodicity=out["periodicity"], + phase=out["phase"], + idivf=out["idivf"], + ) + + def _run_improper_head( + self, features: Mapping[str, torch.Tensor] + ) -> tuple[ImproperHarmonicBag | None, ImproperPeriodicBag | None]: + if self.improper_head is None or "impropers" not in features: + return None, None + out = self.improper_head(features["impropers"]) + harm: ImproperHarmonicBag | None = None + peri: ImproperPeriodicBag | None = None + + if "chi0" in out: + k_h = out.get("k_harmonic", out.get("k")) + if k_h is not None and k_h.ndim == 1: + harm = ImproperHarmonicBag(k=k_h, chi0=out["chi0"]) + + k_p = out.get("k_periodic") + if k_p is None and "phase" in out and "periodicity" in out: + k_candidate = out.get("k") + if k_candidate is not None and k_candidate.ndim == 2: + k_p = k_candidate + if k_p is not None and "phase" in out and "periodicity" in out: + peri = ImproperPeriodicBag( + k=k_p, + periodicity=out["periodicity"], + phase=out["phase"], + idivf=out["idivf"], + ) + return harm, peri + + def _run_atom_head( + self, + features: Mapping[str, torch.Tensor], + batch: Any, + ) -> tuple[LJBag | None, ChargeBag | None]: + if self.atom_head is None or "atoms" not in features: + return None, None + atom_feats = features["atoms"] + atoms_ns = _get_namespace(batch, "atoms") + batch_idx = _ns_tensor(atoms_ns, "batch") if atoms_ns is not None else None + if batch_idx is None and isinstance(batch, Mapping): + batch_idx = batch.get("batch") + kwargs: dict[str, Any] = {} + if batch_idx is not None: + kwargs["batch"] = batch_idx + z = _ns_tensor(atoms_ns, "Z") if atoms_ns is not None else None + if z is not None: + kwargs["Z"] = z + out = self.atom_head(atom_feats, **kwargs) + lj = None + charges = None + if "epsilon" in out and "sigma" in out: + lj = LJBag(epsilon=out["epsilon"], sigma=out["sigma"]) + q = out.get("charge", out.get("q")) + if q is not None: + charges = ChargeBag(q=q) + return lj, charges + + # ------------------------------------------------------------------ + # energy + # ------------------------------------------------------------------ + + def energy( + self, + ir: PotentialIR, + batch: Any, + *, + pos: torch.Tensor | None = None, + ) -> torch.Tensor: + """Evaluate Class-I energy from an IR and geometry. + + Args: + ir: Parameter bags from :meth:`parameterize` (or constructed). + batch: Topology batch with valence namespaces. + pos: Optional positions ``(N, 3)``; else read from batch. + + Returns: + Scalar energy (kcal/mol). + + Notes: + Built-in path uses per-interaction continuous parameters with + identity type indices into the IR bag tables (no discrete type + table inside the composer). When ``evaluator`` was provided at + construction, that callable is used instead. + """ + if self.evaluator is not None: + return self.evaluator(ir, batch, pos=pos) + + p = _pos_from_batch(batch, pos) + device = p.device + dtype = p.dtype + total = torch.zeros((), device=device, dtype=dtype) + terms = self._term_energies(ir, batch, pos=p) + for e in terms.values(): + total = total + e + return total + + def term_energies( + self, + ir: PotentialIR, + batch: Any, + *, + pos: torch.Tensor | None = None, + ) -> dict[str, torch.Tensor]: + """Per-term energy breakdown (kcal/mol scalars).""" + p = _pos_from_batch(batch, pos) + return self._term_energies(ir, batch, pos=p) + + def _term_energies( + self, + ir: PotentialIR, + batch: Any, + *, + pos: torch.Tensor, + ) -> dict[str, torch.Tensor]: + terms: dict[str, torch.Tensor] = {} + device = pos.device + + if ir.bonds is not None: + bond_index = _bond_index_from_batch(batch) + if bond_index is not None and bond_index.size(1) > 0: + n = bond_index.size(1) + pot = BondHarmonic(k=ir.bonds.k, r0=ir.bonds.r0) + terms["bonds"] = pot( + pos=pos, + bond_index=bond_index, + bond_types=_identity_types(n, device), + ) + else: + terms["bonds"] = torch.zeros((), device=device, dtype=pos.dtype) + + if ir.angles is not None: + angle_index = _angle_index_from_batch(batch) + if angle_index is not None and angle_index.size(1) > 0: + n = angle_index.size(1) + pot = AngleHarmonic(k=ir.angles.k, theta0=ir.angles.theta0) + terms["angles"] = pot( + pos=pos, + angle_index=angle_index, + angle_types=_identity_types(n, device), + ) + else: + terms["angles"] = torch.zeros((), device=device, dtype=pos.dtype) + + if ir.propers is not None: + proper_index = _proper_index_from_batch(batch) + if proper_index is not None and proper_index.size(1) > 0: + n = proper_index.size(1) + pot = ProperTorsionPeriodic( + k=ir.propers.k, + periodicity=ir.propers.periodicity, + phase=ir.propers.phase, + idivf=ir.propers.idivf, + ) + terms["propers"] = pot( + pos=pos, + proper_index=proper_index, + proper_types=_identity_types(n, device), + ) + else: + terms["propers"] = torch.zeros((), device=device, dtype=pos.dtype) + + if ir.impropers_harmonic is not None: + improper_index = _improper_index_from_batch(batch) + if improper_index is not None and improper_index.size(1) > 0: + n = improper_index.size(1) + pot = ImproperHarmonic( + k=ir.impropers_harmonic.k, + chi0=ir.impropers_harmonic.chi0, + ) + terms["impropers_harmonic"] = pot( + pos=pos, + improper_index=improper_index, + improper_types=_identity_types(n, device), + ) + else: + terms["impropers_harmonic"] = torch.zeros((), device=device, dtype=pos.dtype) + + if ir.impropers_periodic is not None: + improper_index = _improper_index_from_batch(batch) + if improper_index is not None and improper_index.size(1) > 0: + n = improper_index.size(1) + pot = ImproperPeriodic( + k=ir.impropers_periodic.k, + periodicity=ir.impropers_periodic.periodicity, + phase=ir.impropers_periodic.phase, + idivf=ir.impropers_periodic.idivf, + ) + terms["impropers_periodic"] = pot( + pos=pos, + improper_index=improper_index, + improper_types=_identity_types(n, device), + ) + else: + terms["impropers_periodic"] = torch.zeros((), device=device, dtype=pos.dtype) + + return terms + + # ------------------------------------------------------------------ + # forward (thin composition) + # ------------------------------------------------------------------ + + def forward( + self, + batch: Any, + features: Mapping[str, torch.Tensor], + *, + pos: torch.Tensor | None = None, + return_terms: bool = False, + ) -> dict[str, Any]: + """Thin chain: :meth:`parameterize` then :meth:`energy`. + + Args: + batch: Topology + geometry batch. + features: Feature tensors keyed by interaction class. + pos: Optional positions override. + return_terms: If True, include ``term_energies`` breakdown. + + Returns: + Dict with ``energy`` (scalar), ``ir`` (:class:`PotentialIR`), and + optionally ``term_energies``. + """ + ir = self.parameterize(features, batch) + p = _pos_from_batch(batch, pos) + energy = self.energy(ir, batch, pos=p) + out: dict[str, Any] = {"energy": energy, "ir": ir} + if return_terms: + out["term_energies"] = self.term_energies(ir, batch, pos=p) + return out diff --git a/src/molpot/composition/energy_force.py b/src/molpot/composition/energy_force.py deleted file mode 100644 index 3625cb9..0000000 --- a/src/molpot/composition/energy_force.py +++ /dev/null @@ -1,99 +0,0 @@ -"""Generic energy + force model wrapper (physics home for force derivation). - -Encoders live in ``molzoo``; force derivation lives in ``molpot``. Subclass -:class:`EnergyForceModel` and implement :meth:`energy_forward` to attach a -:class:`~molpot.derivation.ForceDerivation` without inventing a third force -path. ``molzoo.pinet.PiNetPotential`` is the primary consumer. -""" - -from __future__ import annotations - -from typing import Literal - -import torch -import torch.nn as nn -from tensordict import TensorDict - -from molpot.derivation import ForceDerivation - - -class EnergyForceModel(nn.Module): - """Energy model + optional forces via the shared :class:`ForceDerivation`. - - Subclasses implement :meth:`energy_forward` only. Force paths: - - * **eval + functorch**: single ``grad(..., has_aux=True)`` pass. - * **train + functorch**: eager energy (params connected) + force pass. - * **autograd**: always ``ForceDerivation(method="autograd")`` (cuEq-safe). - - Args: - force_method: ``"functorch"`` or ``"autograd"`` (see - :class:`molpot.derivation.ForceDerivation`). - compute_forces: Default for :meth:`forward` when ``compute_forces`` is - not passed explicitly. - """ - - def __init__( - self, - *, - force_method: Literal["functorch", "autograd"] = "functorch", - compute_forces: bool = False, - ) -> None: - super().__init__() - self.force_derivation = ForceDerivation(method=force_method) - self.compute_forces_default = compute_forces - self._compiled_energy_forward = None - - def energy_forward(self, batch: TensorDict) -> dict[str, torch.Tensor]: - """Return at least ``{"energy": (B,)}``; may include auxiliaries.""" - raise NotImplementedError - - def forward( - self, batch: TensorDict, *, compute_forces: bool | None = None - ) -> dict[str, torch.Tensor]: - if compute_forces is None: - compute_forces = self.compute_forces_default - if compute_forces: - return self._forward_with_forces(batch) - energy_forward = self._compiled_energy_forward or self.energy_forward - return energy_forward(batch) - - def _forward_with_forces(self, batch: TensorDict) -> dict[str, torch.Tensor]: - pos = batch["atoms", "pos"].detach() - base = batch.clone() - base["atoms", "pos"] = pos - - if self.force_derivation.method == "functorch" and not self.training: - - def energy_fn_aux(p: torch.Tensor) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: - b = base.clone() - b["atoms", "pos"] = p - out = self.energy_forward(b) - return out["energy"].sum(), out - - forces, out = self.force_derivation(energy_fn_aux, pos, has_aux=True) - out["forces"] = forces - return out - - def energy_fn(p: torch.Tensor) -> torch.Tensor: - b = base.clone() - b["atoms", "pos"] = p - return self.energy_forward(b)["energy"].sum() - - # Training (or autograd): keep an eager energy pass so parameters stay - # connected for a force-supervised loss when using functorch's separate - # force transform; autograd_forces already handles create_graph. - if self.force_derivation.method == "functorch": - out = self.energy_forward(base.clone()) - out["forces"] = self.force_derivation(energy_fn, pos) - return out - - out = self.energy_forward(base.clone()) - out["forces"] = self.force_derivation(energy_fn, pos) - return out - - def compile_energy(self, *, backend: str = "inductor", **kwargs) -> None: - """Compile the energy-only forward (not the force training path).""" - self._compiled_energy_forward = torch.compile( - self.energy_forward, backend=backend, **kwargs - ) diff --git a/src/molpot/composition/heads.py b/src/molpot/composition/heads.py index 7954015..e9c6cd0 100644 --- a/src/molpot/composition/heads.py +++ b/src/molpot/composition/heads.py @@ -6,6 +6,8 @@ import torch.nn as nn import torch.nn.functional as F +from molix import config + class LJParameterHead(nn.Module): """Predict per-atom Lennard-Jones parameters from node features. @@ -31,16 +33,18 @@ def __init__( self.min_epsilon = min_epsilon self.min_sigma = min_sigma self.mlp = nn.Sequential( - nn.Linear(feature_dim, hidden_dim), + nn.Linear(feature_dim, hidden_dim, dtype=config.ftype), nn.SiLU(), - nn.Linear(hidden_dim, 2), + nn.Linear(hidden_dim, 2, dtype=config.ftype), ) - def forward(self, node_features: torch.Tensor) -> dict[str, torch.Tensor]: + def forward(self, node_features: torch.Tensor, **kwargs) -> dict[str, torch.Tensor]: """Predict LJ parameters from node features. Args: node_features: Per-node features ``(N, D)``. + **kwargs: Ignored; accepted for a uniform parameter-head signature + (e.g. MultiHead forwarding ``batch`` / ``Z``). Returns: Dict with ``"epsilon"`` ``(N,)`` and ``"sigma"`` ``(N,)``. @@ -72,9 +76,9 @@ def __init__( self.min_eps = min_eps self.min_lam = min_lam self.mlp = nn.Sequential( - nn.Linear(feature_dim, hidden_dim), + nn.Linear(feature_dim, hidden_dim, dtype=config.ftype), nn.SiLU(), - nn.Linear(hidden_dim, 2), + nn.Linear(hidden_dim, 2, dtype=config.ftype), ) def forward(self, node_features: torch.Tensor, **kwargs) -> dict[str, torch.Tensor]: @@ -115,9 +119,9 @@ def __init__( self.min_eps = min_eps self.min_lam = min_lam self.mlp = nn.Sequential( - nn.Linear(feature_dim, hidden_dim), + nn.Linear(feature_dim, hidden_dim, dtype=config.ftype), nn.SiLU(), - nn.Linear(hidden_dim, 2), + nn.Linear(hidden_dim, 2, dtype=config.ftype), ) def forward(self, node_features: torch.Tensor, **kwargs) -> dict[str, torch.Tensor]: @@ -157,9 +161,9 @@ def __init__( super().__init__() self.total_charge = total_charge self.mlp = nn.Sequential( - nn.Linear(feature_dim, hidden_dim), + nn.Linear(feature_dim, hidden_dim, dtype=config.ftype), nn.SiLU(), - nn.Linear(hidden_dim, 1), + nn.Linear(hidden_dim, 1, dtype=config.ftype), ) def forward( @@ -220,9 +224,9 @@ def __init__( self.register_buffer("alpha_free", alpha_free) self.register_buffer("r_star_free", r_star_free) self.mlp = nn.Sequential( - nn.Linear(feature_dim, hidden_dim), + nn.Linear(feature_dim, hidden_dim, dtype=config.ftype), nn.SiLU(), - nn.Linear(hidden_dim, 1), + nn.Linear(hidden_dim, 1, dtype=config.ftype), ) def forward( diff --git a/src/molpot/composition/mm_heads.py b/src/molpot/composition/mm_heads.py new file mode 100644 index 0000000..12e9560 --- /dev/null +++ b/src/molpot/composition/mm_heads.py @@ -0,0 +1,324 @@ +"""Continuous Class-I MM parameter heads (learnable classical force fields). + +Maps per-interaction (or per-atom) feature vectors to Class-I IR parameter +dicts in CLASS_I_CANONICAL units (kcal/mol, Å, e, rad). Heads are pure MLPs; +endpoint symmetry for bonds/angles/propers is the caller's responsibility via +order-invariant feature pooling (see module notes below). + +Endpoint symmetry (documented contract) +-------------------------------------- +Bond features that reverse atom order ``(i,j)↔(j,i)``, angle features that +reverse ``(i,j,k)↔(k,j,i)``, and proper torsion features that reverse +``(i,j,k,l)↔(l,k,j,i)`` must yield **identical** parameters. Enforce this by +symmetric feature construction at the call site (e.g. sum/mean of endpoint +embeddings). These heads do **not** reorder atoms — they stay pure MLPs. + +Units (CLASS_I_CANONICAL) +------------------------- +- Bond ``k``: kcal/mol/Ų; ``r0``: Å +- Angle ``k``: kcal/mol/rad²; ``theta0``: rad ∈ (0, π) +- Proper / improper periodic ``k``: kcal/mol; ``phase``: rad +- Improper harmonic ``k``: kcal/mol/rad²; ``chi0``: rad + +References: + Spec: learnable-classical-ff-03-mm-heads + OpenMM User Guide §19 "Forces" +""" + +from __future__ import annotations + +import math + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from molix import config + +__all__ = [ + "BondParamHead", + "AngleParamHead", + "ProperTorsionParamHead", + "ImproperParamHead", +] + + +def _mlp(feature_dim: int, hidden_dim: int, out_dim: int) -> nn.Sequential: + return nn.Sequential( + nn.Linear(feature_dim, hidden_dim, dtype=config.ftype), + nn.SiLU(), + nn.Linear(hidden_dim, out_dim, dtype=config.ftype), + ) + + +class BondParamHead(nn.Module): + """Map bond features to harmonic bond parameters. + + Args: + feature_dim: Input feature dimension. + hidden_dim: Hidden layer dimension. + min_k: Positive floor for force constants (kcal/mol/Ų). + min_r0: Positive floor for equilibrium lengths (Å). + + Forward: + features: ``(N_bonds, D)`` → ``{"k": (N,), "r0": (N,)}`` with + ``k > min_k`` and ``r0 > min_r0`` via softplus. + """ + + def __init__( + self, + feature_dim: int, + hidden_dim: int = 64, + min_k: float = 1e-4, + min_r0: float = 1e-4, + ) -> None: + super().__init__() + self.min_k = min_k + self.min_r0 = min_r0 + self.mlp = _mlp(feature_dim, hidden_dim, 2) + + def forward(self, features: torch.Tensor) -> dict[str, torch.Tensor]: + """Predict bond ``k``, ``r0`` from features. + + Args: + features: Bond features ``(N_bonds, feature_dim)``. Prefer + order-invariant pooling so ``(i,j)`` and ``(j,i)`` match. + + Returns: + Dict with ``k`` ``(N,)`` (kcal/mol/Ų) and ``r0`` ``(N,)`` (Å). + """ + raw = self.mlp(features) + k = F.softplus(raw[:, 0]) + self.min_k + r0 = F.softplus(raw[:, 1]) + self.min_r0 + return {"k": k, "r0": r0} + + +class AngleParamHead(nn.Module): + """Map angle features to harmonic angle parameters. + + Args: + feature_dim: Input feature dimension. + hidden_dim: Hidden layer dimension. + min_k: Positive floor for force constants (kcal/mol/rad²). + + Forward: + features: ``(N_angles, D)`` → ``{"k": (N,), "theta0": (N,)}`` with + ``k > min_k`` and ``theta0 ∈ (0, π)`` via scaled sigmoid. + """ + + def __init__( + self, + feature_dim: int, + hidden_dim: int = 64, + min_k: float = 1e-4, + ) -> None: + super().__init__() + self.min_k = min_k + self.mlp = _mlp(feature_dim, hidden_dim, 2) + + def forward(self, features: torch.Tensor) -> dict[str, torch.Tensor]: + """Predict angle ``k``, ``theta0`` from features. + + Args: + features: Angle features ``(N_angles, feature_dim)``. Prefer + endpoint-symmetric pooling for ``(i,j,k)`` vs ``(k,j,i)``. + + Returns: + Dict with ``k`` ``(N,)`` (kcal/mol/rad²) and ``theta0`` ``(N,)`` + in radians, strictly inside ``(0, π)``. + """ + raw = self.mlp(features) + k = F.softplus(raw[:, 0]) + self.min_k + # Open interval (0, π): epsilon margin avoids exact 0/π singularities. + eps = 1e-4 + theta0 = eps + (math.pi - 2.0 * eps) * torch.sigmoid(raw[:, 1]) + return {"k": k, "theta0": theta0} + + +class ProperTorsionParamHead(nn.Module): + """Map proper-torsion features to multi-term cosine parameters. + + Periodicity is a fixed buffer (config-time). Force constants use softplus + (``k ≥ 0``); phases are unconstrained (radians). + + Args: + feature_dim: Input feature dimension. + hidden_dim: Hidden layer dimension. + n_terms: Number of Fourier terms ``T``. + periodicity: Integer periodicities of length ``T`` (shared across + interactions). Defaults to ``(1, 2, …, T)``. + min_k: Softplus floor for barrier heights (kcal/mol). + default_idivf: Default AMBER-style identity divisor per interaction. + """ + + def __init__( + self, + feature_dim: int, + hidden_dim: int = 64, + n_terms: int = 1, + periodicity: tuple[int, ...] | list[int] | None = None, + min_k: float = 0.0, + default_idivf: float = 1.0, + ) -> None: + super().__init__() + if n_terms < 1: + raise ValueError(f"n_terms must be >= 1, got {n_terms}") + if periodicity is None: + periodicity = tuple(range(1, n_terms + 1)) + if len(periodicity) != n_terms: + raise ValueError(f"periodicity length {len(periodicity)} must equal n_terms={n_terms}") + self.n_terms = n_terms + self.min_k = min_k + self.default_idivf = default_idivf + # k and phase per term + self.mlp = _mlp(feature_dim, hidden_dim, 2 * n_terms) + self.register_buffer( + "periodicity", + torch.tensor(list(periodicity), dtype=torch.long), + ) + + def forward(self, features: torch.Tensor) -> dict[str, torch.Tensor]: + """Predict multi-term proper torsion parameters. + + Args: + features: Proper features ``(N_propers, feature_dim)``. + + Returns: + Dict with: + + - ``k``: ``(N, T)`` barriers (kcal/mol), ``≥ min_k`` + - ``phase``: ``(N, T)`` phases (rad), unconstrained + - ``periodicity``: ``(T,)`` integer buffer + - ``idivf``: ``(N,)`` identity divisors (positive) + """ + n = features.shape[0] + raw = self.mlp(features) + t = self.n_terms + k = F.softplus(raw[:, :t]) + self.min_k + phase = raw[:, t : 2 * t] + idivf = torch.full( + (n,), + self.default_idivf, + dtype=features.dtype, + device=features.device, + ) + return { + "k": k, + "phase": phase, + "periodicity": self.periodicity, + "idivf": idivf, + } + + +class ImproperParamHead(nn.Module): + """Map improper features to harmonic and/or periodic improper parameters. + + Config flags select which parameter families to emit (not a multi-method + switch over unrelated kernels — both families are Class-I improper terms). + + When both modes are enabled, keys are namespaced + (``k_harmonic`` / ``chi0`` and ``k_periodic`` / ``phase`` / …) to avoid + collisions. Single-mode outputs use the short IR field names (``k``, + ``chi0`` or ``k``, ``phase``, ``periodicity``, ``idivf``). + + Args: + feature_dim: Input feature dimension. + hidden_dim: Hidden layer dimension. + include_harmonic: Emit harmonic improper params. + include_periodic: Emit multi-term cosine improper params. + n_terms: Fourier terms for the periodic branch. + periodicity: Integer periodicities for the periodic branch. + min_k: Softplus floor for force constants / barriers. + default_idivf: Identity divisor for the periodic branch. + """ + + def __init__( + self, + feature_dim: int, + hidden_dim: int = 64, + *, + include_harmonic: bool = True, + include_periodic: bool = False, + n_terms: int = 1, + periodicity: tuple[int, ...] | list[int] | None = None, + min_k: float = 1e-4, + default_idivf: float = 1.0, + ) -> None: + super().__init__() + if not include_harmonic and not include_periodic: + raise ValueError("ImproperParamHead requires include_harmonic and/or include_periodic") + self.include_harmonic = include_harmonic + self.include_periodic = include_periodic + self.min_k = min_k + self.default_idivf = default_idivf + self.n_terms = n_terms + + out_dim = 0 + if include_harmonic: + out_dim += 2 # k, chi0 + if include_periodic: + if n_terms < 1: + raise ValueError(f"n_terms must be >= 1, got {n_terms}") + if periodicity is None: + periodicity = tuple(2 for _ in range(n_terms)) # common improper n=2 + if len(periodicity) != n_terms: + raise ValueError( + f"periodicity length {len(periodicity)} must equal n_terms={n_terms}" + ) + out_dim += 2 * n_terms # k, phase per term + self.register_buffer( + "periodicity", + torch.tensor(list(periodicity), dtype=torch.long), + ) + else: + # Placeholder so attribute always exists for type checkers. + self.register_buffer("periodicity", torch.zeros(0, dtype=torch.long)) + + self.mlp = _mlp(feature_dim, hidden_dim, out_dim) + + def forward(self, features: torch.Tensor) -> dict[str, torch.Tensor]: + """Predict improper parameters from features. + + Args: + features: Improper features ``(N_impropers, feature_dim)``. + + Returns: + Dict of parameter tensors. See class docstring for key naming when + both harmonic and periodic modes are enabled. + """ + n = features.shape[0] + raw = self.mlp(features) + out: dict[str, torch.Tensor] = {} + offset = 0 + both = self.include_harmonic and self.include_periodic + + if self.include_harmonic: + k_h = F.softplus(raw[:, offset]) + self.min_k + chi0 = raw[:, offset + 1] + offset += 2 + if both: + out["k_harmonic"] = k_h + out["chi0"] = chi0 + else: + out["k"] = k_h + out["chi0"] = chi0 + + if self.include_periodic: + t = self.n_terms + k_p = F.softplus(raw[:, offset : offset + t]) + self.min_k + phase = raw[:, offset + t : offset + 2 * t] + idivf = torch.full( + (n,), + self.default_idivf, + dtype=features.dtype, + device=features.device, + ) + if both: + out["k_periodic"] = k_p + else: + out["k"] = k_p + out["phase"] = phase + out["periodicity"] = self.periodicity + out["idivf"] = idivf + + return out diff --git a/src/molpot/composition/parameterizer.py b/src/molpot/composition/parameterizer.py new file mode 100644 index 0000000..1f4c07a --- /dev/null +++ b/src/molpot/composition/parameterizer.py @@ -0,0 +1,367 @@ +"""ClassicalMMParameterizer: encoder → MM heads → PotentialIR → Class-I E/F. + +End-to-end training-time module for learnable classical force fields. +Accepts a chemical-perception encoder **via Protocol** so ``molpot`` never +imports ``molzoo`` / ``molrep.chem``. Heads and energy evaluation are owned by +:class:`~molpot.composition.classical_mm.ClassicalMMComposer` (spec 03). + +Pipeline primitives (prefer explicit methods over one opaque façade):: + + features = param.encode(batch) + ir = param.parameterize(batch) # or features= + energy = param.energy(batch) # or ir= + out = param(batch, compute_forces=True) # thin composition + +Units: CLASS_I_CANONICAL (kcal/mol, Å, e, rad) inside IR and energy. +Optional eV conversion only via :func:`energy_kcal_to_ev` at the loss edge — +never by mutating IR units. + +Forces: solely :class:`~molpot.derivation.ForceDerivation` +(``F = -∂E/∂r``). No hand-rolled force kernels. + +References: + Spec: learnable-classical-ff-05-neural-parameterizer + OpenMM User Guide §19 "Forces" +""" + +from __future__ import annotations + +from collections.abc import Mapping +from typing import Any, Protocol, runtime_checkable + +import torch +import torch.nn as nn +from tensordict import TensorDict + +from molpot.composition.classical_mm import ClassicalMMComposer +from molpot.derivation import ForceDerivation +from molpot.derivation.protocol import write_energy, write_forces +from molpot.ir import PotentialIR + +__all__ = [ + "ChemEmbeddingsLike", + "ChemEncoderProtocol", + "ClassicalMMParameterizer", + "KCAL_MOL_TO_EV", + "energy_kcal_to_ev", +] + +# kcal/mol → eV. 1 eV = 23.060547830619026 kcal/mol (CODATA / OpenMM-style). +# Used only at the training-loss boundary — never redefines IR units. +KCAL_MOL_TO_EV: float = 1.0 / 23.060547830619026 + + +def energy_kcal_to_ev(energy_kcal: torch.Tensor) -> torch.Tensor: + """Convert Class-I energy from kcal/mol to eV **without mutating IR**. + + Args: + energy_kcal: Energy tensor in kcal/mol (any shape). + + Returns: + Same-shaped tensor in eV (``energy_kcal * KCAL_MOL_TO_EV``). + + Notes: + Call this only at the loss / logging boundary. Parameterize and + classical energy stay in CLASS_I_CANONICAL (kcal/mol). + """ + return energy_kcal * KCAL_MOL_TO_EV + + +@runtime_checkable +class ChemEmbeddingsLike(Protocol): + """Structural type for continuous chem feature payloads. + + Duck-typed so molpot never imports ``molrep.chem.ChemEmbeddings``. + Implementations expose :meth:`interaction_dict` with keys consumable by + :class:`~molpot.composition.classical_mm.ClassicalMMComposer` + (``atoms`` / ``bonds`` / ``angles`` / ``propers`` / ``impropers``). + """ + + def interaction_dict(self) -> Mapping[str, torch.Tensor]: + """Feature mapping keyed for ClassicalMMComposer heads.""" + ... + + +@runtime_checkable +class ChemEncoderProtocol(Protocol): + """Minimal encoder surface for classical MM parameterization. + + Implementations live in molrep/molzoo; molpot only sees this Protocol. + Either :meth:`forward` writing batch chem features, or + :meth:`embeddings` returning a :class:`ChemEmbeddingsLike`, or both. + """ + + def forward(self, td: TensorDict) -> TensorDict: + """Run chemical perception; may write ``*.chem_features`` on ``td``.""" + ... + + +# Optional embeddings method is not required by the Protocol class itself +# (runtime_checkable only checks methods declared on the Protocol). Callers +# and ClassicalMMParameterizer.encode discover embeddings via getattr. + + +_FEATURE_NS_KEYS: tuple[tuple[str, str], ...] = ( + ("atoms", "atoms"), + ("bonds", "bonds"), + ("angles", "angles"), + ("propers", "propers"), + ("impropers", "impropers"), +) + + +def _features_from_batch_chem(batch: Any) -> dict[str, torch.Tensor] | None: + """Extract ``*.chem_features`` written by a chem encoder, if present.""" + if batch is None: + return None + out: dict[str, torch.Tensor] = {} + for ns_name, feat_key in _FEATURE_NS_KEYS: + try: + if ns_name not in batch: + continue + ns = batch[ns_name] + if isinstance(ns, (Mapping, TensorDict)) and "chem_features" in ns: + out[feat_key] = ns["chem_features"] + except Exception: # noqa: BLE001 — TensorDict key miss variants + continue + return out or None + + +def _features_from_embeddings(emb: Any) -> dict[str, torch.Tensor]: + """Normalize ChemEmbeddingsLike / Mapping into composer feature keys.""" + if isinstance(emb, Mapping): + # Accept plural (composer) or singular (ChemEmbeddings.as_dict) keys. + alias = { + "atom": "atoms", + "bond": "bonds", + "angle": "angles", + "proper": "propers", + "improper": "impropers", + } + out: dict[str, torch.Tensor] = {} + for k, v in emb.items(): + out[alias.get(k, k)] = v + return out + interaction = getattr(emb, "interaction_dict", None) + if callable(interaction): + return dict(interaction()) + as_dict = getattr(emb, "as_dict", None) + if callable(as_dict): + return _features_from_embeddings(as_dict()) + raise TypeError("encoder embeddings must be a Mapping or expose interaction_dict()/as_dict()") + + +def _pos_from_batch(batch: Any, pos: torch.Tensor | None) -> torch.Tensor: + if pos is not None: + return pos + if isinstance(batch, (Mapping, TensorDict)): + try: + if "atoms" in batch: + atoms = batch["atoms"] + if isinstance(atoms, (Mapping, TensorDict)) and "pos" in atoms: + return atoms["pos"] + except Exception: # noqa: BLE001 + pass + if "pos" in batch: + return batch["pos"] + raise ValueError("ClassicalMMParameterizer requires pos or batch['atoms','pos']") + + +class ClassicalMMParameterizer(nn.Module): + """encoder (Protocol) → MM heads → PotentialIR → classical E (+ optional F). + + Args: + encoder: Chem-perception module satisfying + :class:`ChemEncoderProtocol` (registered if ``nn.Module``). + composer: :class:`~molpot.composition.classical_mm.ClassicalMMComposer` + owning heads and Class-I energy evaluation. + force_derivation: Optional + :class:`~molpot.derivation.ForceDerivation`. Defaults to + autograd backend (universal). Forces are never hand-rolled. + + Notes: + Prefer :meth:`encode`, :meth:`parameterize`, and :meth:`energy` as + separate primitives. :meth:`forward` is a thin chain for training + loops (``compute_forces`` optional). + """ + + def __init__( + self, + encoder: nn.Module, + composer: ClassicalMMComposer, + force_derivation: ForceDerivation | None = None, + ) -> None: + super().__init__() + if not isinstance(encoder, nn.Module): + raise TypeError("encoder must be an nn.Module (Protocol surface)") + if not isinstance(composer, ClassicalMMComposer): + raise TypeError("composer must be a ClassicalMMComposer") + self.encoder = encoder + self.composer = composer + self.force_derivation = ( + force_derivation if force_derivation is not None else ForceDerivation(method="autograd") + ) + + # ------------------------------------------------------------------ + # encode + # ------------------------------------------------------------------ + + def encode(self, batch: TensorDict | Mapping[str, Any]) -> dict[str, torch.Tensor]: + """Run the encoder Protocol and return composer-keyed features. + + Resolution order: + + 1. If the encoder exposes ``embeddings(batch)``, use that payload. + 2. Else run ``encoder.forward(batch)`` and read ``*.chem_features``. + 3. Else, if ``forward`` returned a Mapping of feature tensors, use it. + + Args: + batch: Nested TensorDict (or Mapping) with topology + geometry. + + Returns: + Feature tensors keyed ``atoms`` / ``bonds`` / ``angles`` / + ``propers`` / ``impropers`` (subset present). + + Raises: + RuntimeError: If no features can be resolved from the encoder. + """ + embeddings_fn = getattr(self.encoder, "embeddings", None) + if callable(embeddings_fn): + # Prefer embeddings() after a forward pass so side-effect writers + # (e.g. ChemEncoder writing chem_features) stay consistent. + forward = getattr(self.encoder, "forward", None) + if callable(forward): + try: + forward(batch) # type: ignore[arg-type] + except TypeError: + # Some fakes only implement embeddings; ignore. + pass + emb = embeddings_fn(batch) + # ChemEncoder.embeddings reads written chem_features; if the + # encoder was not run (or is FakeEncoder with direct embeddings), + # emb is still valid. + return _features_from_embeddings(emb) + + forward = getattr(self.encoder, "forward", None) + if not callable(forward): + raise RuntimeError("encoder must implement forward(batch) and/or embeddings(batch)") + result = forward(batch) # type: ignore[misc] + + # Features written onto the batch under *.chem_features + from_batch = _features_from_batch_chem(batch) + if from_batch is not None: + return from_batch + if result is not None and result is not batch: + from_result = _features_from_batch_chem(result) + if from_result is not None: + return from_result + if isinstance(result, Mapping) and any( + k in result + for k in ("atoms", "bonds", "angles", "propers", "impropers", "atom", "bond") + ): + return _features_from_embeddings(result) + + raise RuntimeError( + "encoder produced no features: expected embeddings() or *.chem_features on batch" + ) + + # ------------------------------------------------------------------ + # parameterize + # ------------------------------------------------------------------ + + def parameterize( + self, + batch: TensorDict | Mapping[str, Any], + features: Mapping[str, torch.Tensor] | None = None, + ) -> PotentialIR: + """Heads → Class-I :class:`~molpot.ir.PotentialIR` (kcal/mol…). + + Args: + batch: Topology batch (valence namespaces). + features: Optional precomputed feature dict. When ``None``, + :meth:`encode` runs the encoder Protocol. + + Returns: + :class:`PotentialIR` with ``unit_system="class_i_canonical"``. + """ + feats = features if features is not None else self.encode(batch) + return self.composer.parameterize(feats, batch) + + # ------------------------------------------------------------------ + # energy + # ------------------------------------------------------------------ + + def energy( + self, + batch: TensorDict | Mapping[str, Any], + ir: PotentialIR | None = None, + *, + pos: torch.Tensor | None = None, + ) -> torch.Tensor: + """Classical Class-I energy sum (kcal/mol). + + Args: + batch: Topology + geometry batch. + ir: Optional precomputed IR; when ``None``, :meth:`parameterize`. + pos: Optional positions ``(N, 3)``; else read from batch. + + Returns: + Scalar energy in kcal/mol. + """ + bag = ir if ir is not None else self.parameterize(batch) + return self.composer.energy(bag, batch, pos=pos) + + # ------------------------------------------------------------------ + # forward (thin composition) + # ------------------------------------------------------------------ + + def forward( + self, + batch: TensorDict, + *, + compute_forces: bool = False, + features: Mapping[str, torch.Tensor] | None = None, + pos: torch.Tensor | None = None, + return_terms: bool = False, + ) -> dict[str, Any]: + """Thin chain: encode → parameterize → energy (+ optional forces). + + Args: + batch: Nested TensorDict with topology and ``atoms.pos``. + compute_forces: If True, derive ``F = -∂E/∂pos`` via + :class:`~molpot.derivation.ForceDerivation` and write + ``atoms.forces``. + features: Optional feature override (skips encoder). + pos: Optional positions override. + return_terms: If True, include per-term energy breakdown. + + Returns: + Dict with ``energy`` (scalar kcal/mol), ``ir`` + (:class:`PotentialIR`), and optionally ``forces`` ``(N, 3)`` and + ``term_energies``. Also writes ``graphs.energy`` (and + ``atoms.forces`` when requested) onto ``batch`` in place. + """ + ir = self.parameterize(batch, features=features) + p = _pos_from_batch(batch, pos) + energy = self.composer.energy(ir, batch, pos=p) + + # Peer keys on the batch (molix schema: graphs.energy). + # Scalar energy → ensure_graphs with batch_size=[] when dim==0. + if isinstance(batch, TensorDict): + write_energy(batch, energy) + + out: dict[str, Any] = {"energy": energy, "ir": ir} + if return_terms: + out["term_energies"] = self.composer.term_energies(ir, batch, pos=p) + + if compute_forces: + + def energy_fn(positions: torch.Tensor) -> torch.Tensor: + return self.composer.energy(ir, batch, pos=positions) + + forces = self.force_derivation(energy_fn, p) + out["forces"] = forces + if isinstance(batch, TensorDict): + write_forces(batch, forces) + + return out diff --git a/src/molpot/composition/sonata.py b/src/molpot/composition/sonata.py index 020101b..17fa965 100644 --- a/src/molpot/composition/sonata.py +++ b/src/molpot/composition/sonata.py @@ -542,93 +542,94 @@ def energy_of_strain(strain: torch.Tensor) -> torch.Tensor: return out + @classmethod + def from_encoder( + cls, + encoder: nn.Module, + *, + sigma: float = 1.0, + dl: float = 2.0, + prefactor: float = 90.4756, + charge: bool = True, + dipole: bool = True, + quadrupole: bool = True, + constrain_total_charge: bool = True, + short_range_head: nn.Module | list[nn.Module] | None = None, + total_charge_key: str = "total_charge", + hidden_dim: int = 128, + avg_num_neighbors: float | None = None, + ) -> "Sonata": + """Build a wired :class:`Sonata` from an encoder and loose hyperparameters. -# --------------------------------------------------------------------------- -# build_sonata factory -# --------------------------------------------------------------------------- - + Distinct from :meth:`__init__`, which takes the sub-modules + already built: this constructs the :class:`PermMultipoleHead` and + :class:`EwaldMultipoleEnergy` for you and validates that the + encoder can actually feed them. -def build_sonata( - encoder: nn.Module, - *, - sigma: float = 1.0, - dl: float = 2.0, - prefactor: float = 90.4756, - charge: bool = True, - dipole: bool = True, - quadrupole: bool = True, - constrain_total_charge: bool = True, - short_range_head: nn.Module | list[nn.Module] | None = None, - total_charge_key: str = "total_charge", - hidden_dim: int = 128, - avg_num_neighbors: float | None = None, -) -> Sonata: - """Build a wired :class:`Sonata` from an encoder and loose hyperparameters. + Args: + encoder: Allegro-style encoder. Must satisfy + ``encoder.expose_tensor_track is True`` and (when ``dipole`` + or ``quadrupole`` is on) ``encoder.l_max >= 2``. + sigma: σ-Gaussian charge-smearing length in Å. Default ``1.0``. + dl: Reciprocal-space grid resolution in Å. Default ``2.0``. + prefactor: Electrostatic prefactor ``1/(2 ε₀)``. Default + ``90.4756`` (eV·Å·e⁻²). + charge: Predict atomic charges. Default ``True``. + dipole: Predict atomic dipoles. Default ``True``. Requires + ``encoder.l_max >= 2``. + quadrupole: Predict atomic quadrupoles. Default ``True``. + Requires ``encoder.l_max >= 2``. + constrain_total_charge: Project per-graph charge sums onto + ``total_charge_key``. Default ``True``. + short_range_head: Optional ``nn.Module`` (or ``list``) writing + ``"energy_short"`` ``(B,)``. + total_charge_key: Per-graph total-charge key under ``graphs``. + Required when ``constrain_total_charge=True``. + hidden_dim: Hidden width of the multipole head's scalar MLPs. + avg_num_neighbors: Dataset-wide ⟨|N(i)|⟩ for the edge→atom + pool normalisation in the multipole head. - Args: - encoder: Allegro-style encoder. Must satisfy - ``encoder.expose_tensor_track is True`` and (when ``dipole`` - or ``quadrupole`` is on) ``encoder.l_max >= 2``. - sigma: σ-Gaussian charge-smearing length in Å. Default ``1.0``. - dl: Reciprocal-space grid resolution in Å. Default ``2.0``. - prefactor: Electrostatic prefactor ``1/(2 ε₀)``. Default - ``90.4756`` (eV·Å·e⁻²). - charge: Predict atomic charges. Default ``True``. - dipole: Predict atomic dipoles. Default ``True``. Requires - ``encoder.l_max >= 2``. - quadrupole: Predict atomic quadrupoles. Default ``True``. - Requires ``encoder.l_max >= 2``. - constrain_total_charge: Project per-graph charge sums onto - ``total_charge_key``. Default ``True``. - short_range_head: Optional ``nn.Module`` (or ``list``) writing - ``"energy_short"`` ``(B,)``. - total_charge_key: Per-graph total-charge key under ``graphs``. - Required when ``constrain_total_charge=True``. - hidden_dim: Hidden width of the multipole head's scalar MLPs. - avg_num_neighbors: Dataset-wide ⟨|N(i)|⟩ for the edge→atom - pool normalisation in the multipole head. + Returns: + A fully wired :class:`Sonata`. - Returns: - A fully wired :class:`Sonata`. + Raises: + ValueError: If ``encoder.expose_tensor_track`` is not ``True``, + or if ``dipole`` / ``quadrupole`` is requested but + ``encoder.l_max < 2``. + """ + if not getattr(encoder, "expose_tensor_track", False): + raise ValueError( + "Sonata requires an encoder built with `expose_tensor_track=True` " + "to expose the equivariant tensor-track features the multipole " + "head consumes. Reconstruct your encoder with this flag set." + ) + if (dipole or quadrupole) and getattr(encoder, "l_max", 0) < 2: + raise ValueError( + "Sonata requires `encoder.l_max >= 2` when `dipole=True` or " + f"`quadrupole=True`; got `l_max={getattr(encoder, 'l_max', None)}`." + ) - Raises: - ValueError: If ``encoder.expose_tensor_track`` is not ``True``, - or if ``dipole`` / ``quadrupole`` is requested but - ``encoder.l_max < 2``. - """ - if not getattr(encoder, "expose_tensor_track", False): - raise ValueError( - "Sonata requires an encoder built with `expose_tensor_track=True` " - "to expose the equivariant tensor-track features the multipole " - "head consumes. Reconstruct your encoder with this flag set." + head = PermMultipoleHead( + input_dim=encoder.output_dim, + avg_num_neighbors=avg_num_neighbors, + charge=charge, + dipole=dipole, + quadrupole=quadrupole, + constrain_total_charge=constrain_total_charge, + total_charge_key=total_charge_key, + hidden_dim=hidden_dim, + tensor_irreps=encoder.tensor_track_irreps, ) - if (dipole or quadrupole) and getattr(encoder, "l_max", 0) < 2: - raise ValueError( - "Sonata requires `encoder.l_max >= 2` when `dipole=True` or " - f"`quadrupole=True`; got `l_max={getattr(encoder, 'l_max', None)}`." + ewald = EwaldMultipoleEnergy( + sigma=sigma, + dl=dl, + prefactor=prefactor, + remove_self_interaction=True, + use_epsilon_r_scaling=False, + ) + return cls( + encoder=encoder, + perm_multipole_head=head, + ewald=ewald, + short_range_head=short_range_head, ) - - head = PermMultipoleHead( - input_dim=encoder.output_dim, - avg_num_neighbors=avg_num_neighbors, - charge=charge, - dipole=dipole, - quadrupole=quadrupole, - constrain_total_charge=constrain_total_charge, - total_charge_key=total_charge_key, - hidden_dim=hidden_dim, - tensor_irreps=encoder.tensor_track_irreps, - ) - ewald = EwaldMultipoleEnergy( - sigma=sigma, - dl=dl, - prefactor=prefactor, - remove_self_interaction=True, - use_epsilon_r_scaling=False, - ) - return Sonata( - encoder=encoder, - perm_multipole_head=head, - ewald=ewald, - short_range_head=short_range_head, - ) diff --git a/src/molpot/derivation/__init__.py b/src/molpot/derivation/__init__.py index 849f5c3..9c0e13f 100644 --- a/src/molpot/derivation/__init__.py +++ b/src/molpot/derivation/__init__.py @@ -1,24 +1,51 @@ -"""Physical quantity derivation modules. +"""Physical quantity derivation: aggregation + in-place derivative readouts. -Modules that derive physical quantities (energy, forces, stress) from -atomic-level predictions. These consume representation outputs and produce -observable physical quantities. +Public call shape (``model`` is always the first argument):: + + # energy only — no Derivative session + batch = EnergyReadout(model, method="func", backward=False)(batch) + + # energy + force as a sequential pair (one energy eval when backward=True) + batch = EnergyReadout(model, method="func", backward=True)(batch) + batch = ForceReadout(model, method="func")(batch) + +Or a potential that fixes forces at init (monomorphic pipeline, preferred):: + + model = PiNetPotential(..., compute_forces=True, method="func") + batch = model(batch) + model.compile() # optional torch.compile of that static forward + +Potentials and modes do not hand-roll the pass body: they bind one of the two +shared batch-level kernels (``energy_core(batch) -> batch`` in, batch with +``graphs.energy`` / ``atoms.forces`` out):: + + batch = grad_force_pass(self._write_energy, batch, detach_energy=False) + batch = func_force_pass(self._write_energy, batch) """ from molpot.derivation.energy import EnergyAggregation +from molpot.derivation.energy_readout import EnergyReadout from molpot.derivation.force import ( ForceDerivation, autograd_forces, + autograd_forces_from_energy, functorch_forces, functorch_forces_with_aux, ) +from molpot.derivation.force_readout import ForceReadout +from molpot.derivation.kernels import func_force_pass, grad_force_pass from molpot.derivation.stress import StressDerivation __all__ = [ "EnergyAggregation", + "EnergyReadout", + "ForceReadout", "ForceDerivation", "StressDerivation", "autograd_forces", + "autograd_forces_from_energy", + "func_force_pass", "functorch_forces", "functorch_forces_with_aux", + "grad_force_pass", ] diff --git a/src/molpot/derivation/derivative.py b/src/molpot/derivation/derivative.py new file mode 100644 index 0000000..69c5b4d --- /dev/null +++ b/src/molpot/derivation/derivative.py @@ -0,0 +1,39 @@ +"""Internal readout session (mode + model + sequential state). + +Not part of the public call shape — created inside :class:`EnergyReadout` / +:class:`ForceReadout` and stashed on the batch. +""" + +from __future__ import annotations + +from typing import Any, Literal + +from tensordict import TensorDict + +from molpot.derivation.modes.func import FuncMode +from molpot.derivation.modes.grad import GradMode + + +class Derivative: + """Session: one mode, one model, state for Energy → Force on a batch.""" + + def __init__( + self, + method: Literal["func", "grad"], + model: Any, + ) -> None: + if method not in ("func", "grad"): + raise ValueError(f"method must be 'func' or 'grad', got {method!r}") + self.method = method + self.model = model + self.mode: FuncMode | GradMode = FuncMode() if method == "func" else GradMode() + self._backward: bool = False + self._energy_ready: bool = False + self._lazy_func: bool = False + self._pos_leaf = None + + def run_energy(self, batch: TensorDict, *, backward: bool = False) -> TensorDict: + return self.mode.run_energy(self, batch, backward=backward) + + def run_forces(self, batch: TensorDict) -> TensorDict: + return self.mode.run_forces(self, batch) diff --git a/src/molpot/derivation/energy_readout.py b/src/molpot/derivation/energy_readout.py new file mode 100644 index 0000000..b6cbe83 --- /dev/null +++ b/src/molpot/derivation/energy_readout.py @@ -0,0 +1,89 @@ +"""EnergyReadout — peer step: write energy on batch (in-place). + +``__init__`` binds a **monomorphic** ``__call__`` from ``method`` / ``backward``. +Energy-only never touches a Derivative session. +""" + +from __future__ import annotations + +from typing import Any, Literal + +from tensordict import TensorDict + +from molpot.derivation.derivative import Derivative +from molpot.derivation.protocol import ( + absorb_model_output, + attach_session, + call_energy, + get_session, + has_energy, +) + + +class EnergyReadout: + """Materialise energy on ``batch``. + + Signature:: + + EnergyReadout(model, *, method="func", backward=False) + + Args: + model: Potential (energy core via ``_write_energy`` or ``forward``). + method: ``"func"`` or ``"grad"`` — must match a following ForceReadout + when ``backward=True``. + backward: If ``True``, prepare for a following ForceReadout (one energy + evaluation for the pair). If ``False``, energy-only — **no session**. + """ + + def __init__( + self, + model: Any, + *, + method: Literal["func", "grad"] = "func", + backward: bool = False, + ) -> None: + self.model = model + self.method = method + self.backward = backward + + # Init-time kernel selection — __call__ is a single bound path. + if not backward: + self._run = self._energy_only + elif method == "func": + self._run = self._func_lazy + elif method == "grad": + self._run = self._grad_with_graph + else: # pragma: no cover + raise ValueError(f"method must be 'func' or 'grad', got {method!r}") + + def __call__(self, batch: TensorDict) -> TensorDict: + return self._run(batch) + + def _energy_only(self, batch: TensorDict) -> TensorDict: + """Static energy path — no Derivative, no session (compile-friendly).""" + out = call_energy(self.model, batch) + batch = absorb_model_output(batch, out) + if not has_energy(batch): + raise RuntimeError( + "energy core must write batch['graphs','energy'] (or return a dict with 'energy')" + ) + return batch + + def _func_lazy(self, batch: TensorDict) -> TensorDict: + """Mark batch for a following ForceReadout (func has_aux flush).""" + session = self._session_for(batch) + return session.run_energy(batch, backward=True) + + def _grad_with_graph(self, batch: TensorDict) -> TensorDict: + """Energy on a requires_grad pos leaf for a following ForceReadout.""" + session = self._session_for(batch) + return session.run_energy(batch, backward=True) + + def _session_for(self, batch: TensorDict) -> Derivative: + session = get_session(batch) + if session is None or session.method != self.method or session.model is not self.model: + session = Derivative(method=self.method, model=self.model) + attach_session(batch, session) + else: + session.model = self.model + return session diff --git a/src/molpot/derivation/force.py b/src/molpot/derivation/force.py index 9ea9ef5..fc2421c 100644 --- a/src/molpot/derivation/force.py +++ b/src/molpot/derivation/force.py @@ -1,26 +1,31 @@ """Force derivation: ``F = -∂E/∂pos``. Single responsibility: atomic forces as the negative gradient of energy w.r.t. -atomic positions. Two **explicit** backends — the caller picks one per model; -there is no auto-detection and no silent fallback: - -* **functorch** (``torch.func.grad``) — for pure-PyTorch models (e.g. PiNet). - The transform is traced into the forward graph, so ``energy → force → loss`` - is a single backward and composes with ``torch.compile(fullgraph=True)``. - It does NOT work on cuEquivariance *fused* kernels: those register a legacy - ``autograd.Function`` without a ``setup_context`` staticmethod, which the - functorch transforms reject (pytorch#170834). Supports ``has_aux=True`` so - energy side-outputs can be returned with a single pass. - -* **autograd** (``torch.autograd.grad``) — for models built on cuEq fused ops - (e.g. MACE), matching the upstream MACE library. Eager-correct, trains (the - force stays connected to parameters for a double-backward force loss), and is - ``torch.compile``-able via ``torch._dynamo.allow_in_graph(torch.autograd.grad)`` - with the outer ``loss.backward()`` run eagerly. +atomic positions. Two **explicit, non-mixing** backends — the caller picks one +per model; there is no auto-detection and no silent fallback between them: + +* **functorch** (``torch.func.grad``) — pure-PyTorch graphs (e.g. PiNet). + Traced into the forward graph; ``energy → force → loss`` is one backward and + composes with ``torch.compile(fullgraph=True)``. Supports ``has_aux=True`` so + energy side-outputs come back with a single energy evaluation. Does **not** + work on cuEquivariance *fused* kernels (legacy ``autograd.Function`` without + ``setup_context`` — pytorch#170834). + +* **autograd** (``torch.autograd.grad``) — cuEq / MACE-style graphs. Eager + correct, ``create_graph`` keeps the force connected to parameters for a + force-supervised double-backward, and is ``torch.compile``-able via + ``torch._dynamo.allow_in_graph(torch.autograd.grad)`` with the outer + ``loss.backward()`` run eagerly. Also exposes + :func:`autograd_forces_from_energy` for the 1-pass composition when energy + was already materialised on a ``requires_grad`` position leaf. + +:class:`ForceDerivation` only **dispatches** to the chosen backend. It never +routes a ``functorch`` model through ``autograd`` helpers, or the reverse. Pick the backend explicitly: ``ForceDerivation(method="functorch")`` for PiNet, -``ForceDerivation(method="autograd")`` for MACE. The default is ``"autograd"`` -(correct for every model; only pure-torch graphs gain anything from functorch). +``ForceDerivation(method="autograd")`` for MACE. Default ``"autograd"`` (safe +for every model; choose ``"functorch"`` only for pure-torch graphs you want to +``torch.compile(fullgraph)``). Example: >>> deriv = ForceDerivation(method="autograd") @@ -36,6 +41,10 @@ import torch import torch.nn as nn +# --------------------------------------------------------------------------- +# functorch backend (only) +# --------------------------------------------------------------------------- + def functorch_forces( energy_fn: Callable[[torch.Tensor], torch.Tensor], @@ -57,22 +66,28 @@ def functorch_forces_with_aux( """``F = -∂E/∂pos`` plus auxiliary outputs via ``torch.func.grad(..., has_aux=True)``. ``energy_fn`` must return ``(scalar_energy, aux)``. Returns ``(forces, aux)``. + One energy evaluation — the 1-pass path for the functorch backend. """ grad, aux = torch.func.grad(energy_fn, has_aux=True)(pos) return -grad, aux +# --------------------------------------------------------------------------- +# autograd backend (only) +# --------------------------------------------------------------------------- + + def autograd_forces( energy_fn: Callable[[torch.Tensor], torch.Tensor], pos: torch.Tensor, ) -> torch.Tensor: """``F = -∂E/∂pos`` via ``torch.autograd.grad`` (cuEq / MACE). - ``create_graph`` follows the ambient grad state: training (grad enabled) - keeps the force connected to the parameters (mixed 2nd derivative - ``∂²E/∂pos∂θ``) so a force-loss ``.backward()`` reaches them; pure inference - (``torch.no_grad()``) detaches. Works on cuEq fused ops (they support - ordinary double backward). + Re-runs ``energy_fn`` on a fresh ``requires_grad`` leaf. ``create_graph`` + follows the ambient grad state: training (grad enabled) keeps the force + connected to parameters (mixed 2nd derivative ``∂²E/∂pos∂θ``) so a + force-loss ``.backward()`` reaches them; pure inference (``torch.no_grad()``) + detaches. Works on cuEq fused ops (ordinary double backward). """ create_graph = torch.is_grad_enabled() with torch.enable_grad(): @@ -81,6 +96,34 @@ def autograd_forces( return -grad +def autograd_forces_from_energy( + energy: torch.Tensor, + pos: torch.Tensor, + *, + create_graph: bool | None = None, +) -> torch.Tensor: + """``F = -∂E/∂pos`` when energy was already computed on ``pos`` (autograd 1-pass). + + Autograd-backend only. Pair with one ``energy_forward`` that used + ``pos`` as a ``requires_grad`` leaf — no second energy evaluation. + + Args: + energy: Per-graph energies ``(B,)`` or a scalar total. + pos: Position leaf that ``energy`` depends on; must require grad. + create_graph: Keep force connected to parameters for a force-supervised + ``loss.backward()``. Defaults to ambient ``torch.is_grad_enabled()``. + + Returns: + Atomic forces ``(N, 3)``. + """ + if create_graph is None: + create_graph = torch.is_grad_enabled() + scalar = energy if energy.ndim == 0 else energy.sum() + with torch.enable_grad(): + (grad,) = torch.autograd.grad(scalar, pos, create_graph=create_graph) + return -grad + + _BACKENDS: dict[str, Callable] = { "functorch": functorch_forces, "autograd": autograd_forces, @@ -88,14 +131,11 @@ def autograd_forces( class ForceDerivation(nn.Module): - """Compute forces as ``F = -∂E/∂pos`` with an explicit backend (no fallback). + """Dispatch ``F = -∂E/∂pos`` to exactly one backend — never both. Args: - method: ``"functorch"`` (``torch.func.grad``; pure-torch models such as - PiNet — compile-friendly single backward) or ``"autograd"`` - (``torch.autograd.grad``; cuEq/MACE models). Default ``"autograd"``, - correct for every model; choose ``"functorch"`` only for a pure-torch - energy graph you want to ``torch.compile(fullgraph)``. + method: ``"functorch"`` or ``"autograd"``. Fixed at construction; every + call on this instance stays on that backend. """ def __init__(self, method: Literal["functorch", "autograd"] = "autograd"): @@ -139,8 +179,8 @@ def forward( through it. pos: Atomic positions ``(N, 3)``. Does not need ``requires_grad``. has_aux: If ``True``, ``energy_fn`` returns ``(energy, aux)`` and this - method returns ``(forces, aux)``. Only supported for - ``method="functorch"`` (single-pass eval path for PiNet). + method returns ``(forces, aux)``. **functorch only** (1-pass + energy + side outputs). Returns: Atomic forces ``(N, 3)``, or ``(forces, aux)`` when ``has_aux=True``. @@ -148,8 +188,27 @@ def forward( if has_aux: if self.method != "functorch": raise ValueError( - "has_aux=True requires ForceDerivation(method='functorch'); " - f"got method={self.method!r}" + f"has_aux=True is functorch-only; got ForceDerivation(method={self.method!r})" ) return functorch_forces_with_aux(energy_fn, pos) return _BACKENDS[self.method](energy_fn, pos) + + def forces_from_energy( + self, + energy: torch.Tensor, + pos: torch.Tensor, + *, + create_graph: bool | None = None, + ) -> torch.Tensor: + """Autograd-backend 1-pass: energy already on ``pos``, return forces. + + Raises: + ValueError: If this instance is not ``method="autograd"``. + """ + if self.method != "autograd": + raise ValueError( + "forces_from_energy is autograd-only; " + f"got ForceDerivation(method={self.method!r}). " + "For functorch use forward(..., has_aux=True)." + ) + return autograd_forces_from_energy(energy, pos, create_graph=create_graph) diff --git a/src/molpot/derivation/force_readout.py b/src/molpot/derivation/force_readout.py new file mode 100644 index 0000000..5ab6b68 --- /dev/null +++ b/src/molpot/derivation/force_readout.py @@ -0,0 +1,67 @@ +"""ForceReadout — peer step: write forces on batch (in-place). + +``__init__`` binds a monomorphic ``__call__`` from ``method``. Prefer a prior +``EnergyReadout(..., backward=True)`` so the pair is one potential forward. +""" + +from __future__ import annotations + +from typing import Any, Literal + +from tensordict import TensorDict + +from molpot.derivation.derivative import Derivative +from molpot.derivation.protocol import attach_session, detach_session, get_session + + +class ForceReadout: + """Materialise forces on ``batch``. + + Signature:: + + ForceReadout(model, *, method="func") + + Prefer a prior ``EnergyReadout(model, method=..., backward=True)`` so the + pair is a single potential forward. + """ + + def __init__( + self, + model: Any, + *, + method: Literal["func", "grad"] = "func", + ) -> None: + self.model = model + self.method = method + if method == "func": + self._run = self._run_func + elif method == "grad": + self._run = self._run_grad + else: # pragma: no cover + raise ValueError(f"method must be 'func' or 'grad', got {method!r}") + + def __call__(self, batch: TensorDict) -> TensorDict: + return self._run(batch) + + def _run_func(self, batch: TensorDict) -> TensorDict: + session = self._session_for(batch) + try: + return session.run_forces(batch) + finally: + detach_session(batch) + + def _run_grad(self, batch: TensorDict) -> TensorDict: + session = self._session_for(batch) + try: + return session.run_forces(batch) + finally: + detach_session(batch) + + def _session_for(self, batch: TensorDict) -> Derivative: + session = get_session(batch) + if session is None or session.method != self.method or session.model is not self.model: + session = Derivative(method=self.method, model=self.model) + attach_session(batch, session) + else: + session.model = self.model + return session diff --git a/src/molpot/derivation/kernels.py b/src/molpot/derivation/kernels.py new file mode 100644 index 0000000..197baf2 --- /dev/null +++ b/src/molpot/derivation/kernels.py @@ -0,0 +1,220 @@ +"""Batch-level force passes shared by potentials and sequential readout modes. + +Where this sits relative to :mod:`molpot.derivation.force`: + +* ``force.py`` is the **tensor-level** layer — ``energy_fn(pos) -> scalar`` in, + ``forces (N, 3)`` out. It owns the only ``torch.autograd.grad`` / + ``torch.func.grad`` calls in this package. +* ``kernels.py`` (this module) is the **batch-level** layer — an *energy core* + ``energy_core(batch) -> batch`` in, the same batch back with + ``graphs.energy`` and ``atoms.forces`` written in place. It adds the three + things a batch pass needs and a tensor pass cannot express: position-leaf + ownership, the post-collate key writes, and the ``detach_energy`` policy. + +The kernels **compose** ``force.py``; they never re-derive a gradient. Adding a +third hand-rolled force path inside an encoder or potential is forbidden (see +``.claude/notes/notes.md`` §"Force derivation — dual explicit backends"). + +Backend boundary (pick one per model — no auto-detect, no fallback): + +* :func:`grad_force_pass` — ``torch.autograd.grad`` behind + :func:`~molpot.derivation.force.autograd_forces_from_energy`. One energy + forward, then one backward on the position leaf. Works on cuEquivariance + *fused* kernels (legacy ``autograd.Function`` without ``setup_context``, + pytorch#170834) and is the MACE-shaped path. +* :func:`func_force_pass` — ``torch.func.grad(..., has_aux=True)`` behind + :func:`~molpot.derivation.force.functorch_forces_with_aux`. The derivative is + traced into the forward graph, so ``energy → force → loss`` is a single + backward and the whole pass composes with ``torch.compile(fullgraph=True)``. + Pure-PyTorch energy graphs only (PiNet). Rejects cuEq fused kernels. + +Units (repo-wide): positions Å, energies eV, forces eV/Å. The kernels are +geometry- and unit-agnostic — they only differentiate whatever the energy core +wrote — so a core that emits kcal/mol yields kcal/(mol·Å) forces silently. Keep +the core on the repo convention. + +Compile constraints (hard): :func:`func_force_pass` sits on PiNet's +``fullgraph=True`` path. It must stay free of ``set_non_tensor``, ``.item()``, +tensor-value-dependent branching and session-dict lookups. The Python-level +branches in :func:`grad_force_pass` (``energy_core is None``, +``detach_energy``) are static per call site and never appear in the func path. +Bind ``energy_core`` once at construction (a bound method, or +``functools.partial(call_energy, model)``) rather than building a fresh lambda +per call, so the traced callable stays stable. + +State: neither kernel holds state, allocates modules, or touches the +``protocol._SESSIONS`` side channel — session bookkeeping stays in +:mod:`molpot.derivation.modes`. Both mutate ``batch`` in place and return it. +""" + +from __future__ import annotations + +from collections.abc import Callable + +import torch +from tensordict import TensorDict + +from molpot.derivation.force import autograd_forces_from_energy, functorch_forces_with_aux +from molpot.derivation.protocol import ( + ENERGY_KEY, + POS_KEY, + absorb_model_output, + has_energy, + write_forces, +) + +__all__ = ["func_force_pass", "grad_force_pass"] + +#: An energy core: reads ``atoms.pos`` off the batch and writes ``graphs.energy`` +#: (optionally ``atoms.energy``). May return the batch, a ``dict`` with an +#: ``"energy"`` entry, or ``None`` when it wrote in place. +EnergyCore = Callable[[TensorDict], "TensorDict | dict | None"] + + +def grad_force_pass( + energy_core: EnergyCore | None, + batch: TensorDict, + *, + create_graph: bool | None = None, + detach_energy: bool | None = None, +) -> TensorDict: + """One energy forward + ``torch.autograd.grad`` on the position leaf. + + Writes ``graphs.energy`` ``(B,)`` in eV (via the core) and + ``atoms.forces`` ``(N, 3)`` in eV/Å onto ``batch``, in place. + + Position-leaf ownership: if ``atoms.pos`` ``(N, 3)`` is not already a live + ``requires_grad`` **leaf**, this kernel makes one + (``pos.detach().requires_grad_(True)``) and writes it back to the batch, so + the energy the core computes is differentiable w.r.t. it. A leaf supplied by + the caller is reused as is — no second leaf, no lost graph. + + Args: + energy_core: Energy core ``energy_core(batch)`` writing ``graphs.energy``. + ``None`` means the energy is **already materialised** on a live + position leaf (the sequential ``EnergyReadout(backward=True)`` → + ``ForceReadout`` protocol): the forward is skipped and the existing + graph is differentiated, which is what makes that pair a single + model forward. + batch: Post-collate batch. Must carry ``atoms.pos``; ``graphs.energy`` + must already be present when ``energy_core is None``. + create_graph: Keep the force connected to the parameters (mixed second + derivative ``∂²E/∂pos∂θ``) so a force-supervised ``loss.backward()`` + reaches them. ``None`` follows the ambient ``torch.is_grad_enabled()`` + *at entry* — resolved before the internal ``enable_grad`` scope, so + calling under ``torch.no_grad()`` yields detached forces rather than + silently building a graph. + detach_energy: Leaf-ownership knob for the returned ``graphs.energy``. + + * ``False`` — never detach; energy stays attached for an energy loss + (PiNet and ``GradMode``). + * ``True`` — always detach (energy is a reported quantity only). + * ``None`` — detach **iff this call created the position leaf**, i.e. + the caller had no graph to lose. This is MACE's ``get_outputs`` + shape (``mace_matpes.py:368-377``). + + Returns: + The same ``batch`` object, with ``atoms.forces`` ``(N, 3)`` eV/Å written. + + Raises: + RuntimeError: If ``energy_core is None`` but ``atoms.pos`` is not a live + ``requires_grad`` leaf (nothing to differentiate), or if the core ran + and left no ``graphs.energy`` behind. + + Note: + The energy forward runs inside ``torch.enable_grad()``, so the pass + returns correct forces even under an ambient ``torch.no_grad()`` — the + gradient is a *value*, not a side effect of the caller's grad mode. + """ + pos = batch[POS_KEY] + owns_leaf = not (pos.requires_grad and pos.is_leaf) + + if energy_core is None and owns_leaf: + raise RuntimeError( + "grad_force_pass(energy_core=None) needs atoms.pos to be a live " + "requires_grad leaf carrying the energy graph; got a tensor that is " + "detached or non-leaf. Pass an energy_core, or materialise the " + "energy on a position leaf first." + ) + + # Resolve before the enable_grad scope: create_graph must reflect the + # caller's grad mode, not the one this kernel forces on for the derivative. + resolved_create_graph = torch.is_grad_enabled() if create_graph is None else create_graph + + if owns_leaf: + pos = pos.detach().requires_grad_(True) + batch[POS_KEY] = pos + + with torch.enable_grad(): + if energy_core is not None: + batch = absorb_model_output(batch, energy_core(batch)) + if not has_energy(batch): + raise RuntimeError( + "energy core must write batch['graphs','energy'] " + "(graphs.energy) or return a dict with 'energy'" + ) + forces = autograd_forces_from_energy( + batch[ENERGY_KEY], pos, create_graph=resolved_create_graph + ) + + should_detach = owns_leaf if detach_energy is None else detach_energy + if should_detach: + batch[ENERGY_KEY] = batch[ENERGY_KEY].detach() + write_forces(batch, forces) + return batch + + +def func_force_pass(energy_core: EnergyCore, batch: TensorDict) -> TensorDict: + """Single ``torch.func.grad(..., has_aux=True)`` pass — energy and force at once. + + The energy core is evaluated **once**, inside the functorch transform, on a + cloned batch whose ``atoms.pos`` is the differentiation variable; the filled + clone comes back as ``aux`` so ``graphs.energy`` ``(B,)`` eV (and + ``atoms.energy`` ``(N,)`` eV when the core writes it) can be copied onto the + real batch alongside ``atoms.forces`` ``(N, 3)`` eV/Å. + + Because the derivative is traced into the forward graph, an + ``energy → force → loss`` objective needs one ordinary ``backward()`` and the + whole pass survives ``torch.compile(fullgraph=True)``. There is deliberately + no ``detach_energy`` knob here: the only callers want the energy attached, + and cuEq fused kernels — the ones that need leaf-detaching — cannot use + functorch at all (pytorch#170834). + + Args: + energy_core: Energy core ``energy_core(batch)`` writing ``graphs.energy``. + Must be a pure-PyTorch graph and must recompute every + position-derived quantity (edge vectors, distances) from the + ``atoms.pos`` it is handed, or the gradient path is cut. + batch: Post-collate batch carrying ``atoms.pos`` ``(N, 3)`` Å. + ``requires_grad`` is not needed — ``torch.func.grad`` tracks the + input itself. + + Returns: + The same ``batch`` object, with ``graphs.energy``, optionally + ``atoms.energy``, and ``atoms.forces`` written in place. + + Raises: + RuntimeError: If the core leaves no ``graphs.energy`` on the batch. + """ + pos = batch[POS_KEY].detach() + base = batch.clone() + base[POS_KEY] = pos + + def energy_fn_aux(p: torch.Tensor) -> tuple[torch.Tensor, TensorDict]: + b = base.clone() + b[POS_KEY] = p + b = absorb_model_output(b, energy_core(b)) + if not has_energy(b): + raise RuntimeError( + "energy core must write batch['graphs','energy'] (graphs.energy) " + "inside the func energy path" + ) + return b[ENERGY_KEY].sum(), b + + # functorch_forces_with_aux already returns -grad — do not negate again. + forces, filled = functorch_forces_with_aux(energy_fn_aux, pos) + batch[ENERGY_KEY] = filled[ENERGY_KEY] + if "atoms" in filled.keys() and "energy" in filled["atoms"].keys(): + batch["atoms", "energy"] = filled["atoms", "energy"] + write_forces(batch, forces) + return batch diff --git a/src/molpot/derivation/modes/__init__.py b/src/molpot/derivation/modes/__init__.py new file mode 100644 index 0000000..48bd9b3 --- /dev/null +++ b/src/molpot/derivation/modes/__init__.py @@ -0,0 +1,6 @@ +"""Differentiation modes: FuncMode (torch.func) and GradMode (torch.autograd).""" + +from molpot.derivation.modes.func import FuncMode +from molpot.derivation.modes.grad import GradMode + +__all__ = ["FuncMode", "GradMode"] diff --git a/src/molpot/derivation/modes/func.py b/src/molpot/derivation/modes/func.py new file mode 100644 index 0000000..bdabf9b --- /dev/null +++ b/src/molpot/derivation/modes/func.py @@ -0,0 +1,72 @@ +"""FuncMode — energy/force peers via ``torch.func`` only. + +In-place batch contract. ``EnergyReadout(backward=True)`` is **lazy** +(no model call); ``ForceReadout`` flushes a single ``grad(..., has_aux=True)`` +pass that writes energy and forces. Energy-only uses ``backward=False``. +""" + +from __future__ import annotations + +from functools import partial + +from tensordict import TensorDict + +from molpot.derivation.kernels import func_force_pass +from molpot.derivation.protocol import ( + absorb_model_output, + call_energy, + has_energy, +) + + +class FuncMode: + """torch.func-only mode for sequential Energy/Force readouts.""" + + name: str = "func" + + def run_energy( + self, + deriv: object, + batch: TensorDict, + *, + backward: bool, + ) -> TensorDict: + model = deriv.model + if model is None: + raise RuntimeError("session model is not set") + + if backward: + if getattr(deriv, "_energy_ready", False) and not getattr(deriv, "_lazy_func", False): + raise RuntimeError( + "Energy already materialised without lazy-func state; " + "for method='func' use EnergyReadout(backward=True) " + "then ForceReadout (lazy+flush), or Energy only with " + "backward=False" + ) + deriv._lazy_func = True + deriv._backward = True + deriv._energy_ready = False + return batch + + deriv._lazy_func = False + deriv._backward = False + out = call_energy(model, batch) + batch = absorb_model_output(batch, out) + if not has_energy(batch): + raise RuntimeError( + "model.forward must write batch['graphs','energy'] (or return a dict with 'energy')" + ) + deriv._energy_ready = True + return batch + + def run_forces(self, deriv: object, batch: TensorDict) -> TensorDict: + model = deriv.model + if model is None: + raise RuntimeError("session model is not set") + + batch = func_force_pass(partial(call_energy, model), batch) + + deriv._lazy_func = False + deriv._energy_ready = True + deriv._backward = True + return batch diff --git a/src/molpot/derivation/modes/grad.py b/src/molpot/derivation/modes/grad.py new file mode 100644 index 0000000..e58b0e6 --- /dev/null +++ b/src/molpot/derivation/modes/grad.py @@ -0,0 +1,76 @@ +"""GradMode — energy/force peers via ``torch.autograd`` only. + +In-place batch contract. Sequential Energy then Force is one model forward when +``backward=True`` left a live ``E ← pos`` graph. +""" + +from __future__ import annotations + +from tensordict import TensorDict + +from molpot.derivation.kernels import grad_force_pass +from molpot.derivation.protocol import ( + POS_KEY, + absorb_model_output, + call_energy, + has_energy, +) + + +class GradMode: + """Autograd-only mode for sequential Energy/Force readouts.""" + + name: str = "grad" + + def run_energy( + self, + deriv: object, + batch: TensorDict, + *, + backward: bool, + ) -> TensorDict: + model = deriv.model + if model is None: + raise RuntimeError("session model is not set") + + pos = batch[POS_KEY] + if backward: + pos = pos.detach().requires_grad_(True) + batch[POS_KEY] = pos + deriv._pos_leaf = pos + deriv._backward = True + else: + deriv._backward = False + deriv._pos_leaf = None + + out = call_energy(model, batch) + batch = absorb_model_output(batch, out) + if not has_energy(batch): + raise RuntimeError( + "model.forward must write batch['graphs','energy'] (or return a dict with 'energy')" + ) + deriv._energy_ready = True + deriv._lazy_func = False + return batch + + def run_forces(self, deriv: object, batch: TensorDict) -> TensorDict: + if not getattr(deriv, "_energy_ready", False) or not getattr(deriv, "_backward", False): + batch = self.run_energy(deriv, batch, backward=True) + + pos = getattr(deriv, "_pos_leaf", None) + if pos is None: + pos = batch[POS_KEY] + if not pos.requires_grad: + raise RuntimeError( + "ForceReadout (grad mode) needs positions with requires_grad; " + "call EnergyReadout(..., backward=True) first" + ) + + # Energy is already materialised on the leaf run_energy(backward=True) + # installed — energy_core=None keeps the pair at one model forward. + return grad_force_pass( + None, + batch, + create_graph=bool(getattr(deriv.model, "training", False)), + detach_energy=False, + ) diff --git a/src/molpot/derivation/protocol.py b/src/molpot/derivation/protocol.py new file mode 100644 index 0000000..b01a383 --- /dev/null +++ b/src/molpot/derivation/protocol.py @@ -0,0 +1,184 @@ +"""Shared keys and helpers for writing energies and forces onto a batch. + +Every model in this repository communicates through one object: the +*post-collate batch*, a nested :class:`~tensordict.TensorDict` whose ``atoms`` +/ ``edges`` / ``graphs`` sub-dictionaries hold per-atom, per-edge and per-graph +tensors. :func:`molix.data.collate.collate_molecules` builds that object and +owns its schema (CLAUDE.md, "Post-collate batch schema"). This module owns only +the handful of keys a *potential* — a model mapping atomic positions to an +energy — writes back into it, plus the helpers that perform those writes in +place. + +No differentiation logic lives here: force values arrive already computed. +``F = -dE/dpos`` is taken by the two *modes* (``molpot.derivation.modes.func`` +and ``.grad`` — the :mod:`torch.func` and :mod:`torch.autograd` backends). Those +modes call a potential's **energy core**, the private ``_write_energy`` entry +point that computes the energy and nothing else, never the public ``forward``, +which may append further readouts a force pass must not re-run; +:func:`call_energy` is that dispatch. + +Readout sessions live in a module-level dictionary keyed by ``id(batch)`` rather +than on the batch itself. TensorDict's ``set_non_tensor`` would place a plain +Python object inside a tensor container, which Dynamo — the Python-bytecode +tracer behind :func:`torch.compile` — cannot trace, so it breaks compilation. +""" + +from __future__ import annotations + +from typing import Any, Protocol, runtime_checkable + +import torch +from tensordict import TensorDict + +# Nested post-collate keys (in-place writes). The schema is a molix contract; +# the keys live in molix.schema (single owner) and are re-exported here for +# potential-side code. +from molix.schema import ( + ATOMIC_ENERGY_KEY, + ENERGY_KEY, + FORCES_KEY, + POS_KEY, + has_energy, + has_forces, +) + +__all__ = [ + "ATOMIC_ENERGY_KEY", + "ENERGY_KEY", + "FORCES_KEY", + "POS_KEY", + "PotentialModule", + "absorb_model_output", + "attach_session", + "call_energy", + "detach_session", + "ensure_graphs", + "get_session", + "has_energy", + "has_forces", + "write_energy", + "write_forces", +] + +# Session side-channel: never batch.set_non_tensor (breaks torch.compile / Dynamo). +_SESSIONS: dict[int, Any] = {} + + +@runtime_checkable +class PotentialModule(Protocol): + """Potential that can write energy peers onto a batch.""" + + def forward(self, batch: TensorDict) -> TensorDict: + """Public entry (may compose readouts). Not used by modes for the core.""" + ... + + +def call_energy(model: Any, batch: TensorDict) -> TensorDict: + """Run the energy core only. + + Prefer ``model._write_energy(batch)`` when present (full potentials that + compose readouts in ``forward``). Fall back to ``model(batch)`` for toys + whose ``forward`` *is* the energy path. + """ + write = getattr(model, "_write_energy", None) + if callable(write): + return write(batch) + return model(batch) + + +def ensure_graphs(batch: TensorDict, num_graphs: int | None = None) -> TensorDict: + """Ensure a ``graphs`` sub-TensorDict exists for writing energy. + + An existing ``graphs`` namespace is returned untouched — its ``batch_size`` + is never rewritten, even when ``num_graphs`` disagrees with it. The + canonical producer of that namespace is + :func:`molix.data.collate.collate_molecules`; this helper only fills the + gap for batches assembled by hand. + + Args: + batch: Post-collate root batch, mutated in place. + num_graphs: Number of graphs ``B``. Creates the namespace with the + schema-conforming ``batch_size=[B]``. ``None`` (the default) + creates it with ``batch_size=[]`` instead, which does **not** + conform to CLAUDE.md's ``"graphs": TensorDict(batch_size=[B])`` + schema: a consumer reading ``batch["graphs"].batch_size[0]`` + (e.g. ``src/molzoo/pinet/potential.py``) raises ``IndexError`` on + it. The fallback exists only as transitional backward + compatibility for out-of-tree adapters — every in-tree caller + passes ``num_graphs``. + + Returns: + The same ``batch``, with a ``graphs`` namespace guaranteed present. + """ + if "graphs" not in batch.keys(): + batch["graphs"] = TensorDict(batch_size=[] if num_graphs is None else [num_graphs]) + return batch + + +def write_energy( + batch: TensorDict, + energy: torch.Tensor, + *, + atomic_energy: torch.Tensor | None = None, +) -> TensorDict: + """Write peer energy keys onto ``batch`` (in-place). + + Args: + batch: Post-collate root batch, mutated in place. + energy: Per-graph energy ``(B,)`` in eV. A 0-dim tensor is accepted and + leaves the created ``graphs`` namespace at ``batch_size=[]``. + atomic_energy: Optional per-atom energy ``(N,)`` in eV, written under + ``("atoms", "energy")``. + + Returns: + The same ``batch``. + """ + # B comes from the static shape of a (B,) energy — no host sync, and the + # same static-B convention as MACEPotential.energy_core(num_graphs=...). + ensure_graphs(batch, energy.shape[0] if energy.dim() == 1 else None) + batch[ENERGY_KEY] = energy + if atomic_energy is not None: + batch[ATOMIC_ENERGY_KEY] = atomic_energy + return batch + + +def write_forces(batch: TensorDict, forces: torch.Tensor) -> TensorDict: + """Write peer forces onto ``batch`` (in-place).""" + batch[FORCES_KEY] = forces + return batch + + +def absorb_model_output(batch: TensorDict, out: TensorDict | dict) -> TensorDict: + """Merge a model return value into ``batch``.""" + if out is batch: + return batch + if isinstance(out, TensorDict): + if has_energy(out): + batch[ENERGY_KEY] = out[ENERGY_KEY] + if "atoms" in out.keys() and "energy" in out["atoms"].keys(): + batch[ATOMIC_ENERGY_KEY] = out["atoms", "energy"] + return batch + if isinstance(out, dict): + if "energy" in out: + write_energy( + batch, + out["energy"], + atomic_energy=out.get("atomic_energy"), + ) + return batch + return batch + + +def attach_session(batch: TensorDict, session: Any) -> None: + """Store a readout session for ``batch`` (side-channel, not on TensorDict).""" + _SESSIONS[id(batch)] = session + + +def get_session(batch: TensorDict) -> Any | None: + """Return the readout session if EnergyReadout attached one.""" + return _SESSIONS.get(id(batch)) + + +def detach_session(batch: TensorDict) -> None: + """Drop session for ``batch`` (call after ForceReadout finishes).""" + _SESSIONS.pop(id(batch), None) diff --git a/src/molpot/graph/radius.py b/src/molpot/graph/radius.py deleted file mode 100644 index cfc24d9..0000000 --- a/src/molpot/graph/radius.py +++ /dev/null @@ -1,53 +0,0 @@ -"""Radius graph neighbor search for molpot.""" - -import torch - -from molix.F.locality import get_neighbor_pairs - - -def radius_graph( - pos: torch.Tensor, - batch: torch.Tensor, - cutoff: float, -) -> tuple[torch.Tensor, torch.Tensor]: - """Compute neighbor list for a set of positions. - - Args: - pos: Atomic positions [N, 3] - batch: Batch indices [N] - cutoff: Cutoff radius - - Returns: - edge_index: [num_edges, 2] - edge_vec: [num_edges, 3] (pos_j - pos_i) - """ - # Note: get_neighbor_pairs expects (positions, cutoff, ...) - # It returns (neighbors, deltas, distances, number_found_pairs) - # neighbors is [num_pairs, 2] - # deltas is [num_pairs, 3] (pos_j - pos_i) - - # We currently don't support batching in get_neighbor_pairs directly if - # it doesn't take batch index. But we can mask or handle it. - # Actually, molix backend usually handles PBC but maybe not multiple molecules - # unless they are separated by PBC. - - # For now, let's assume it handles the whole batch if positions are concatenated. - # But we need to ensure we don't find neighbors across different molecules. - - # A simple but potentially slow way is to use torch_cluster.radius_graph if available, - # but we should stick to molix as it's our backend. - - # Let's use get_neighbor_pairs and then mask pairs from different molecules. - neighbors, deltas, _, _ = get_neighbor_pairs(pos, cutoff) - - # Mask pairs that are across different molecules - # neighbors: [num_pairs, 2] - node_i = neighbors[:, 0] - node_j = neighbors[:, 1] - - mask = batch[node_i] == batch[node_j] - - edge_index = neighbors[mask] # [num_edges, 2] - edge_vec = deltas[mask] # [num_edges, 3] - - return edge_index, edge_vec diff --git a/src/molpot/heads/__init__.py b/src/molpot/heads/__init__.py index 14b7d55..fed1b35 100644 --- a/src/molpot/heads/__init__.py +++ b/src/molpot/heads/__init__.py @@ -22,6 +22,12 @@ PermMultipoleHead, PermMultipoleHeadSpec, ) +from molpot.heads.provenance import ( + CoverageRegime, + ParameterProvenance, + SupportClassifier, + attach_provenance, +) from molpot.heads.rescale import GlobalRescale, PerSpeciesScaleShift from molpot.heads.type import TypeHead @@ -30,6 +36,7 @@ "AtomicReferenceEnergy", "BondChargeHead", "ChargeResponseHead", + "CoverageRegime", "DipoleHead", "EdgeEnergyHead", "ElementAlphaTable", @@ -38,10 +45,13 @@ "GlobalRescale", "HardnessHead", "HardnessHeadSpec", + "ParameterProvenance", "PermMultipoleHead", "PermMultipoleHeadSpec", "PolarizabilityHead", "PolarizabilityHeadSpec", "PerSpeciesScaleShift", + "SupportClassifier", "TypeHead", + "attach_provenance", ] diff --git a/src/molpot/heads/charge_bond.py b/src/molpot/heads/charge_bond.py index a8a22fc..31302cf 100644 --- a/src/molpot/heads/charge_bond.py +++ b/src/molpot/heads/charge_bond.py @@ -60,6 +60,7 @@ import torch import torch.nn as nn +from molix import config from molix.F.scatter import scatter_sum from molpot.heads._common import graph_counts as _graph_counts @@ -78,7 +79,7 @@ class BondChargeHead(nn.Module): the bond distance scalar. hidden_dim: Hidden dimension of the bond-charge MLP. full_neighbor_list: When ``True`` (default, matches - :class:`molix.nn.locality.NeighborList`'s ``symmetry=True`` + :class:`molix.data.tasks.neighbor.NeighborList`'s ``symmetry=True`` default), the edge list contains both ``(i, j)`` and ``(j, i)`` and ``q_{ij}`` is scattered only to the source atom — antisymmetry of the bidirectional pair guarantees @@ -141,11 +142,11 @@ def __init__( in_dim = 2 * node_dim + 1 + self.edge_dim self.mlp = nn.Sequential( - nn.Linear(in_dim, hidden_dim), + nn.Linear(in_dim, hidden_dim, dtype=config.ftype), nn.SiLU(), - nn.Linear(hidden_dim, hidden_dim), + nn.Linear(hidden_dim, hidden_dim, dtype=config.ftype), nn.SiLU(), - nn.Linear(hidden_dim, 1), + nn.Linear(hidden_dim, 1, dtype=config.ftype), ) def _bond_charges( diff --git a/src/molpot/heads/charge_response.py b/src/molpot/heads/charge_response.py index 493ae30..0bf7071 100644 --- a/src/molpot/heads/charge_response.py +++ b/src/molpot/heads/charge_response.py @@ -23,10 +23,6 @@ ANG2BOHR = 1.8897259886 -def _scatter_sum(src: torch.Tensor, index: torch.Tensor, dim_size: int) -> torch.Tensor: - return scatter_sum(src, index, dim_size=dim_size) - - def _relative_atom_indices( batch: torch.Tensor, num_graphs: int ) -> tuple[torch.Tensor, torch.Tensor, int]: @@ -100,21 +96,21 @@ def __init__( self.epsilon = float(epsilon) self.atom_diag_mlp = nn.Sequential( - nn.Linear(node_scalar_dim, hidden_dim), + nn.Linear(node_scalar_dim, hidden_dim, dtype=config.ftype), nn.SiLU(), - nn.Linear(hidden_dim, 1), + nn.Linear(hidden_dim, 1, dtype=config.ftype), ) self.edge_scalar_mlp = nn.Sequential( - nn.Linear(edge_scalar_dim, hidden_dim), + nn.Linear(edge_scalar_dim, hidden_dim, dtype=config.ftype), nn.SiLU(), - nn.Linear(hidden_dim, 1), + nn.Linear(hidden_dim, 1, dtype=config.ftype), ) self.edge_vector_mlp = nn.Linear(edge_vector_dim, 1, dtype=config.ftype) if self.iso: self.iso_mlp = nn.Sequential( - nn.Linear(node_scalar_dim, hidden_dim), + nn.Linear(node_scalar_dim, hidden_dim, dtype=config.ftype), nn.SiLU(), - nn.Linear(hidden_dim, 1), + nn.Linear(hidden_dim, 1, dtype=config.ftype), ) _default_sigma = {1: 0.312, 6: 0.730, 7: 0.709, 8: 0.661, 16: 1.048, 17: 1.016} @@ -208,7 +204,7 @@ def forward( if self.iso: alpha_iso_atom = F.softplus(self.iso_mlp(node_scalars).squeeze(-1)) - alpha_iso = _scatter_sum(alpha_iso_atom, atom_batch, num_graphs) + alpha_iso = scatter_sum(alpha_iso_atom, atom_batch, dim_size=num_graphs) eye3 = torch.eye(3, dtype=alpha.dtype, device=alpha.device) alpha_iso_tensor = alpha_iso[:, None, None] * eye3[None] alpha = alpha + alpha_iso_tensor @@ -276,16 +272,25 @@ def _make_etainv(self, atom_diag, edge_index, edge_response, atom_batch, rel, co def _make_eem(self, atom_diag, sigma, pos, atom_batch, num_graphs, counts, nmax): eta = self._make_eta(atom_diag, sigma, pos, atom_batch, num_graphs, counts, nmax) - chi_blocks = [] - for b_idx, n in enumerate(counts.tolist()): - block = eta[b_idx, :n, :n] - inv = torch.linalg.inv(block) - chi_blocks.append(self._make_lrf(inv.unsqueeze(0))[0]) - padded_chi = torch.zeros_like(eta) - for b_idx, chi_b in enumerate(chi_blocks): - n = chi_b.shape[0] - padded_chi[b_idx, :n, :n] = chi_b - return padded_chi, eta + + # One batched inverse instead of ``num_graphs`` Python-level + # ``torch.linalg.inv`` calls (each a separate LAPACK/cuSOLVER launch, + # plus a ``counts.tolist()`` device→host sync to drive the loop). + # + # ``_make_eta`` masks each padded matrix to its active block, leaving a + # zero diagonal on the padding rows — singular, so it cannot be + # inverted as-is. Writing 1 onto the inactive diagonal makes it exactly + # block-diagonal ``[[A, 0], [0, I]]``, whose inverse is + # ``[[A⁻¹, 0], [0, I]]``; masking the identity block back to zero + # recovers the per-molecule ``A⁻¹`` padded with zeros. ``_make_lrf`` is + # row-sum based, so those zero rows contribute nothing and it yields + # the same result the per-block loop produced. + idx = torch.arange(nmax, device=eta.device) + active = idx.unsqueeze(0) < counts.unsqueeze(1) + eta_padded = eta + torch.diag_embed((~active).to(eta.dtype)) + inv = torch.linalg.inv(eta_padded) + block_mask = (active[:, :, None] & active[:, None, :]).to(eta.dtype) + return self._make_lrf(inv * block_mask), eta def _make_eta(self, atom_diag, sigma, pos, atom_batch, num_graphs, counts, nmax): dense_pos, _ = _dense_positions(pos, atom_batch, num_graphs) @@ -330,4 +335,4 @@ def _local_alpha(self, edge_index, edge_diff, edge_response, atom_batch, num_gra edge_vec.unsqueeze(-1) * edge_vec.unsqueeze(-2) ) edge_batch = atom_batch[src] - return _scatter_sum(weighted, edge_batch, num_graphs) + return scatter_sum(weighted, edge_batch, dim_size=num_graphs) diff --git a/src/molpot/heads/dipole.py b/src/molpot/heads/dipole.py index 714f4c4..01109fb 100644 --- a/src/molpot/heads/dipole.py +++ b/src/molpot/heads/dipole.py @@ -16,10 +16,6 @@ from molpot.heads._common import graph_counts as _graph_counts -def _scatter_sum(src: torch.Tensor, index: torch.Tensor, dim_size: int) -> torch.Tensor: - return scatter_sum(src, index, dim_size=dim_size) - - def _charge_neutralize( charges: torch.Tensor, batch: torch.Tensor, @@ -27,12 +23,12 @@ def _charge_neutralize( *, total_charge: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - pre = _scatter_sum(charges, batch, num_graphs) + pre = scatter_sum(charges, batch, dim_size=num_graphs) if total_charge is None: total_charge = torch.zeros_like(pre) correction = (pre - total_charge.view_as(pre)) / _graph_counts(batch, num_graphs) charges = charges - correction[batch] - post = _scatter_sum(charges, batch, num_graphs) + post = scatter_sum(charges, batch, dim_size=num_graphs) return charges, pre, post @@ -95,9 +91,9 @@ def __init__( if self.uses_ac: self.charge_mlp = nn.Sequential( - nn.Linear(node_scalar_dim, hidden_dim), + nn.Linear(node_scalar_dim, hidden_dim, dtype=config.ftype), nn.SiLU(), - nn.Linear(hidden_dim, 1), + nn.Linear(hidden_dim, 1, dtype=config.ftype), ) if self.uses_ad or self.uses_os: if node_vector_dim is None: @@ -112,9 +108,9 @@ def __init__( if edge_scalar_dim is None or edge_vector_dim is None: raise ValueError("edge_scalar_dim and edge_vector_dim required for BC variant.") self.bond_scalar_mlp = nn.Sequential( - nn.Linear(edge_scalar_dim, hidden_dim), + nn.Linear(edge_scalar_dim, hidden_dim, dtype=config.ftype), nn.SiLU(), - nn.Linear(hidden_dim, 1), + nn.Linear(hidden_dim, 1, dtype=config.ftype), ) self.bond_vector_mlp = nn.Linear( edge_vector_dim, @@ -180,27 +176,27 @@ def forward( out["charge_sum_pre_proj"] = pre out["charge_sum_post_proj"] = post out["atomic_charges"] = charges - dipole = dipole + _scatter_sum( + dipole = dipole + scatter_sum( charges.unsqueeze(-1) * pos, atom_batch, - num_graphs, + dim_size=num_graphs, ) if self.uses_os: assert oxidation is not None ox = oxidation.to(dtype=pos.dtype) out["oxidation_charges"] = ox - dipole = dipole + _scatter_sum( + dipole = dipole + scatter_sum( ox.unsqueeze(-1) * pos, atom_batch, - num_graphs, + dim_size=num_graphs, ) if self.uses_ad or self.uses_os: assert node_vectors is not None atomic_dipoles = self.atomic_dipole_gate(node_vectors).squeeze(-1) out["atomic_dipoles"] = atomic_dipoles - dipole = dipole + _scatter_sum(atomic_dipoles, atom_batch, num_graphs) + dipole = dipole + scatter_sum(atomic_dipoles, atom_batch, dim_size=num_graphs) if self.uses_bc: assert edge_scalars is not None and edge_index is not None @@ -214,7 +210,7 @@ def forward( out["bond_charges"] = bond_charge bond_dipoles = bond_charge.unsqueeze(-1) * edge_diff out["bond_dipoles"] = bond_dipoles - dipole = dipole + _scatter_sum(bond_dipoles, edge_batch, num_graphs) + dipole = dipole + scatter_sum(bond_dipoles, edge_batch, dim_size=num_graphs) if self.regularization: out["bond_charge_l2"] = bond_charge.square().mean() diff --git a/src/molpot/heads/energy.py b/src/molpot/heads/energy.py index 7d6af6f..4fc72b0 100644 --- a/src/molpot/heads/energy.py +++ b/src/molpot/heads/energy.py @@ -97,9 +97,9 @@ class AtomicEnergyMLP(nn.Module): def __init__(self, hidden_dim: int = 64): super().__init__() self.mlp = nn.Sequential( - nn.Linear(hidden_dim, hidden_dim), + nn.Linear(hidden_dim, hidden_dim, dtype=config.ftype), nn.SiLU(), - nn.Linear(hidden_dim, 1), + nn.Linear(hidden_dim, 1, dtype=config.ftype), ) def forward(self, atoms_h: torch.Tensor) -> torch.Tensor: @@ -125,9 +125,9 @@ class EnergyHead(nn.Module): def __init__(self, hidden_dim: int = 64): super().__init__() self.atomic_mlp = nn.Sequential( - nn.Linear(hidden_dim, hidden_dim), + nn.Linear(hidden_dim, hidden_dim, dtype=config.ftype), nn.SiLU(), - nn.Linear(hidden_dim, 1), + nn.Linear(hidden_dim, 1, dtype=config.ftype), ) def forward(self, atoms_h: torch.Tensor, graph_batch: torch.Tensor) -> torch.Tensor: diff --git a/src/molpot/heads/multipole.py b/src/molpot/heads/multipole.py index f9f4c23..bce5630 100644 --- a/src/molpot/heads/multipole.py +++ b/src/molpot/heads/multipole.py @@ -296,6 +296,7 @@ def __init__( irreps_in=cue.Irreps(G, [(self._u, "1o")]), irreps_out=cue.Irreps(G, [(1, "1o")]), layout=cue.ir_mul, + dtype=config.ftype, ) # Θ readout (l=2, parity even): slice 2e block, gate, collapse u·2e → 1·2e. @@ -313,6 +314,7 @@ def __init__( irreps_in=cue.Irreps(G, [(self._u, "2e")]), irreps_out=cue.Irreps(G, [(1, "2e")]), layout=cue.ir_mul, + dtype=config.ftype, ) @classmethod diff --git a/src/molpot/heads/provenance/__init__.py b/src/molpot/heads/provenance/__init__.py new file mode 100644 index 0000000..8bc0324 --- /dev/null +++ b/src/molpot/heads/provenance/__init__.py @@ -0,0 +1,22 @@ +"""Confidence, coverage, and parameter provenance surfaces. + +Thin audit types for learnable Class-I force fields. No active-learning +loop lives here — consumers attach records / regimes externally. + +Public API: + CoverageRegime, SupportClassifier, ParameterProvenance, attach_provenance + +Reference: + Spec: learnable-classical-ff-09-provenance +""" + +from molpot.heads.provenance.classifier import SupportClassifier +from molpot.heads.provenance.parameter import ParameterProvenance, attach_provenance +from molpot.heads.provenance.regime import CoverageRegime + +__all__ = [ + "CoverageRegime", + "ParameterProvenance", + "SupportClassifier", + "attach_provenance", +] diff --git a/src/molpot/heads/provenance/classifier.py b/src/molpot/heads/provenance/classifier.py new file mode 100644 index 0000000..f55ad96 --- /dev/null +++ b/src/molpot/heads/provenance/classifier.py @@ -0,0 +1,132 @@ +"""Map type confidences + chemical support → :class:`CoverageRegime`. + +Accepts confidences produced by +:meth:`~molrep.heads.type.TypeHead.decode_with_confidence` (or any equivalent +``(type_ids, confidence)`` pair). This module does **not** reimplement +softmax-max. + +Policy (defaults ``conf_in=0.8``, ``conf_near=0.5``): + +* ``conf < conf_near`` → :attr:`~CoverageRegime.UNKNOWN` +* inside support, ``conf >= conf_in`` → :attr:`~CoverageRegime.IN_SUPPORT` +* inside support, ``conf_near <= conf < conf_in`` → + :attr:`~CoverageRegime.NEAR_SUPPORT` +* outside support, ``conf >= conf_near`` → + :attr:`~CoverageRegime.EXTRAPOLATING` (never ``IN_SUPPORT``) +* ``support is None`` → confidence-only (membership treated as in-support) + +Reference: + Spec: learnable-classical-ff-09-provenance +""" + +from __future__ import annotations + +import math + +import torch + +from molpot.heads.provenance.regime import CoverageRegime +from molrep.embedding.support import ChemicalSupportIndex + +__all__ = ["SupportClassifier"] + + +class SupportClassifier: + """Map confidences + support index → :class:`CoverageRegime` per row. + + Args: + support: Optional :class:`~molrep.embedding.support.ChemicalSupportIndex`. + When ``None``, classification is confidence-only (membership + assumed in-support). + conf_in: Minimum confidence for :attr:`~CoverageRegime.IN_SUPPORT`. + conf_near: Minimum usable confidence; below → + :attr:`~CoverageRegime.UNKNOWN`. Mid band + ``[conf_near, conf_in)`` → :attr:`~CoverageRegime.NEAR_SUPPORT` + when in support. + """ + + def __init__( + self, + support: ChemicalSupportIndex | None, + *, + conf_in: float = 0.8, + conf_near: float = 0.5, + ) -> None: + if not 0.0 <= conf_near <= conf_in <= 1.0: + raise ValueError( + f"need 0 <= conf_near <= conf_in <= 1; got conf_near={conf_near}, conf_in={conf_in}" + ) + self._support = support + self.conf_in = float(conf_in) + self.conf_near = float(conf_near) + + @property + def support(self) -> ChemicalSupportIndex | None: + """Attached chemical support index, if any.""" + return self._support + + def classify( + self, + type_ids: torch.Tensor, + confidence: torch.Tensor, + *, + query: torch.Tensor | None = None, + ) -> list[CoverageRegime]: + """Classify each row into a coverage regime. + + Args: + type_ids: Predicted discrete ids ``(m,)`` (long / int). Used for + membership when ``query`` is omitted and a type-id bank is + attached. + confidence: Per-row confidence in ``[0, 1]``, shape ``(m,)``. + Prefer values from + :meth:`~molrep.heads.type.TypeHead.decode_with_confidence`. + query: Optional continuous embeddings ``(m, dim)``. When given, + membership uses L2 :meth:`~ChemicalSupportIndex.contains` + instead of type-id lookup. + + Returns: + List of :class:`CoverageRegime` of length ``m``. + """ + if type_ids.ndim != 1: + raise ValueError(f"type_ids must be 1-D; got {tuple(type_ids.shape)}") + if confidence.shape != type_ids.shape: + raise ValueError( + f"confidence shape {tuple(confidence.shape)} != " + f"type_ids shape {tuple(type_ids.shape)}" + ) + + m = int(type_ids.shape[0]) + if m == 0: + return [] + + in_support = self._membership(type_ids, query=query) + regimes: list[CoverageRegime] = [] + for i in range(m): + conf = float(confidence[i].item()) + inside = bool(in_support[i].item()) + regimes.append(self._regime_one(conf, inside)) + return regimes + + def _membership( + self, + type_ids: torch.Tensor, + *, + query: torch.Tensor | None, + ) -> torch.Tensor: + if self._support is None: + return torch.ones(type_ids.shape[0], dtype=torch.bool, device=type_ids.device) + if query is not None: + return self._support.contains(query) + return self._support.contains_type_ids(type_ids) + + def _regime_one(self, conf: float, inside: bool) -> CoverageRegime: + # NaN confidence is unusable. + if math.isnan(conf) or conf < self.conf_near: + return CoverageRegime.UNKNOWN + if inside: + if conf >= self.conf_in: + return CoverageRegime.IN_SUPPORT + return CoverageRegime.NEAR_SUPPORT + # Outside support: usable confidence → extrapolating (never IN_SUPPORT). + return CoverageRegime.EXTRAPOLATING diff --git a/src/molpot/heads/provenance/parameter.py b/src/molpot/heads/provenance/parameter.py new file mode 100644 index 0000000..81d2436 --- /dev/null +++ b/src/molpot/heads/provenance/parameter.py @@ -0,0 +1,91 @@ +"""Frozen parameter provenance records for Class-I FF audit trails. + +Reference: + Spec: learnable-classical-ff-09-provenance +""" + +from __future__ import annotations + +from collections.abc import MutableMapping, Sequence +from dataclasses import asdict, dataclass +from typing import Any + +from molpot.heads.provenance.regime import CoverageRegime +from molrep.condensation.classes import InteractionClass + +__all__ = ["ParameterProvenance", "attach_provenance"] + + +@dataclass(frozen=True) +class ParameterProvenance: + """Immutable audit record for one predicted / exported parameter row. + + Construct with :class:`ParameterProvenance` directly (no factory). Keep + lists of records **alongside** IR / ForceSpec payloads — do not fold into + a mega context bag. + + Args: + interaction: Interaction class name (e.g. ``"bond"``) or + :class:`~molrep.condensation.classes.InteractionClass`. + type_id: Discrete condensed type id, if any. + confidence: Softmax-max confidence from + :meth:`~molrep.heads.type.TypeHead.decode_with_confidence`, if any. + regime: Chemical-space coverage regime. + source: Provenance tag — e.g. ``"neural_continuous"``, + ``"condensed_type"``, ``"symbolic"``. + pattern: Optional SMARTS / pattern id when known. + ir_units: Unit system tag (default Class-I canonical). + notes: Free-form short note. + """ + + interaction: str | InteractionClass + type_id: int | None + confidence: float | None + regime: CoverageRegime + source: str + pattern: str | None = None + ir_units: str = "class_i_canonical" + notes: str = "" + + def as_dict(self) -> dict[str, Any]: + """JSON-friendly dict (enums → stable string names).""" + d = asdict(self) + inter = self.interaction + d["interaction"] = inter.value if isinstance(inter, InteractionClass) else str(inter) + d["regime"] = self.regime.name + return d + + +def attach_provenance( + target: MutableMapping[str, Any] | Any, + records: Sequence[ParameterProvenance], + *, + key: str = "provenance", +) -> Any: + """Attach provenance records as metadata **alongside** a payload. + + Does not invent a god context object: writes a JSON-friendly list under + ``key`` on a mapping, or on an object that exposes a ``metadata`` dict + (e.g. :class:`~molix.ff_export.force_spec.ForceSpec`). + + Args: + target: Mutable mapping, or object with ``.metadata`` mapping, or a + mapping-like ForceSpec-style record. + records: Provenance rows to attach. + key: Metadata key (default ``"provenance"``). + + Returns: + The same ``target`` (mutated) for call chaining. + """ + payload = [r.as_dict() for r in records] + if isinstance(target, MutableMapping): + target[key] = payload + return target + meta = getattr(target, "metadata", None) + if isinstance(meta, MutableMapping): + meta[key] = payload + return target + raise TypeError( + "attach_provenance requires a mutable mapping or an object " + f"with mutable .metadata; got {type(target)!r}" + ) diff --git a/src/molpot/heads/provenance/regime.py b/src/molpot/heads/provenance/regime.py new file mode 100644 index 0000000..c1ebefa --- /dev/null +++ b/src/molpot/heads/provenance/regime.py @@ -0,0 +1,32 @@ +"""Coverage regimes for chemical-space support of Class-I parameters. + +Reference: + Spec: learnable-classical-ff-09-provenance +""" + +from __future__ import annotations + +from enum import Enum + +__all__ = ["CoverageRegime"] + + +class CoverageRegime(Enum): + """Where a predicted parameter sits relative to chemical support. + + Members: + IN_SUPPORT: Type id / embedding inside the support index **and** + confidence at or above ``conf_in``. + NEAR_SUPPORT: Inside support with mid confidence + (``conf_near <= conf < conf_in``), or within a configured + distance margin when using continuous banks. + EXTRAPOLATING: Finite prediction outside support (confidence still + usable: ``conf >= conf_near``) — treat as extrapolation. + UNKNOWN: Below the confidence floor (``conf < conf_near``), missing + topology, or otherwise unusable for a regime call. + """ + + IN_SUPPORT = "in_support" + NEAR_SUPPORT = "near_support" + EXTRAPOLATING = "extrapolating" + UNKNOWN = "unknown" diff --git a/src/molpot/heads/type.py b/src/molpot/heads/type.py index 40e542e..178f5a6 100644 --- a/src/molpot/heads/type.py +++ b/src/molpot/heads/type.py @@ -3,6 +3,8 @@ import torch import torch.nn as nn +from molix import config + class TypeHead(nn.Module): """Predict atom types from atomic representations.""" @@ -16,9 +18,9 @@ def __init__(self, hidden_dim: int = 64, num_types: int = 100): """ super().__init__() self.module = nn.Sequential( - nn.Linear(hidden_dim, hidden_dim), + nn.Linear(hidden_dim, hidden_dim, dtype=config.ftype), nn.SiLU(), - nn.Linear(hidden_dim, num_types), + nn.Linear(hidden_dim, num_types, dtype=config.ftype), ) def forward(self, atoms_h: torch.Tensor) -> torch.Tensor: @@ -31,3 +33,23 @@ def forward(self, atoms_h: torch.Tensor) -> torch.Tensor: Type logits [N, num_types] """ return self.module(atoms_h) + + def decode(self, logits: torch.Tensor) -> torch.Tensor: + """Decode logits to type indices.""" + return logits.argmax(dim=-1) + + def decode_with_confidence( + self, + logits: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Decode with confidence scores (softmax max). + + Args: + logits: Type logits ``(N, num_types)``. + + Returns: + ``(indices, confidence)`` each of shape ``(N,)``. + """ + probs = torch.softmax(logits, dim=-1) + confidence, indices = probs.max(dim=-1) + return indices, confidence diff --git a/src/molpot/ir/__init__.py b/src/molpot/ir/__init__.py new file mode 100644 index 0000000..a723f1a --- /dev/null +++ b/src/molpot/ir/__init__.py @@ -0,0 +1,33 @@ +"""Potential Intermediate Representation (Class-I molecular mechanics). + +Names interaction bags, canonical units, and nonbonded scaling for the +learnable classical force-field stack. Bags are parameter containers only — +evaluation lives in :mod:`molpot.potentials`. +""" + +from molpot.ir.bags import ( + AngleBag, + BondBag, + ChargeBag, + ImproperHarmonicBag, + ImproperPeriodicBag, + LJBag, + ProperTorsionBag, +) +from molpot.ir.potential_ir import PotentialIR +from molpot.ir.scaling import NonbondedScaling +from molpot.ir.units import CLASS_I_CANONICAL, UnitTag + +__all__ = [ + "UnitTag", + "CLASS_I_CANONICAL", + "BondBag", + "AngleBag", + "ProperTorsionBag", + "ImproperPeriodicBag", + "ImproperHarmonicBag", + "LJBag", + "ChargeBag", + "NonbondedScaling", + "PotentialIR", +] diff --git a/src/molpot/ir/bags.py b/src/molpot/ir/bags.py new file mode 100644 index 0000000..8c5e50a --- /dev/null +++ b/src/molpot/ir/bags.py @@ -0,0 +1,203 @@ +"""Typed parameter bags for the Class-I Potential Intermediate Representation. + +Bags hold parameter tables only — they do not evaluate energy. Index +conventions (bond_index [2, N], proper_index [4, N], …) live on the kernel +side; bags validate only internal field-shape consistency. + +References: + OpenMM User Guide §19 "Forces" + SMIRNOFF specification (OpenFF) — proper / improper / nonbonded sections +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch + +__all__ = [ + "BondBag", + "AngleBag", + "ProperTorsionBag", + "ImproperPeriodicBag", + "ImproperHarmonicBag", + "LJBag", + "ChargeBag", +] + + +def _require_equal_1d(a: torch.Tensor, b: torch.Tensor, name_a: str, name_b: str) -> None: + if a.shape != b.shape: + raise ValueError( + f"{name_a} and {name_b} must have equal shape, got {tuple(a.shape)} vs {tuple(b.shape)}" + ) + + +def _validate_cosine_term_tables( + *, + k: torch.Tensor, + periodicity: torch.Tensor, + phase: torch.Tensor, + idivf: torch.Tensor, + bag_name: str, +) -> None: + """Validate multi-term cosine torsion tables. + + Expected shapes + --------------- + k, phase : ``[n_types, n_terms]`` + periodicity : ``[n_terms]`` + idivf : ``[n_types]`` + """ + if k.ndim != 2: + raise ValueError(f"{bag_name}.k must be [n_types, n_terms], got shape {tuple(k.shape)}") + if phase.shape != k.shape: + raise ValueError( + f"{bag_name}.phase must match k shape {tuple(k.shape)}, got {tuple(phase.shape)}" + ) + n_types, n_terms = k.shape + if periodicity.ndim != 1 or periodicity.shape[0] != n_terms: + raise ValueError( + f"{bag_name}.periodicity must be [n_terms]={n_terms}, " + f"got shape {tuple(periodicity.shape)}" + ) + if idivf.ndim != 1 or idivf.shape[0] != n_types: + raise ValueError( + f"{bag_name}.idivf must be [n_types]={n_types}, got shape {tuple(idivf.shape)}" + ) + if not bool((periodicity > 0).all()): + raise ValueError(f"{bag_name}.periodicity entries must be > 0, got {periodicity.tolist()}") + + +@dataclass +class BondBag: + """Harmonic bond parameters: ``E = ½ k (r − r₀)²``. + + Attributes: + k: Force constants ``[n_types]`` (or per-interaction). + r0: Equilibrium lengths ``[n_types]`` (same shape as ``k``). + """ + + k: torch.Tensor + r0: torch.Tensor + + def __post_init__(self) -> None: + _require_equal_1d(self.k, self.r0, "k", "r0") + + +@dataclass +class AngleBag: + """Harmonic angle parameters: ``E = ½ k (θ − θ₀)²``. + + Attributes: + k: Force constants ``[n_types]``. + theta0: Equilibrium angles in radians ``[n_types]``. + """ + + k: torch.Tensor + theta0: torch.Tensor + + def __post_init__(self) -> None: + _require_equal_1d(self.k, self.theta0, "k", "theta0") + + +@dataclass +class ProperTorsionBag: + """Class-I multi-term cosine proper torsion parameters. + + Energy per interaction type ``t`` and term ``m``:: + + E = Σ_m (k[t,m] / idivf[t]) * [1 + cos(n[m] * φ − γ[t,m])] + + Attributes: + k: Barrier heights ``[n_types, n_terms]``. + periodicity: Integer periodicities ``[n_terms]`` (shared across types). + phase: Phase offsets γ ``[n_types, n_terms]`` (radians). + idivf: AMBER scale / identity divisor ``s`` per type ``[n_types]``. + """ + + k: torch.Tensor + periodicity: torch.Tensor + phase: torch.Tensor + idivf: torch.Tensor + + def __post_init__(self) -> None: + _validate_cosine_term_tables( + k=self.k, + periodicity=self.periodicity, + phase=self.phase, + idivf=self.idivf, + bag_name="ProperTorsionBag", + ) + + +@dataclass +class ImproperPeriodicBag: + """Class-I multi-term cosine improper torsion parameters. + + Same field layout as :class:`ProperTorsionBag`. Central atom of the + improper is at row index 1 of ``improper_index`` (SMIRNOFF trefoil). + + Attributes: + k: Barrier heights ``[n_types, n_terms]``. + periodicity: Integer periodicities ``[n_terms]``. + phase: Phase offsets γ ``[n_types, n_terms]`` (radians). + idivf: Scale / identity divisor ``s`` per type ``[n_types]``. + """ + + k: torch.Tensor + periodicity: torch.Tensor + phase: torch.Tensor + idivf: torch.Tensor + + def __post_init__(self) -> None: + _validate_cosine_term_tables( + k=self.k, + periodicity=self.periodicity, + phase=self.phase, + idivf=self.idivf, + bag_name="ImproperPeriodicBag", + ) + + +@dataclass +class ImproperHarmonicBag: + """Harmonic improper parameters: ``E = ½ k (χ − χ₀)²``. + + Attributes: + k: Force constants ``[n_types]``. + chi0: Equilibrium improper angles in radians ``[n_types]``. + """ + + k: torch.Tensor + chi0: torch.Tensor + + def __post_init__(self) -> None: + _require_equal_1d(self.k, self.chi0, "k", "chi0") + + +@dataclass +class LJBag: + """Lennard-Jones 12-6 atom (or type) parameters. + + Attributes: + epsilon: Well depths ``[n]``. + sigma: Collision diameters ``[n]`` (same shape as ``epsilon``). + """ + + epsilon: torch.Tensor + sigma: torch.Tensor + + def __post_init__(self) -> None: + _require_equal_1d(self.epsilon, self.sigma, "epsilon", "sigma") + + +@dataclass +class ChargeBag: + """Partial charges. + + Attributes: + q: Charges in elementary charge units ``[n_atoms]`` (or per type). + """ + + q: torch.Tensor diff --git a/src/molpot/ir/potential_ir.py b/src/molpot/ir/potential_ir.py new file mode 100644 index 0000000..1763abf --- /dev/null +++ b/src/molpot/ir/potential_ir.py @@ -0,0 +1,62 @@ +"""PotentialIR — aggregate Class-I interaction bags + unit system. + +Missing bags are allowed (zero contribution from that term). The unit system +defaults to ``"class_i_canonical"``; unknown labels raise. + +References: + OpenMM User Guide §19 "Forces" + Spec: learnable-classical-ff-01-ir-kernels +""" + +from __future__ import annotations + +from dataclasses import dataclass, field + +from molpot.ir.bags import ( + AngleBag, + BondBag, + ChargeBag, + ImproperHarmonicBag, + ImproperPeriodicBag, + LJBag, + ProperTorsionBag, +) +from molpot.ir.scaling import NonbondedScaling + +__all__ = ["PotentialIR", "KNOWN_UNIT_SYSTEMS"] + +KNOWN_UNIT_SYSTEMS: frozenset[str] = frozenset({"class_i_canonical"}) + + +@dataclass +class PotentialIR: + """Aggregate holding optional Class-I parameter bags and scaling. + + Attributes: + bonds: Optional :class:`~molpot.ir.bags.BondBag`. + angles: Optional :class:`~molpot.ir.bags.AngleBag`. + propers: Optional :class:`~molpot.ir.bags.ProperTorsionBag`. + impropers_periodic: Optional :class:`~molpot.ir.bags.ImproperPeriodicBag`. + impropers_harmonic: Optional :class:`~molpot.ir.bags.ImproperHarmonicBag`. + lj: Optional :class:`~molpot.ir.bags.LJBag`. + charges: Optional :class:`~molpot.ir.bags.ChargeBag`. + scaling: Optional :class:`~molpot.ir.scaling.NonbondedScaling`. + unit_system: Unit-system label. Default ``"class_i_canonical"``. + """ + + bonds: BondBag | None = None + angles: AngleBag | None = None + propers: ProperTorsionBag | None = None + impropers_periodic: ImproperPeriodicBag | None = None + impropers_harmonic: ImproperHarmonicBag | None = None + lj: LJBag | None = None + charges: ChargeBag | None = None + scaling: NonbondedScaling | None = None + unit_system: str = field(default="class_i_canonical") + + def __post_init__(self) -> None: + if self.unit_system not in KNOWN_UNIT_SYSTEMS: + raise ValueError( + f"Unknown unit_system {self.unit_system!r}; " + f"known systems: {sorted(KNOWN_UNIT_SYSTEMS)}" + ) diff --git a/src/molpot/ir/scaling.py b/src/molpot/ir/scaling.py new file mode 100644 index 0000000..0b88ad3 --- /dev/null +++ b/src/molpot/ir/scaling.py @@ -0,0 +1,66 @@ +"""Nonbonded 1–2 / 1–3 / 1–4 scaling factors for Class-I force fields. + +Defaults follow AMBER / GAFF / SMIRNOFF Class-I conventions: + +| relation | scale_q | scale_lj | +|----------|---------|----------| +| 1–2 | 0 | 0 | +| 1–3 | 0 | 0 | +| 1–4 | 5/6 | 0.5 | + +References: + OpenMM User Guide §19 / SMIRNOFF nonbonded section + Cornell et al., JACS 1995 DOI 10.1021/ja00124a002 (AMBER) +""" + +from __future__ import annotations + +from dataclasses import dataclass + +__all__ = ["NonbondedScaling"] + + +@dataclass +class NonbondedScaling: + """Pair-exclusion scale factors for Coulomb and LJ interactions. + + Attributes: + scale_q_12: Coulomb scale for 1–2 (bonded) pairs. Default ``0``. + scale_q_13: Coulomb scale for 1–3 (angle) pairs. Default ``0``. + scale_q_14: Coulomb scale for 1–4 (proper torsion) pairs. Default ``5/6``. + scale_lj_12: LJ scale for 1–2 pairs. Default ``0``. + scale_lj_13: LJ scale for 1–3 pairs. Default ``0``. + scale_lj_14: LJ scale for 1–4 pairs. Default ``0.5``. + """ + + scale_q_12: float = 0.0 + scale_q_13: float = 0.0 + scale_q_14: float = 5.0 / 6.0 + scale_lj_12: float = 0.0 + scale_lj_13: float = 0.0 + scale_lj_14: float = 0.5 + + def scales_for(self, relation: str) -> tuple[float, float]: + """Return ``(scale_q, scale_lj)`` for a bonded-path relation. + + Args: + relation: One of ``"1-2"``, ``"1-3"``, ``"1-4"`` (also accepts + underscore forms ``"1_2"`` …). + + Returns: + Coulomb and LJ scale factors for that relation. + + Raises: + ValueError: Unknown relation label. + """ + key = relation.strip().replace("_", "-") + table = { + "1-2": (self.scale_q_12, self.scale_lj_12), + "1-3": (self.scale_q_13, self.scale_lj_13), + "1-4": (self.scale_q_14, self.scale_lj_14), + } + if key not in table: + raise ValueError( + f"Unknown nonbonded relation {relation!r}; expected one of {sorted(table)}" + ) + return table[key] diff --git a/src/molpot/ir/units.py b/src/molpot/ir/units.py new file mode 100644 index 0000000..f5395bb --- /dev/null +++ b/src/molpot/ir/units.py @@ -0,0 +1,36 @@ +"""Canonical unit tags for the Class-I Potential Intermediate Representation. + +Internal Class-I units (SI-free): kcal/mol, angstrom, elementary charge e, +radians. Downstream export (kJ/mol·nm) is owned by a later sub-spec. + +References: + OpenMM User Guide §19 "Forces" + Cornell et al., JACS 1995 DOI 10.1021/ja00124a002 (AMBER) +""" + +from enum import Enum +from types import MappingProxyType + +__all__ = ["UnitTag", "CLASS_I_CANONICAL"] + + +class UnitTag(Enum): + """Physical dimensions named by the Potential IR.""" + + energy = "energy" + length = "length" + charge = "charge" + angle = "angle" + force_const_bond = "force_const_bond" + force_const_angle = "force_const_angle" + torsion_barrier = "torsion_barrier" + + +CLASS_I_CANONICAL: MappingProxyType[str, str] = MappingProxyType( + { + "energy": "kcal/mol", + "length": "angstrom", + "charge": "e", + "angle": "radian", + } +) diff --git a/src/molpot/pooling/graph.py b/src/molpot/pooling/graph.py index 763ea31..83ce482 100644 --- a/src/molpot/pooling/graph.py +++ b/src/molpot/pooling/graph.py @@ -124,14 +124,17 @@ def forward( if x.dim() == 1: out = torch.full((dim_size,), float("-inf"), dtype=x.dtype, device=x.device) + index = batch else: out = torch.full((dim_size, x.shape[1]), float("-inf"), dtype=x.dtype, device=x.device) - - for mol_idx in range(dim_size): - mask = batch == mol_idx - if mask.any(): - out[mol_idx] = x[mask].max(dim=0)[0] if x.dim() > 1 else x[mask].max() - + index = batch.unsqueeze(-1).expand_as(x) + + # One fused scatter instead of a Python loop over molecules. The loop + # form issued ``dim_size`` kernel launches *and* a ``mask.any()`` + # device→host sync per molecule, making pooling cost scale with batch + # size. ``include_self=True`` against the ``-inf`` fill reproduces the + # loop's semantics exactly, including ``-inf`` rows for empty graphs. + out.scatter_reduce_(0, index, x, reduce="amax", include_self=True) return out def __repr__(self) -> str: diff --git a/src/molpot/potentials/__init__.py b/src/molpot/potentials/__init__.py index a407a25..edf4488 100644 --- a/src/molpot/potentials/__init__.py +++ b/src/molpot/potentials/__init__.py @@ -3,21 +3,26 @@ - LJ126: Lennard-Jones 12-6 - BondHarmonic: Harmonic bond stretching - AngleHarmonic: Harmonic angle bending -- DihedralHarmonic: Harmonic dihedral torsion +- DihedralHarmonic: Harmonic dihedral / improper-style form (not Class-I proper) +- ProperTorsionPeriodic: Class-I multi-term cosine proper torsion +- ImproperPeriodic: Class-I multi-term cosine improper +- ImproperHarmonic: Harmonic improper torsion - RepulsionExp6: Buckingham-style exponential repulsion - DispersionC6: Tang-Toennies C6 dispersion - ChargeTransfer: Charge-transfer potential - Polarization: Self-consistent induced-dipole polarization +- ZBLRepulsion: Ziegler-Biersack-Littmark screened nuclear repulsion """ from molpot.potentials.angles import AngleHarmonic from molpot.potentials.base import BasePotential from molpot.potentials.bonds import BondHarmonic -from molpot.potentials.dihedrals import DihedralHarmonic +from molpot.potentials.dihedrals import DihedralHarmonic, ProperTorsionPeriodic from molpot.potentials.elec import ( EwaldMultipoleEnergy, EwaldMultipoleEnergySpec, ) +from molpot.potentials.impropers import ImproperHarmonic, ImproperPeriodic from molpot.potentials.mixing import geometric_arithmetic_mixing from molpot.potentials.nonbonded import ( ChargeTransfer, @@ -28,6 +33,7 @@ repulsion_mixing, ) from molpot.potentials.polarization import Polarization +from molpot.potentials.repulsion import ZBLRepulsion from molpot.potentials.vdw import LJ126, lorentz_berthelot __all__ = [ @@ -37,6 +43,9 @@ "BondHarmonic", "AngleHarmonic", "DihedralHarmonic", + "ProperTorsionPeriodic", + "ImproperPeriodic", + "ImproperHarmonic", "EwaldMultipoleEnergy", "EwaldMultipoleEnergySpec", "geometric_arithmetic_mixing", @@ -47,4 +56,5 @@ "dispersion_mixing", "ct_mixing", "Polarization", + "ZBLRepulsion", ] diff --git a/src/molpot/potentials/dihedrals/__init__.py b/src/molpot/potentials/dihedrals/__init__.py index 19ff4dc..4f78840 100644 --- a/src/molpot/potentials/dihedrals/__init__.py +++ b/src/molpot/potentials/dihedrals/__init__.py @@ -1,5 +1,6 @@ """Dihedral potentials.""" from molpot.potentials.dihedrals.harmonic import DihedralHarmonic +from molpot.potentials.dihedrals.periodic import ProperTorsionPeriodic -__all__ = ["DihedralHarmonic"] +__all__ = ["DihedralHarmonic", "ProperTorsionPeriodic"] diff --git a/src/molpot/potentials/dihedrals/harmonic.py b/src/molpot/potentials/dihedrals/harmonic.py index 32a2228..c5b2907 100644 --- a/src/molpot/potentials/dihedrals/harmonic.py +++ b/src/molpot/potentials/dihedrals/harmonic.py @@ -6,18 +6,23 @@ class DihedralHarmonic(BasePotential): - """Harmonic dihedral torsion potential. + """Harmonic dihedral / improper-style form ``E = ½ k (φ − φ₀)²``. + + **Not** a Class-I proper torsion. For AMBER / GAFF / SMIRNOFF propers use + :class:`~molpot.potentials.dihedrals.periodic.ProperTorsionPeriodic`. + + Energy formula:: - Energy formula: E = 0.5 * k * (phi - phi0)^2 Parameters are stored as type-indexed vectors: - k[dihedral_type]: Force constant - phi0[dihedral_type]: Equilibrium dihedral angle (in radians) + + - ``k[dihedral_type]``: Force constant + - ``phi0[dihedral_type]``: Equilibrium dihedral angle (radians) Attributes: - k: Force constants [num_dihedral_types] - phi0: Equilibrium dihedral angles in radians [num_dihedral_types] + k: Force constants ``[num_dihedral_types]``. + phi0: Equilibrium dihedral angles in radians ``[num_dihedral_types]``. """ name = "dihedral_harmonic_torch" diff --git a/src/molpot/potentials/dihedrals/periodic.py b/src/molpot/potentials/dihedrals/periodic.py new file mode 100644 index 0000000..b3d0ad4 --- /dev/null +++ b/src/molpot/potentials/dihedrals/periodic.py @@ -0,0 +1,195 @@ +"""Class-I multi-term cosine proper torsion. + +Energy formula (OpenMM §19.4 / SMIRNOFF / AMBER):: + + E = Σ_torsions Σ_m (k[t,m] / s[t]) * [1 + cos(n[m] * φ − γ[t,m])] + +where ``t = proper_types[i]``, ``s = idivf``, ``n = periodicity``, +``γ = phase``, and ``φ`` is the i-j-k-l dihedral angle from the standard +atan2(n1, n2) construction (same geometry as :class:`DihedralHarmonic`). + +References: + OpenMM User Guide §19.4 + SMIRNOFF specification (OpenFF) — proper torsions + Cornell et al., JACS 1995 DOI 10.1021/ja00124a002 (AMBER) +""" + +from __future__ import annotations + +from typing import Any + +import torch + +from molpot.potentials.base import BasePotential + + +def _dihedral_phi( + pos: torch.Tensor, + index: torch.Tensor, +) -> torch.Tensor: + """Compute dihedral angles for COO ``index`` of shape ``[4, N]``. + + Args: + pos: Positions ``[n_atoms, 3]``. + index: Torsion indices ``[4, N]`` (i-j-k-l). + + Returns: + Dihedral angles ``φ`` of shape ``[N]`` in ``[-π, π]``. + """ + pos_i = pos[index[0]] + pos_j = pos[index[1]] + pos_k = pos[index[2]] + pos_l = pos[index[3]] + + b1 = pos_j - pos_i + b2 = pos_k - pos_j + b3 = pos_l - pos_k + + n1 = torch.cross(b1, b2, dim=-1) + n2 = torch.cross(b2, b3, dim=-1) + + n1_norm = n1 / (torch.norm(n1, dim=-1, keepdim=True) + 1e-8) + n2_norm = n2 / (torch.norm(n2, dim=-1, keepdim=True) + 1e-8) + + cos_phi = torch.sum(n1_norm * n2_norm, dim=-1) + cos_phi = torch.clamp(cos_phi, -1.0, 1.0) + + b2_norm = b2 / (torch.norm(b2, dim=-1, keepdim=True) + 1e-8) + cross_n1_n2 = torch.cross(n1_norm, n2_norm, dim=-1) + sin_phi = torch.sum(cross_n1_n2 * b2_norm, dim=-1) + + return torch.atan2(sin_phi, cos_phi) + + +def _require_index_4_by_n(index: torch.Tensor, name: str) -> None: + if index.ndim != 2 or index.shape[0] != 4: + raise ValueError( + f"{name} must be COO [4, N] (i-j-k-l); got shape {tuple(index.shape)}. " + "A geometric edge_index [E, 2] or row-major [N, 4] is not a torsion list." + ) + + +class ProperTorsionPeriodic(BasePotential): + """Class-I multi-term cosine proper torsion. + + Energy:: + + E = Σ_n (k_n / s) * [1 + cos(n * φ − γ_n)] + + Parameters are type-indexed tables: + + - ``k``: ``[n_types, n_terms]`` + - ``periodicity``: ``[n_terms]`` (shared integer periods) + - ``phase``: ``[n_types, n_terms]`` (radians) + - ``idivf``: ``[n_types]`` AMBER scale factor ``s`` + + Attributes: + k: Barrier heights ``[n_types, n_terms]``. + periodicity: Periodicities ``[n_terms]``. + phase: Phase offsets ``[n_types, n_terms]``. + idivf: Identity / scale divisors ``[n_types]``. + """ + + name = "proper_torsion_periodic_torch" + type = "dihedral" + + k: torch.Tensor + periodicity: torch.Tensor + phase: torch.Tensor + idivf: torch.Tensor + + def __init__( + self, + k: torch.Tensor, + periodicity: torch.Tensor, + phase: torch.Tensor, + idivf: torch.Tensor, + ) -> None: + """Initialize ProperTorsionPeriodic. + + Args: + k: Barrier heights ``[n_types, n_terms]``. + periodicity: Integer periods ``[n_terms]`` (must be > 0). + phase: Phase offsets γ in radians ``[n_types, n_terms]``. + idivf: Scale factors ``s`` per type ``[n_types]``. + """ + super().__init__() + + if k.ndim != 2: + raise ValueError(f"k must be [n_types, n_terms], got shape {tuple(k.shape)}") + if phase.shape != k.shape: + raise ValueError(f"phase must match k shape {tuple(k.shape)}, got {tuple(phase.shape)}") + n_types, n_terms = k.shape + if periodicity.ndim != 1 or periodicity.shape[0] != n_terms: + raise ValueError( + f"periodicity must be [n_terms]={n_terms}, got shape {tuple(periodicity.shape)}" + ) + if idivf.ndim != 1 or idivf.shape[0] != n_types: + raise ValueError(f"idivf must be [n_types]={n_types}, got shape {tuple(idivf.shape)}") + if not bool((periodicity > 0).all()): + raise ValueError(f"periodicity entries must be > 0, got {periodicity.tolist()}") + + self.register_buffer("k", k) + self.register_buffer("periodicity", periodicity) + self.register_buffer("phase", phase) + self.register_buffer("idivf", idivf) + + def forward(self, data: dict[str, Any] | None = None, **kwargs: Any) -> torch.Tensor: + """Compute Class-I multi-term proper torsion energy. + + Args: + data: Optional dictionary with molecular fields. + **kwargs: Explicit tensors: + - pos: Positions ``[n_atoms, 3]``. + - proper_index: COO proper indices ``[4, N]`` (i-j-k-l). + - proper_types: Proper types ``[N]``. + + Returns: + Total proper torsion energy (scalar). + + Raises: + ValueError: Missing inputs or ``proper_index`` not ``[4, N]``. + """ + pos = kwargs.get("pos") + proper_index = kwargs.get("proper_index") + proper_types = kwargs.get("proper_types") + + if data is not None and isinstance(data, dict): + if pos is None: + pos = data.get("pos") + if pos is None and isinstance(data.get("atoms"), dict): + pos = data["atoms"].get("pos") + if proper_index is None: + proper_index = data.get("proper_index") + if proper_types is None: + proper_types = data.get("proper_types") + + if pos is None or proper_index is None or proper_types is None: + raise ValueError("ProperTorsionPeriodic requires pos, proper_index, and proper_types.") + + if not isinstance(pos, torch.Tensor): + pos = torch.from_numpy(pos).float() + proper_index = torch.from_numpy(proper_index).long() + proper_types = torch.from_numpy(proper_types).long() + + _require_index_4_by_n(proper_index, "proper_index") + + if proper_index.size(1) == 0: + return torch.tensor(0.0, device=pos.device, dtype=pos.dtype) + + phi = _dihedral_phi(pos, proper_index) # [N] + + # Look up per-torsion parameter rows: [N, n_terms] + k_t = self.k[proper_types] + phase_t = self.phase[proper_types] + s_t = self.idivf[proper_types] # [N] + n = self.periodicity.to(dtype=pos.dtype) # [n_terms] + + # Broadcast: phi [N, 1], n [1, n_terms] → [N, n_terms] + arg = n.unsqueeze(0) * phi.unsqueeze(1) - phase_t + term = (k_t / s_t.unsqueeze(1)) * (1.0 + torch.cos(arg)) + return term.sum() + + def __repr__(self) -> str: + n_types, n_terms = self.k.shape + return f"ProperTorsionPeriodic(n_types={n_types}, n_terms={n_terms})" diff --git a/src/molpot/potentials/elec/lib/kspace_filter.py b/src/molpot/potentials/elec/lib/kspace_filter.py index 24fd928..7393890 100644 --- a/src/molpot/potentials/elec/lib/kspace_filter.py +++ b/src/molpot/potentials/elec/lib/kspace_filter.py @@ -1,11 +1,19 @@ """K-space convolution filter for reciprocal-space potential calculations.""" +import os from typing import Optional import torch from molpot.potentials.elec.lib.kvectors import generate_kvectors_for_mesh +#: Default for :attr:`KSpaceFilter.check_nan`, read once at import. +#: The guard is a full-mesh reduction — a device sync barrier on the PME hot +#: path — so it is off unless ``MOLNEX_PME_CHECK_NAN=1`` is exported. Set +#: ``filter.check_nan = True`` (or pass ``check_nan=True``) to enable it for a +#: single instance without touching the environment. +CHECK_NAN_DEFAULT: bool = os.environ.get("MOLNEX_PME_CHECK_NAN", "") == "1" + class KSpaceKernel(torch.nn.Module): """Base class for a reciprocal-space convolution kernel. @@ -48,6 +56,12 @@ class KSpaceFilter(torch.nn.Module): ``"ortho"``). ifft_norm: Normalization for inverse FFT (``"forward"``, ``"backward"``, ``"ortho"``). + check_nan: Raise :class:`ValueError` when :meth:`forward` produces NaNs + (usually an unsuitable ``mesh_spacing``). The scan is a full-mesh + reduction and therefore a device sync barrier on the PME hot path, + so it defaults to :data:`CHECK_NAN_DEFAULT` (off unless + ``MOLNEX_PME_CHECK_NAN=1``). Enable it per instance when debugging + a grid; the attribute stays writable after construction. """ def __init__( @@ -57,11 +71,13 @@ def __init__( kernel: KSpaceKernel, fft_norm: str = "ortho", ifft_norm: str = "ortho", + check_nan: bool | None = None, ): super().__init__() self._fft_norm = fft_norm self._ifft_norm = ifft_norm + self.check_nan = CHECK_NAN_DEFAULT if check_nan is None else bool(check_nan) if fft_norm not in ["ortho", "forward", "backward"]: raise ValueError(f"Invalid option '{fft_norm}' for the `fft_norm` parameter.") if ifft_norm not in ["ortho", "forward", "backward"]: @@ -119,15 +135,9 @@ def forward(self, mesh_values: torch.Tensor) -> torch.Tensor: s=mesh_values.shape[-3:], ) - # Full-mesh NaN scan is a device sync barrier on the PME hot path. - # Opt in only when debugging (env MOLNEX_PME_CHECK_NAN=1) or via - # ``self.check_nan = True`` on the instance. - check_nan = getattr(self, "check_nan", False) - if not check_nan: - import os - - check_nan = os.environ.get("MOLNEX_PME_CHECK_NAN", "") == "1" - if check_nan and torch.isnan(result).any(): + # Full-mesh NaN scan is a device sync barrier on the PME hot path — + # opt in per instance (see ``check_nan`` in the class docstring). + if self.check_nan and torch.isnan(result).any(): raise ValueError( "NaNs detected in the k-space filter result. This are probably caused " "by an unsuitable `mesh_spacing`, resulting in a problematic grid of " @@ -174,6 +184,7 @@ class P3MKSpaceFilter(KSpaceFilter): ifft_norm: Inverse FFT normalization. mode: 0 for potential, 1 for energy, 2 for dipolar torque, 3 for dipolar force. differential_order: Order of the difference operator (1-6). + check_nan: See :class:`KSpaceFilter`. Reference: Deserno, M. & Holm, C. J. Chem. Phys. 109, 7678–7693 (1998) @@ -189,6 +200,7 @@ def __init__( ifft_norm: str = "ortho", mode: int = 0, differential_order: int = 2, + check_nan: bool | None = None, ): self.interpolation_nodes = interpolation_nodes if mode not in [0, 1, 2, 3]: @@ -200,7 +212,7 @@ def __init__( ) self.differential_order = differential_order - super().__init__(cell, ns_mesh, kernel, fft_norm, ifft_norm) + super().__init__(cell, ns_mesh, kernel, fft_norm, ifft_norm, check_nan) self.register_buffer( "_diff_coeff", torch.tensor( diff --git a/src/molpot/potentials/impropers/__init__.py b/src/molpot/potentials/impropers/__init__.py new file mode 100644 index 0000000..e236a9b --- /dev/null +++ b/src/molpot/potentials/impropers/__init__.py @@ -0,0 +1,6 @@ +"""Improper torsion potentials (periodic cosine + harmonic).""" + +from molpot.potentials.impropers.harmonic import ImproperHarmonic +from molpot.potentials.impropers.periodic import ImproperPeriodic + +__all__ = ["ImproperPeriodic", "ImproperHarmonic"] diff --git a/src/molpot/potentials/impropers/harmonic.py b/src/molpot/potentials/impropers/harmonic.py new file mode 100644 index 0000000..49c3a80 --- /dev/null +++ b/src/molpot/potentials/impropers/harmonic.py @@ -0,0 +1,123 @@ +"""Harmonic improper torsion: ``E = ½ k (χ − χ₀)²``. + +``improper_index [4, N]`` uses molrs layout ``[center, i, j, k]`` (center at +**row 0**). χ is the dihedral angle of the ordered quartet +``(i, center, j, k)``. + +This is the CHARMM / some-SMIRNOFF harmonic improper form — **not** a Class-I +proper torsion. For AMBER/GAFF/SMIRNOFF propers use +:class:`~molpot.potentials.dihedrals.periodic.ProperTorsionPeriodic`. + +References: + OpenMM User Guide §19.5 (harmonic impropers) + CHARMM force field improper form + molrs ``Topology`` impropers ``[center, i, j, k]`` +""" + +from __future__ import annotations + +from typing import Any + +import torch + +from molpot.potentials.base import BasePotential +from molpot.potentials.dihedrals.periodic import _dihedral_phi, _require_index_4_by_n +from molpot.potentials.impropers.periodic import _improper_to_dihedral_index + + +class ImproperHarmonic(BasePotential): + """Harmonic improper torsion potential. + + Energy formula:: + + E = 0.5 * k * (χ − χ₀)² + + Parameters are type-indexed vectors: + + - ``k``: force constants ``[n_types]`` + - ``chi0``: equilibrium improper angles in radians ``[n_types]`` + + Central atom is at **row index 0** of ``improper_index`` (molrs layout). + + Attributes: + k: Force constants ``[n_types]``. + chi0: Equilibrium improper angles in radians ``[n_types]``. + """ + + name = "improper_harmonic_torch" + type = "improper" + + k: torch.Tensor + chi0: torch.Tensor + + def __init__(self, k: torch.Tensor, chi0: torch.Tensor) -> None: + """Initialize ImproperHarmonic. + + Args: + k: Force constant vector ``[n_types]``. + chi0: Equilibrium improper angles in radians ``[n_types]``. + """ + super().__init__() + + if k.shape != chi0.shape: + raise ValueError( + f"k and chi0 must have same shape, got k: {k.shape}, chi0: {chi0.shape}" + ) + if k.ndim != 1: + raise ValueError(f"k must be 1D vector [n_types], got shape {k.shape}") + + self.register_buffer("k", k) + self.register_buffer("chi0", chi0) + + def forward(self, data: dict[str, Any] | None = None, **kwargs: Any) -> torch.Tensor: + """Compute harmonic improper energy. + + Args: + data: Optional dictionary with molecular fields. + **kwargs: Explicit tensors: + - pos: Positions ``[n_atoms, 3]``. + - improper_index: COO improper indices ``[4, N]`` molrs layout + ``[center, i, j, k]`` (center at **row 0**). + - improper_types: Improper types ``[N]``. + + Returns: + Total improper energy (scalar). + + Raises: + ValueError: Missing inputs or ``improper_index`` not ``[4, N]``. + """ + pos = kwargs.get("pos") + improper_index = kwargs.get("improper_index") + improper_types = kwargs.get("improper_types") + + if data is not None and isinstance(data, dict): + if pos is None: + pos = data.get("pos") + if pos is None and isinstance(data.get("atoms"), dict): + pos = data["atoms"].get("pos") + if improper_index is None: + improper_index = data.get("improper_index") + if improper_types is None: + improper_types = data.get("improper_types") + + if pos is None or improper_index is None or improper_types is None: + raise ValueError("ImproperHarmonic requires pos, improper_index, and improper_types.") + + if not isinstance(pos, torch.Tensor): + pos = torch.from_numpy(pos).float() + improper_index = torch.from_numpy(improper_index).long() + improper_types = torch.from_numpy(improper_types).long() + + _require_index_4_by_n(improper_index, "improper_index") + + if improper_index.size(1) == 0: + return torch.tensor(0.0, device=pos.device, dtype=pos.dtype) + + chi = _dihedral_phi(pos, _improper_to_dihedral_index(improper_index)) + k_imp = self.k[improper_types] + chi0_imp = self.chi0[improper_types] + energy_per = 0.5 * k_imp * (chi - chi0_imp) ** 2 + return energy_per.sum() + + def __repr__(self) -> str: + return f"ImproperHarmonic(n_types={len(self.k)})" diff --git a/src/molpot/potentials/impropers/periodic.py b/src/molpot/potentials/impropers/periodic.py new file mode 100644 index 0000000..7f14326 --- /dev/null +++ b/src/molpot/potentials/impropers/periodic.py @@ -0,0 +1,162 @@ +"""Class-I multi-term cosine improper torsion. + +Same cosine energy form as :class:`~molpot.potentials.dihedrals.periodic.ProperTorsionPeriodic`:: + + E = Σ_impropers Σ_m (k[t,m] / s[t]) * [1 + cos(n[m] * φ − γ[t,m])] + +**Index layout (molrs Topology source of truth):** ``improper_index`` is +``[4, N]`` with rows ``[center, i, j, k]`` — **center at row 0**, matching +molrs ``Topology`` impropers and molpy force-field topology. Internally the +dihedral angle is evaluated on the ordered quartet ``(i, center, j, k)`` +(rows 1, 0, 2, 3). OpenFF trefoil reordering is an export adapter, not a +second in-kernel layout. + +References: + OpenMM User Guide §19.5 + SMIRNOFF specification (OpenFF) — impropers + molrs ``Topology`` impropers ``[center, i, j, k]`` +""" + +from __future__ import annotations + +from typing import Any + +import torch + +from molpot.potentials.base import BasePotential +from molpot.potentials.dihedrals.periodic import _dihedral_phi, _require_index_4_by_n + + +def _improper_to_dihedral_index(improper_index: torch.Tensor) -> torch.Tensor: + """Map molrs ``[center, i, j, k]`` → dihedral order ``[i, center, j, k]``.""" + return torch.stack( + ( + improper_index[1], + improper_index[0], + improper_index[2], + improper_index[3], + ), + dim=0, + ) + + +class ImproperPeriodic(BasePotential): + """Class-I multi-term cosine improper torsion. + + Energy:: + + E = Σ_n (k_n / s) * [1 + cos(n * φ − γ_n)] + + Central atom of each improper is at **row index 0** of ``improper_index`` + (molrs ``[center, i, j, k]``). + + Attributes: + k: Barrier heights ``[n_types, n_terms]``. + periodicity: Periodicities ``[n_terms]``. + phase: Phase offsets ``[n_types, n_terms]``. + idivf: Scale factors ``[n_types]``. + """ + + name = "improper_periodic_torch" + type = "improper" + + k: torch.Tensor + periodicity: torch.Tensor + phase: torch.Tensor + idivf: torch.Tensor + + def __init__( + self, + k: torch.Tensor, + periodicity: torch.Tensor, + phase: torch.Tensor, + idivf: torch.Tensor, + ) -> None: + """Initialize ImproperPeriodic. + + Args: + k: Barrier heights ``[n_types, n_terms]``. + periodicity: Integer periods ``[n_terms]`` (must be > 0). + phase: Phase offsets γ in radians ``[n_types, n_terms]``. + idivf: Scale factors ``s`` per type ``[n_types]``. + """ + super().__init__() + + if k.ndim != 2: + raise ValueError(f"k must be [n_types, n_terms], got shape {tuple(k.shape)}") + if phase.shape != k.shape: + raise ValueError(f"phase must match k shape {tuple(k.shape)}, got {tuple(phase.shape)}") + n_types, n_terms = k.shape + if periodicity.ndim != 1 or periodicity.shape[0] != n_terms: + raise ValueError( + f"periodicity must be [n_terms]={n_terms}, got shape {tuple(periodicity.shape)}" + ) + if idivf.ndim != 1 or idivf.shape[0] != n_types: + raise ValueError(f"idivf must be [n_types]={n_types}, got shape {tuple(idivf.shape)}") + if not bool((periodicity > 0).all()): + raise ValueError(f"periodicity entries must be > 0, got {periodicity.tolist()}") + + self.register_buffer("k", k) + self.register_buffer("periodicity", periodicity) + self.register_buffer("phase", phase) + self.register_buffer("idivf", idivf) + + def forward(self, data: dict[str, Any] | None = None, **kwargs: Any) -> torch.Tensor: + """Compute Class-I multi-term improper torsion energy. + + Args: + data: Optional dictionary with molecular fields. + **kwargs: Explicit tensors: + - pos: Positions ``[n_atoms, 3]``. + - improper_index: COO improper indices ``[4, N]`` molrs layout + ``[center, i, j, k]`` (center at **row 0**). + - improper_types: Improper types ``[N]``. + + Returns: + Total improper torsion energy (scalar). + + Raises: + ValueError: Missing inputs or ``improper_index`` not ``[4, N]``. + """ + pos = kwargs.get("pos") + improper_index = kwargs.get("improper_index") + improper_types = kwargs.get("improper_types") + + if data is not None and isinstance(data, dict): + if pos is None: + pos = data.get("pos") + if pos is None and isinstance(data.get("atoms"), dict): + pos = data["atoms"].get("pos") + if improper_index is None: + improper_index = data.get("improper_index") + if improper_types is None: + improper_types = data.get("improper_types") + + if pos is None or improper_index is None or improper_types is None: + raise ValueError("ImproperPeriodic requires pos, improper_index, and improper_types.") + + if not isinstance(pos, torch.Tensor): + pos = torch.from_numpy(pos).float() + improper_index = torch.from_numpy(improper_index).long() + improper_types = torch.from_numpy(improper_types).long() + + _require_index_4_by_n(improper_index, "improper_index") + + if improper_index.size(1) == 0: + return torch.tensor(0.0, device=pos.device, dtype=pos.dtype) + + # molrs [center,i,j,k] → dihedral (i,center,j,k) for atan2 path. + phi = _dihedral_phi(pos, _improper_to_dihedral_index(improper_index)) + + k_t = self.k[improper_types] + phase_t = self.phase[improper_types] + s_t = self.idivf[improper_types] + n = self.periodicity.to(dtype=pos.dtype) + + arg = n.unsqueeze(0) * phi.unsqueeze(1) - phase_t + term = (k_t / s_t.unsqueeze(1)) * (1.0 + torch.cos(arg)) + return term.sum() + + def __repr__(self) -> str: + n_types, n_terms = self.k.shape + return f"ImproperPeriodic(n_types={n_types}, n_terms={n_terms})" diff --git a/src/molpot/potentials/repulsion.py b/src/molpot/potentials/repulsion.py new file mode 100644 index 0000000..463714b --- /dev/null +++ b/src/molpot/potentials/repulsion.py @@ -0,0 +1,120 @@ +"""Ziegler-Biersack-Littmark screened-nuclear-repulsion pair term. + +Short-range repulsion that keeps a learned potential physical when two nuclei +approach closer than any training configuration ever did. MACE's foundation +models add it to the interaction energy so high-temperature MD and structure +search cannot fall into an unphysical attractive well at small ``r``. + +Reference: + Ziegler, Biersack, Littmark. "The Stopping and Range of Ions in Solids" + Pergamon, 1985. + Batatia et al. "A foundation model for atomistic materials chemistry" + (MACE-MP-0), § pair repulsion. https://arxiv.org/abs/2401.00096 +""" + +from __future__ import annotations + +import torch +import torch.nn as nn + +from molix import config +from molix.F.scatter import scatter_sum_compile_safe as _scatter_sum +from molrep.embedding.covalent import covalent_radii +from molrep.embedding.cutoff import PolynomialCutoff + +# Universal ZBL screening function coefficients: phi(x) = sum_k c_k exp(-b_k x). +_SCREENING_C = (0.1818, 0.5099, 0.2802, 0.02817) +_SCREENING_B = (3.2, 0.9423, 0.4028, 0.2016) +# e^2 / (4 pi eps_0) in eV*Angstrom. +_COULOMB_CONSTANT = 14.3996 +# Bohr radius in Angstrom, the length unit of the ZBL screening length. +_BOHR = 0.529 + + +class ZBLRepulsion(nn.Module): + """Per-atom ZBL repulsion energy with a polynomial cutoff envelope. + + For each edge ``(i, j)``: + + .. math:: + + a_{ij} &= \\frac{a_p \\cdot a_0}{Z_i^{a_e} + Z_j^{a_e}} \\\\ + V_{ij} &= \\frac{k\\,Z_i Z_j}{r_{ij}} + \\sum_k c_k e^{-b_k r_{ij} / a_{ij}} \\; u(r_{ij}; R_i + R_j) + + where ``u`` is the polynomial envelope (:meth:`PolynomialCutoff.envelope`) + over the pair's summed covalent radii. Each edge contributes ``V_ij / 2`` to + its target atom, so a bidirectional edge list double-counts exactly back to + the pair energy. + + Args: + exponent: Polynomial-envelope exponent ``p``. + a_exp: Exponent of ``Z`` in the screening length. + a_prefactor: Prefactor of the screening length. + max_z: Largest atomic number in the covalent-radius table. + trainable: If ``True``, ``a_exp`` / ``a_prefactor`` become learnable. + """ + + def __init__( + self, + *, + exponent: int = 6, + a_exp: float = 0.300, + a_prefactor: float = 0.4543, + max_z: int = 118, + trainable: bool = False, + ) -> None: + super().__init__() + ftype = config.ftype + self.exponent = int(exponent) + + self.register_buffer("c", torch.tensor(_SCREENING_C, dtype=ftype)) + self.c: torch.Tensor + self.register_buffer("covalent_radii", covalent_radii(max_z)) + self.covalent_radii: torch.Tensor + + for name, value in (("a_exp", a_exp), ("a_prefactor", a_prefactor)): + tensor = torch.tensor(float(value), dtype=ftype) + if trainable: + setattr(self, name, nn.Parameter(tensor)) + else: + self.register_buffer(name, tensor) + self.a_exp: torch.Tensor + self.a_prefactor: torch.Tensor + + def forward( + self, + r: torch.Tensor, + Z: torch.Tensor, + edge_index: torch.Tensor, + ) -> torch.Tensor: + """Compute the per-atom ZBL repulsion energy. + + Args: + r: Edge distances ``(E,)``. + Z: Atomic numbers ``(N,)``. + edge_index: ``(E, 2)`` with ``[:, 0]`` = source, ``[:, 1]`` = target + (the repo-wide edge convention). + + Returns: + Per-atom energy ``(N,)`` in eV. + """ + z = Z.long() + source, target = edge_index[:, 0], edge_index[:, 1] + z_source = z[source].to(r.dtype) + z_target = z[target].to(r.dtype) + + screening_length = ( + self.a_prefactor + * _BOHR + / (torch.pow(z_source, self.a_exp) + torch.pow(z_target, self.a_exp)) + ) + x = r / screening_length + phi = sum(self.c[k] * torch.exp(-_SCREENING_B[k] * x) for k in range(len(_SCREENING_B))) + + radii = self.covalent_radii + pair_cutoff = radii[z[source]] + radii[z[target]] + envelope = PolynomialCutoff.envelope(r, pair_cutoff, self.exponent) + + v_edges = 0.5 * (_COULOMB_CONSTANT * z_source * z_target) / r * phi * envelope + return _scatter_sum(v_edges, target, Z.shape[0]) diff --git a/src/molrep/analysis/__init__.py b/src/molrep/analysis/__init__.py new file mode 100644 index 0000000..06967f8 --- /dev/null +++ b/src/molrep/analysis/__init__.py @@ -0,0 +1,17 @@ +"""molrep.analysis — latent atom embedding diagnostics (Validation D/E light). + +Must not import molpot / molzoo / condensation. +""" + +from molrep.analysis.artifacts import LatentAnalysisArtifacts +from molrep.analysis.latent_store import AtomLatentTable +from molrep.analysis.projection import LatentPCA2D +from molrep.analysis.type_purity import NearestNeighbourTypePurity, TypePurityReport + +__all__ = [ + "AtomLatentTable", + "NearestNeighbourTypePurity", + "TypePurityReport", + "LatentPCA2D", + "LatentAnalysisArtifacts", +] diff --git a/src/molrep/analysis/artifacts.py b/src/molrep/analysis/artifacts.py new file mode 100644 index 0000000..578fe15 --- /dev/null +++ b/src/molrep/analysis/artifacts.py @@ -0,0 +1,73 @@ +"""Write latent analysis metrics and point artifacts.""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any + +from molrep.analysis.latent_store import AtomLatentTable +from molrep.analysis.projection import LatentPCA2D +from molrep.analysis.type_purity import NearestNeighbourTypePurity + +__all__ = ["LatentAnalysisArtifacts"] + + +class LatentAnalysisArtifacts: + """Emit purity metrics + latent_points.jsonl + type_purity.txt. + + Optionally uses :class:`~molix.io.metrics.MetricsWriter` when available. + """ + + def __init__(self, out_dir: str | Path, *, k: int = 1) -> None: + self.out_dir = Path(out_dir) + self.k = k + + def write(self, table: AtomLatentTable) -> dict[str, Any]: + self.out_dir.mkdir(parents=True, exist_ok=True) + purity = NearestNeighbourTypePurity(k=self.k).score(table) + coords = LatentPCA2D().project(table) + + metrics = { + "nn_type_purity_mean": purity.mean_purity, + "nn_type_purity_k": purity.k, + "n_atoms": table.n_atoms, + "n_labeled": purity.n_labeled, + "n_scored": purity.n_scored, + } + + # points jsonl + points_path = self.out_dir / "latent_points.jsonl" + with points_path.open("w") as fh: + for i in range(table.n_atoms): + rec = { + "i": i, + "molecule_id": table.molecule_id[i], + "x": float(coords[i, 0]), + "y": float(coords[i, 1]), + } + if table.ref_atom_type is not None: + rec["ref_atom_type"] = int(table.ref_atom_type[i]) + fh.write(json.dumps(rec) + "\n") + + purity_path = self.out_dir / "type_purity.txt" + purity_path.write_text( + f"mean_purity={purity.mean_purity}\nk={purity.k}\n" + f"n_scored={purity.n_scored}\nn_labeled={purity.n_labeled}\n" + ) + + # Optional MetricsWriter (record root); always also dump a plain jsonl line. + try: + from molix.io.metrics import MetricsWriter + + mw = MetricsWriter(self.out_dir) + for key, value in metrics.items(): + if value is None: + continue + if isinstance(value, (int, float)): + mw.scalar(key, float(value)) + except Exception: + pass + (self.out_dir / "metrics_summary.jsonl").write_text(json.dumps(metrics) + "\n") + + return metrics diff --git a/src/molrep/analysis/latent_store.py b/src/molrep/analysis/latent_store.py new file mode 100644 index 0000000..9536ed0 --- /dev/null +++ b/src/molrep/analysis/latent_store.py @@ -0,0 +1,43 @@ +"""Store per-atom latent embeddings with molecule and optional type labels.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch + +__all__ = ["AtomLatentTable"] + + +@dataclass +class AtomLatentTable: + """In-memory table of per-atom representations. + + Attributes: + features: Embedding matrix ``(N, D)``. + molecule_id: Per-atom molecule id strings length ``N``. + ref_atom_type: Optional integer type labels ``(N,)``; ``-1`` = unlabeled. + """ + + features: torch.Tensor + molecule_id: list[str] + ref_atom_type: torch.Tensor | None = None + + def __post_init__(self) -> None: + if self.features.ndim != 2: + raise ValueError(f"features must be (N, D), got {tuple(self.features.shape)}") + n = self.features.shape[0] + if len(self.molecule_id) != n: + raise ValueError("molecule_id length must match N") + if self.ref_atom_type is not None: + if self.ref_atom_type.numel() != n: + raise ValueError("ref_atom_type length must match N") + self.ref_atom_type = self.ref_atom_type.long().reshape(-1) + + @property + def n_atoms(self) -> int: + return int(self.features.shape[0]) + + @property + def dim(self) -> int: + return int(self.features.shape[1]) diff --git a/src/molrep/analysis/projection.py b/src/molrep/analysis/projection.py new file mode 100644 index 0000000..25182d5 --- /dev/null +++ b/src/molrep/analysis/projection.py @@ -0,0 +1,35 @@ +"""Deterministic 2-D PCA projection for latent tables (no plotting).""" + +from __future__ import annotations + +import torch + +from molrep.analysis.latent_store import AtomLatentTable + +__all__ = ["LatentPCA2D"] + + +class LatentPCA2D: + """Project atom features to 2-D via SVD (deterministic).""" + + def project(self, table: AtomLatentTable) -> torch.Tensor: + """Return coordinates ``(N, 2)``. + + Raises: + ValueError: If fewer than 2 atoms or feature dim < 1. + """ + x = table.features.detach().float() + if x.shape[0] < 1: + raise ValueError("need at least 1 atom") + x = x - x.mean(dim=0, keepdim=True) + # SVD of (N, D) + # Use full_matrices=False + if x.shape[0] == 1 or x.shape[1] == 1: + out = torch.zeros(x.shape[0], 2, dtype=x.dtype, device=x.device) + out[:, 0] = x[:, 0] if x.shape[1] >= 1 else 0.0 + return out + _, _, vh = torch.linalg.svd(x, full_matrices=False) + comps = vh[:2].T # (D, 2) + if comps.shape[1] == 1: + comps = torch.nn.functional.pad(comps, (0, 1)) + return x @ comps diff --git a/src/molrep/analysis/type_purity.py b/src/molrep/analysis/type_purity.py new file mode 100644 index 0000000..19532a9 --- /dev/null +++ b/src/molrep/analysis/type_purity.py @@ -0,0 +1,65 @@ +"""Leave-one-out k-NN type purity on atom latents.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch + +from molrep.analysis.latent_store import AtomLatentTable + +__all__ = ["NearestNeighbourTypePurity", "TypePurityReport"] + + +@dataclass(frozen=True) +class TypePurityReport: + mean_purity: float | None + k: int + n_scored: int + n_labeled: int + + +class NearestNeighbourTypePurity: + """Leave-one-out L2 k-NN type purity. + + For each labeled atom, find k nearest other labeled atoms (by L2 on + features) and score the fraction whose type matches. Mean over scored + atoms. Unlabeled ``y_i == -1`` are excluded. + + Args: + k: Neighbourhood size (default 1). + """ + + def __init__(self, k: int = 1) -> None: + if k < 1: + raise ValueError("k must be >= 1") + self.k = k + + def score(self, table: AtomLatentTable) -> TypePurityReport: + if table.ref_atom_type is None: + return TypePurityReport(mean_purity=None, k=self.k, n_scored=0, n_labeled=0) + y = table.ref_atom_type + labeled = y >= 0 + n_labeled = int(labeled.sum().item()) + if n_labeled < 2: + return TypePurityReport(mean_purity=None, k=self.k, n_scored=0, n_labeled=n_labeled) + feats = table.features + idx = torch.where(labeled)[0] + purities: list[float] = [] + for i in idx.tolist(): + others = idx[idx != i] + if others.numel() == 0: + continue + d = torch.norm(feats[others] - feats[i], dim=-1) + kk = min(self.k, int(others.numel())) + nn = others[torch.topk(d, k=kk, largest=False).indices] + match = (y[nn] == y[i]).float().mean().item() + purities.append(float(match)) + if not purities: + return TypePurityReport(mean_purity=None, k=self.k, n_scored=0, n_labeled=n_labeled) + return TypePurityReport( + mean_purity=sum(purities) / len(purities), + k=self.k, + n_scored=len(purities), + n_labeled=n_labeled, + ) diff --git a/src/molrep/chem/__init__.py b/src/molrep/chem/__init__.py new file mode 100644 index 0000000..490336b --- /dev/null +++ b/src/molrep/chem/__init__.py @@ -0,0 +1,32 @@ +"""molrep.chem — continuous chemical perception (no energy / no molpot). + +Public surface: + AtomChemEmbedding, BondChemEmbedding, context builders, ChemEmbeddings, + ChemEncoder. +""" + +from molrep.chem.context import ( + AngleContext, + BondContext, + ImproperContext, + ProperContext, +) +from molrep.chem.embed import AtomChemEmbedding, BondChemEmbedding +from molrep.chem.encoder import ChemEncoder +from molrep.chem.features import ChemEmbeddings +from molrep.chem.typing_metrics import TypingRecoveryMetrics, TypingRecoveryReport +from molrep.chem.typing_probe import AtomTypeReadout + +__all__ = [ + "AtomChemEmbedding", + "BondChemEmbedding", + "BondContext", + "AngleContext", + "ProperContext", + "ImproperContext", + "ChemEmbeddings", + "ChemEncoder", + "AtomTypeReadout", + "TypingRecoveryMetrics", + "TypingRecoveryReport", +] diff --git a/src/molrep/chem/context.py b/src/molrep/chem/context.py new file mode 100644 index 0000000..186c984 --- /dev/null +++ b/src/molrep/chem/context.py @@ -0,0 +1,231 @@ +"""Symmetry-aware interaction context builders for valence terms. + +Contexts pool atom (and bond) embeddings into per-interaction feature +vectors used by continuous MM parameter heads. Each builder enforces the +documented reverse / outer-swap symmetry of its interaction class. + +Symmetry contracts +------------------ +- Bond: ``(i,j) ↔ (j,i)`` +- Angle: ``(i,j,k) ↔ (k,j,i)`` +- Proper: ``(i,j,k,l) ↔ (l,k,j,i)`` +- Improper: outer-leg permutations with center fixed (``atomi`` = center, + molrs ``[center, i, j, k]`` layout) + +No energy, no molpot imports. +""" + +from __future__ import annotations + +import torch +import torch.nn as nn + +from molix import config +from molrep.chem.embed import BondChemEmbedding + + +def _mlp(in_dim: int, hidden_dim: int, out_dim: int) -> nn.Sequential: + return nn.Sequential( + nn.Linear(in_dim, hidden_dim, dtype=config.ftype), + nn.SiLU(), + nn.Linear(hidden_dim, out_dim, dtype=config.ftype), + ) + + +class BondContext(nn.Module): + """Bond interaction context — symmetric endpoint pooling. + + Thin wrapper around :class:`~molrep.chem.embed.BondChemEmbedding` so all + four valence contexts share a uniform ``forward(h_atom, *indices)`` surface. + """ + + def __init__( + self, + *, + atom_dim: int, + bond_dim: int = 32, + hidden_dim: int | None = None, + num_bond_types: int = 0, + ) -> None: + super().__init__() + self.bond_dim = int(bond_dim) + self.embed = BondChemEmbedding( + atom_dim=atom_dim, + bond_dim=bond_dim, + hidden_dim=hidden_dim, + num_bond_types=num_bond_types, + ) + + def forward( + self, + h_atom: torch.Tensor, + atomi: torch.Tensor, + atomj: torch.Tensor, + bond_type: torch.Tensor | None = None, + ) -> torch.Tensor: + """Build bond context features ``(N_bonds, bond_dim)``. + + Args: + h_atom: Atom features ``(N, atom_dim)``. + atomi: Endpoint i indices ``(N_bonds,)``. + atomj: Endpoint j indices ``(N_bonds,)``. + bond_type: Optional bond types ``(N_bonds,)``. + + Returns: + Bond features invariant under ``(i,j) → (j,i)``. + """ + return self.embed(h_atom, atomi, atomj, bond_type=bond_type) + + +class AngleContext(nn.Module): + """Angle context with reverse symmetry ``(i,j,k) ↔ (k,j,i)``. + + Mean-pools direction-specific MLPs on ``[h_i, h_j, h_k]`` and + ``[h_k, h_j, h_i]``. + """ + + def __init__( + self, + *, + atom_dim: int, + angle_dim: int = 32, + hidden_dim: int | None = None, + ) -> None: + super().__init__() + self.atom_dim = int(atom_dim) + self.angle_dim = int(angle_dim) + hid = int(hidden_dim) if hidden_dim is not None else max(3 * self.atom_dim, self.angle_dim) + self.dir_mlp = _mlp(3 * self.atom_dim, hid, self.angle_dim) + + def forward( + self, + h_atom: torch.Tensor, + atomi: torch.Tensor, + atomj: torch.Tensor, + atomk: torch.Tensor, + ) -> torch.Tensor: + """Build angle features ``(N_angles, angle_dim)``. + + Args: + h_atom: Atom features ``(N, atom_dim)``. + atomi: Endpoint i ``(N_angles,)``. + atomj: Central atom j ``(N_angles,)``. + atomk: Endpoint k ``(N_angles,)``. + + Returns: + Features invariant under ``(i,j,k) → (k,j,i)``. + """ + n = int(atomi.shape[0]) + if n == 0: + return h_atom.new_zeros((0, self.angle_dim)) + hi = h_atom[atomi.long()] + hj = h_atom[atomj.long()] + hk = h_atom[atomk.long()] + fwd = self.dir_mlp(torch.cat([hi, hj, hk], dim=-1)) + rev = self.dir_mlp(torch.cat([hk, hj, hi], dim=-1)) + return 0.5 * (fwd + rev) + + +class ProperContext(nn.Module): + """Proper-torsion context with reverse symmetry ``(i,j,k,l) ↔ (l,k,j,i)``. + + Mean-pools direction-specific MLPs on the four endpoint embeddings. + """ + + def __init__( + self, + *, + atom_dim: int, + proper_dim: int = 32, + hidden_dim: int | None = None, + ) -> None: + super().__init__() + self.atom_dim = int(atom_dim) + self.proper_dim = int(proper_dim) + hid = int(hidden_dim) if hidden_dim is not None else max(4 * self.atom_dim, self.proper_dim) + self.dir_mlp = _mlp(4 * self.atom_dim, hid, self.proper_dim) + + def forward( + self, + h_atom: torch.Tensor, + atomi: torch.Tensor, + atomj: torch.Tensor, + atomk: torch.Tensor, + atoml: torch.Tensor, + ) -> torch.Tensor: + """Build proper-torsion features ``(N_propers, proper_dim)``. + + Args: + h_atom: Atom features ``(N, atom_dim)``. + atomi: Atom i ``(N_propers,)``. + atomj: Atom j ``(N_propers,)``. + atomk: Atom k ``(N_propers,)``. + atoml: Atom l ``(N_propers,)``. + + Returns: + Features invariant under ``(i,j,k,l) → (l,k,j,i)``. + """ + n = int(atomi.shape[0]) + if n == 0: + return h_atom.new_zeros((0, self.proper_dim)) + hi = h_atom[atomi.long()] + hj = h_atom[atomj.long()] + hk = h_atom[atomk.long()] + hl = h_atom[atoml.long()] + fwd = self.dir_mlp(torch.cat([hi, hj, hk, hl], dim=-1)) + rev = self.dir_mlp(torch.cat([hl, hk, hj, hi], dim=-1)) + return 0.5 * (fwd + rev) + + +class ImproperContext(nn.Module): + """Improper context invariant under outer-leg swaps (center fixed). + + Index convention (molrs / batch schema): ``atomi`` is the **center**, + ``atomj`` / ``atomk`` / ``atoml`` are the three outer legs. Features are + built as ``MLP([h_center, sum(h_outer)])``, which is fully symmetric in + the outer legs. + """ + + def __init__( + self, + *, + atom_dim: int, + improper_dim: int = 32, + hidden_dim: int | None = None, + ) -> None: + super().__init__() + self.atom_dim = int(atom_dim) + self.improper_dim = int(improper_dim) + hid = ( + int(hidden_dim) if hidden_dim is not None else max(2 * self.atom_dim, self.improper_dim) + ) + # [center || sum(outer)] → improper_dim + self.mlp = _mlp(2 * self.atom_dim, hid, self.improper_dim) + + def forward( + self, + h_atom: torch.Tensor, + atomi: torch.Tensor, + atomj: torch.Tensor, + atomk: torch.Tensor, + atoml: torch.Tensor, + ) -> torch.Tensor: + """Build improper features ``(N_impropers, improper_dim)``. + + Args: + h_atom: Atom features ``(N, atom_dim)``. + atomi: Center atom indices ``(N_impropers,)`` (molrs center-first). + atomj: Outer leg j ``(N_impropers,)``. + atomk: Outer leg k ``(N_impropers,)``. + atoml: Outer leg l ``(N_impropers,)``. + + Returns: + Features invariant under any permutation of ``(j,k,l)`` with + center ``i`` fixed. + """ + n = int(atomi.shape[0]) + if n == 0: + return h_atom.new_zeros((0, self.improper_dim)) + h_c = h_atom[atomi.long()] + h_outer = h_atom[atomj.long()] + h_atom[atomk.long()] + h_atom[atoml.long()] + return self.mlp(torch.cat([h_c, h_outer], dim=-1)) diff --git a/src/molrep/chem/embed.py b/src/molrep/chem/embed.py new file mode 100644 index 0000000..29491ce --- /dev/null +++ b/src/molrep/chem/embed.py @@ -0,0 +1,155 @@ +"""Atom and bond continuous chemical embeddings. + +Reuses :class:`~molrep.embedding.node.JointEmbedding` for discrete atomic +numbers (and optional continuous channels). Bond embeddings enforce endpoint +symmetry ``h_ij = h_ji`` by mean-pooling direction-specific MLPs. +""" + +from __future__ import annotations + +import torch +import torch.nn as nn + +from molix import config +from molrep.embedding.node import DiscreteEmbeddingSpec, JointEmbedding + + +def _mlp(in_dim: int, hidden_dim: int, out_dim: int) -> nn.Sequential: + """Two-layer SiLU MLP baked to project float dtype.""" + return nn.Sequential( + nn.Linear(in_dim, hidden_dim, dtype=config.ftype), + nn.SiLU(), + nn.Linear(hidden_dim, out_dim, dtype=config.ftype), + ) + + +class AtomChemEmbedding(nn.Module): + """Map atomic numbers (and optional discrete channels) to atom features. + + Inputs: + ``Z`` of shape ``(N,)`` — atomic numbers. + Optional keyword tensors for extra channels configured at construct + time (future extension; v1 is Z-only via :class:`JointEmbedding`). + + Outputs: + ``h_atom`` of shape ``(N, atom_dim)``. + + Args: + atom_dim: Output feature dimension ``D_a``. + num_elements: Vocabulary size for the atomic-number table (must cover + the largest Z that will be looked up; default covers H–Og). + emb_dim: Intermediate embedding width inside :class:`JointEmbedding`. + Defaults to ``atom_dim``. + """ + + def __init__( + self, + *, + atom_dim: int = 32, + num_elements: int = 119, + emb_dim: int | None = None, + ) -> None: + super().__init__() + if atom_dim <= 0: + raise ValueError(f"atom_dim must be positive, got {atom_dim}") + if num_elements <= 0: + raise ValueError(f"num_elements must be positive, got {num_elements}") + self.atom_dim = int(atom_dim) + self.num_elements = int(num_elements) + z_emb = int(emb_dim) if emb_dim is not None else self.atom_dim + self.joint = JointEmbedding( + embedding_specs=[ + DiscreteEmbeddingSpec( + input_key="Z", + num_classes=self.num_elements, + emb_dim=z_emb, + ), + ], + out_dim=self.atom_dim, + ) + + def forward(self, Z: torch.Tensor) -> torch.Tensor: + """Embed atomic numbers. + + Args: + Z: Atomic numbers ``(N,)`` long. + + Returns: + Atom features ``(N, atom_dim)``. Empty ``N=0`` returns + ``(0, atom_dim)`` without a JointEmbedding call. + """ + if Z.numel() == 0: + return Z.new_zeros((0, self.atom_dim), dtype=config.ftype) + return self.joint(Z=Z.long()) + + +class BondChemEmbedding(nn.Module): + """Endpoint-symmetric bond embedding from atom features. + + Implements ``h_ij = ½ (MLP([h_i, h_j]) + MLP([h_j, h_i]))`` so + reversing bond endpoints leaves the feature unchanged. + + Args: + atom_dim: Input atom feature dimension ``D_a``. + bond_dim: Output bond feature dimension ``D_b``. + hidden_dim: Hidden width of the direction MLP. Defaults to + ``max(atom_dim * 2, bond_dim)``. + num_bond_types: If ``> 0``, an optional discrete bond-type table is + added (summed into both directions before the direction MLP). + ``0`` disables type conditioning. + """ + + def __init__( + self, + *, + atom_dim: int, + bond_dim: int = 32, + hidden_dim: int | None = None, + num_bond_types: int = 0, + ) -> None: + super().__init__() + if atom_dim <= 0 or bond_dim <= 0: + raise ValueError("atom_dim and bond_dim must be positive") + self.atom_dim = int(atom_dim) + self.bond_dim = int(bond_dim) + self.num_bond_types = int(num_bond_types) + hid = int(hidden_dim) if hidden_dim is not None else max(self.atom_dim * 2, self.bond_dim) + self.dir_mlp = _mlp(2 * self.atom_dim, hid, self.bond_dim) + if self.num_bond_types > 0: + self.type_emb = nn.Embedding(self.num_bond_types, self.bond_dim, dtype=config.ftype) + else: + self.type_emb = None + + def forward( + self, + h_atom: torch.Tensor, + atomi: torch.Tensor, + atomj: torch.Tensor, + bond_type: torch.Tensor | None = None, + ) -> torch.Tensor: + """Compute symmetric bond features. + + Args: + h_atom: Atom features ``(N, atom_dim)``. + atomi: Source atom indices ``(N_bonds,)``. + atomj: Target atom indices ``(N_bonds,)``. + bond_type: Optional bond type indices ``(N_bonds,)`` when + ``num_bond_types > 0``. + + Returns: + Bond features ``(N_bonds, bond_dim)`` with ``h_ij == h_ji``. + """ + n_bonds = int(atomi.shape[0]) + if n_bonds == 0: + return h_atom.new_zeros((0, self.bond_dim)) + + hi = h_atom[atomi.long()] + hj = h_atom[atomj.long()] + # Direction-specific then mean → endpoint symmetry + fwd = self.dir_mlp(torch.cat([hi, hj], dim=-1)) + rev = self.dir_mlp(torch.cat([hj, hi], dim=-1)) + h_bond = 0.5 * (fwd + rev) + + if self.type_emb is not None and bond_type is not None: + h_bond = h_bond + self.type_emb(bond_type.long()) + return h_bond diff --git a/src/molrep/chem/encoder.py b/src/molrep/chem/encoder.py new file mode 100644 index 0000000..76a49b9 --- /dev/null +++ b/src/molrep/chem/encoder.py @@ -0,0 +1,316 @@ +"""ChemEncoder — topology-aware continuous chemical perception. + +Reads valence namespaces from a nested :class:`~tensordict.TensorDict` batch +and writes continuous chem features under ``*.chem_features``. No energy, +no force, no molpot import. +""" + +from __future__ import annotations + +from typing import Any, Mapping + +import torch +import torch.nn as nn +from tensordict import TensorDict + +from molix import config +from molrep.chem.context import ( + AngleContext, + BondContext, + ImproperContext, + ProperContext, +) +from molrep.chem.embed import AtomChemEmbedding +from molrep.chem.features import ChemEmbeddings + +__all__ = ["ChemEncoder"] + + +def _ns(batch: Any, name: str) -> Any | None: + """Fetch a top-level namespace from TensorDict / Mapping.""" + if batch is None: + return None + try: + if name in batch: + return batch[name] + except Exception: # noqa: BLE001 — TensorDict key miss variants + return None + return None + + +def _get_tensor(ns: Any, *keys: str) -> torch.Tensor | None: + if ns is None: + return None + for key in keys: + try: + if isinstance(ns, (Mapping, TensorDict)) and key in ns: + return ns[key] + except Exception: # noqa: BLE001 + continue + return None + + +def _bond_endpoints(bonds: Any) -> tuple[torch.Tensor, torch.Tensor] | None: + """Return ``(atomi, atomj)`` from column keys or COO ``bond_index``.""" + if bonds is None: + return None + atomi = _get_tensor(bonds, "atomi") + atomj = _get_tensor(bonds, "atomj") + if atomi is not None and atomj is not None: + return atomi.long(), atomj.long() + packed = _get_tensor(bonds, "bond_index") + if packed is None: + return None + # Contract: bond_index is COO (2, N_bonds) + if packed.dim() != 2: + raise ValueError(f"bond_index must be 2-D, got shape {tuple(packed.shape)}") + if packed.shape[0] == 2: + return packed[0].long(), packed[1].long() + if packed.shape[1] == 2: + # Tolerate (N, 2) layout + return packed[:, 0].long(), packed[:, 1].long() + raise ValueError(f"bond_index must be (2, N) or (N, 2), got shape {tuple(packed.shape)}") + + +def _empty(dim: int, like: torch.Tensor) -> torch.Tensor: + return like.new_zeros((0, dim), dtype=config.ftype) + + +class ChemEncoder(nn.Module): + """Topology-aware chemical perception encoder. + + Reads ``atoms.Z`` and valence namespaces (``bonds`` / ``angles`` / + ``propers`` / ``impropers``) and writes continuous features under + ``*.chem_features``. + + Primary path: + ``forward(batch) -> batch`` (mutates / returns TensorDict with features). + + Secondary path: + ``compose(batch) -> ChemEmbeddings`` then + ``write_batch(batch, emb)``; ``embeddings(batch)`` views written fields. + + Args: + atom_dim: Atom feature dimension ``D_a``. + bond_dim: Bond feature dimension ``D_b``. + angle_dim: Angle feature dimension. + proper_dim: Proper-torsion feature dimension. + improper_dim: Improper feature dimension. + num_elements: Atomic-number vocabulary size. + hidden_dim: Shared hidden width for context MLPs (optional). + num_bond_types: Bond-type table size; ``0`` disables type conditioning. + """ + + in_keys = [ + ("atoms", "Z"), + ] + out_keys = [ + ("atoms", "chem_features"), + ("bonds", "chem_features"), + ("angles", "chem_features"), + ("propers", "chem_features"), + ("impropers", "chem_features"), + ] + + def __init__( + self, + *, + atom_dim: int = 32, + bond_dim: int = 32, + angle_dim: int = 32, + proper_dim: int = 32, + improper_dim: int = 32, + num_elements: int = 119, + hidden_dim: int | None = None, + num_bond_types: int = 0, + ) -> None: + super().__init__() + self.atom_dim = int(atom_dim) + self.bond_dim = int(bond_dim) + self.angle_dim = int(angle_dim) + self.proper_dim = int(proper_dim) + self.improper_dim = int(improper_dim) + self.num_elements = int(num_elements) + + self.atom_embed = AtomChemEmbedding( + atom_dim=self.atom_dim, + num_elements=self.num_elements, + ) + self.bond_context = BondContext( + atom_dim=self.atom_dim, + bond_dim=self.bond_dim, + hidden_dim=hidden_dim, + num_bond_types=num_bond_types, + ) + self.angle_context = AngleContext( + atom_dim=self.atom_dim, + angle_dim=self.angle_dim, + hidden_dim=hidden_dim, + ) + self.proper_context = ProperContext( + atom_dim=self.atom_dim, + proper_dim=self.proper_dim, + hidden_dim=hidden_dim, + ) + self.improper_context = ImproperContext( + atom_dim=self.atom_dim, + improper_dim=self.improper_dim, + hidden_dim=hidden_dim, + ) + + # ------------------------------------------------------------------ + # compose / write / view + # ------------------------------------------------------------------ + + def compose(self, batch: TensorDict | Mapping[str, Any]) -> ChemEmbeddings: + """Build :class:`ChemEmbeddings` from a valence TensorDict batch. + + Args: + batch: Nested batch with ``atoms.Z`` and optional valence + namespaces using column keys ``atomi``/``atomj``/… (or + ``bonds.bond_index`` COO). + + Returns: + :class:`ChemEmbeddings` with counts matching present topology. + Missing optional namespaces yield empty ``(0, D)`` tensors. + """ + atoms = _ns(batch, "atoms") + if atoms is None: + raise KeyError("ChemEncoder requires batch['atoms']") + z = _get_tensor(atoms, "Z") + if z is None: + raise KeyError("ChemEncoder requires batch['atoms', 'Z']") + + h_atom = self.atom_embed(z.long()) + ref = h_atom if h_atom.numel() else z + + # Bonds + bonds = _ns(batch, "bonds") + endpoints = _bond_endpoints(bonds) + if endpoints is None: + h_bond = _empty(self.bond_dim, ref) + else: + atomi, atomj = endpoints + btype = _get_tensor(bonds, "bond_types", "type", "bond_type") + h_bond = self.bond_context(h_atom, atomi, atomj, bond_type=btype) + + # Angles + angles = _ns(batch, "angles") + ai = _get_tensor(angles, "atomi") + aj = _get_tensor(angles, "atomj") + ak = _get_tensor(angles, "atomk") + if ai is None or aj is None or ak is None: + h_angle = _empty(self.angle_dim, ref) + else: + h_angle = self.angle_context(h_atom, ai, aj, ak) + + # Propers + propers = _ns(batch, "propers") + pi = _get_tensor(propers, "atomi") + pj = _get_tensor(propers, "atomj") + pk = _get_tensor(propers, "atomk") + pl = _get_tensor(propers, "atoml") + if pi is None or pj is None or pk is None or pl is None: + h_proper = _empty(self.proper_dim, ref) + else: + h_proper = self.proper_context(h_atom, pi, pj, pk, pl) + + # Impropers (atomi = center) + impropers = _ns(batch, "impropers") + ii = _get_tensor(impropers, "atomi") + ij = _get_tensor(impropers, "atomj") + ik = _get_tensor(impropers, "atomk") + il = _get_tensor(impropers, "atoml") + if ii is None or ij is None or ik is None or il is None: + h_improper = _empty(self.improper_dim, ref) + else: + h_improper = self.improper_context(h_atom, ii, ij, ik, il) + + return ChemEmbeddings( + atom=h_atom, + bond=h_bond, + angle=h_angle, + proper=h_proper, + improper=h_improper, + ) + + def write_batch( + self, + batch: TensorDict, + emb: ChemEmbeddings, + ) -> TensorDict: + """Write chem features onto the batch under ``*.chem_features``. + + Creates missing valence namespaces when the corresponding feature + tensor is empty so all five out-keys are always present after a + write. + + Args: + batch: Nested TensorDict batch (mutated in place). + emb: Feature payload from :meth:`compose`. + + Returns: + The same ``batch`` with features written. + """ + if "atoms" not in batch: + raise KeyError("write_batch requires batch['atoms']") + batch["atoms", "chem_features"] = emb.atom + + self._ensure_ns(batch, "bonds", emb.bond.shape[0]) + batch["bonds", "chem_features"] = emb.bond + + self._ensure_ns(batch, "angles", emb.angle.shape[0]) + batch["angles", "chem_features"] = emb.angle + + self._ensure_ns(batch, "propers", emb.proper.shape[0]) + batch["propers", "chem_features"] = emb.proper + + self._ensure_ns(batch, "impropers", emb.improper.shape[0]) + batch["impropers", "chem_features"] = emb.improper + return batch + + @staticmethod + def _ensure_ns(batch: TensorDict, name: str, n: int) -> None: + if name not in batch: + batch[name] = TensorDict({}, batch_size=[n] if n > 0 else []) + + def embeddings(self, batch: TensorDict | Mapping[str, Any]) -> ChemEmbeddings: + """View written ``*.chem_features`` as :class:`ChemEmbeddings`. + + Args: + batch: Batch previously passed through :meth:`forward` / + :meth:`write_batch`. + + Returns: + :class:`ChemEmbeddings` viewing the feature fields. + + Raises: + KeyError: If required chem feature keys are missing. + """ + if isinstance(batch, TensorDict): + return ChemEmbeddings( + atom=batch["atoms", "chem_features"], + bond=batch["bonds", "chem_features"], + angle=batch["angles", "chem_features"], + proper=batch["propers", "chem_features"], + improper=batch["impropers", "chem_features"], + ) + return ChemEmbeddings( + atom=batch["atoms"]["chem_features"], + bond=batch["bonds"]["chem_features"], + angle=batch["angles"]["chem_features"], + proper=batch["propers"]["chem_features"], + improper=batch["impropers"]["chem_features"], + ) + + def forward(self, batch: TensorDict) -> TensorDict: + """Compose features and write them back onto ``batch``. + + Args: + batch: Nested TensorDict with ``atoms.Z`` and valence topology. + + Returns: + The same batch with ``*.chem_features`` populated. + """ + emb = self.compose(batch) + return self.write_batch(batch, emb) diff --git a/src/molrep/chem/features.py b/src/molrep/chem/features.py new file mode 100644 index 0000000..81f667e --- /dev/null +++ b/src/molrep/chem/features.py @@ -0,0 +1,92 @@ +"""ChemEmbeddings — first-class continuous chemical-perception payload. + +Holds per-atom and per-interaction feature tensors consumed by Classical MM +heads (``molpot``) and neural parameterizers. No energy, no forces. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Mapping + +import torch + + +@dataclass +class ChemEmbeddings: + """Continuous chemical perception feature tensors. + + Attributes: + atom: Per-atom features ``(N, D_a)``. + bond: Per-bond features ``(N_bonds, D_b)``. + angle: Per-angle features ``(N_angles, D_ang)``. + proper: Per-proper-torsion features ``(N_propers, D_p)``. + improper: Per-improper features ``(N_impropers, D_imp)``. Empty + topology uses shape ``(0, D)`` rather than omitting the field. + + Notes: + Counts must align with the valence namespaces on the batch + (``bonds`` / ``angles`` / ``propers`` / ``impropers``). + """ + + atom: torch.Tensor + bond: torch.Tensor + angle: torch.Tensor + proper: torch.Tensor + improper: torch.Tensor + + def as_dict(self) -> dict[str, torch.Tensor]: + """Return a plain mapping of the five feature tensors. + + Returns: + Dict with keys ``atom``, ``bond``, ``angle``, ``proper``, + ``improper``. + """ + return { + "atom": self.atom, + "bond": self.bond, + "angle": self.angle, + "proper": self.proper, + "improper": self.improper, + } + + def interaction_dict(self) -> dict[str, torch.Tensor]: + """Feature mapping keyed for ClassicalMMComposer. + + Returns: + Dict with keys ``atoms``, ``bonds``, ``angles``, ``propers``, + ``impropers`` (plural namespace names used by MM heads). + """ + return { + "atoms": self.atom, + "bonds": self.bond, + "angles": self.angle, + "propers": self.proper, + "impropers": self.improper, + } + + @classmethod + def from_mapping(cls, data: Mapping[str, torch.Tensor]) -> ChemEmbeddings: + """Build from a fixed-key mapping. + + Args: + data: Mapping with ``atom``/``bond``/``angle``/``proper``/ + ``improper`` keys (or plural ``atoms``/``bonds``/…). + + Returns: + A :class:`ChemEmbeddings` instance. + """ + + def _get(*names: str) -> torch.Tensor: + for name in names: + if name in data: + return data[name] + raise KeyError(f"missing ChemEmbeddings field among {names}") + + return cls( + atom=_get("atom", "atoms"), + bond=_get("bond", "bonds"), + angle=_get("angle", "angles"), + proper=_get("proper", "propers"), + improper=_get("improper", "impropers"), + ) diff --git a/src/molrep/chem/typing_metrics.py b/src/molrep/chem/typing_metrics.py new file mode 100644 index 0000000..031fbbf --- /dev/null +++ b/src/molrep/chem/typing_metrics.py @@ -0,0 +1,150 @@ +"""Typing recovery metrics for GAFF atom-type evaluation probes. + +Eval-only reporting: overall accuracy, per-element accuracy, rare-type +accuracy, confusion matrix, and molecule-level error counts. + +This module is **not** part of production :class:`ClassicalMMParameterizer`. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field + +import torch + +__all__ = ["TypingRecoveryMetrics", "TypingRecoveryReport"] + + +@dataclass(frozen=True) +class TypingRecoveryReport: + """Immutable typing recovery report. + + Attributes: + overall_accuracy: Fraction of correctly predicted atom types in ``[0, 1]``. + per_element_accuracy: Map atomic number → accuracy. + rare_type_accuracy: Accuracy on types with support ≤ ``rare_max_count``. + confusion: Confusion matrix ``(n_types, n_types)`` as int64 tensor. + molecule_error_counts: Map molecule_id → number of mis-typed atoms. + n_atoms: Number of labeled atoms scored. + n_molecules: Number of molecules with at least one labeled atom. + """ + + overall_accuracy: float + per_element_accuracy: dict[int, float] + rare_type_accuracy: float | None + confusion: torch.Tensor + molecule_error_counts: dict[str, int] + n_atoms: int + n_molecules: int + + +@dataclass +class TypingRecoveryMetrics: + """Accumulate type predictions for a :class:`TypingRecoveryReport`. + + Args: + num_types: Size of the discrete type vocabulary (confusion matrix). + rare_max_count: Types with ≤ this many labels count as rare. + """ + + num_types: int + rare_max_count: int = 5 + _pred: list[torch.Tensor] = field(default_factory=list, repr=False) + _true: list[torch.Tensor] = field(default_factory=list, repr=False) + _Z: list[torch.Tensor] = field(default_factory=list, repr=False) + _mol: list[str] = field(default_factory=list, repr=False) + + def update( + self, + pred: torch.Tensor, + target: torch.Tensor, + *, + Z: torch.Tensor | None = None, + molecule_ids: list[str] | None = None, + ) -> None: + """Accumulate one batch of integer type ids. + + Args: + pred: Predicted type indices ``(N,)``. + target: Reference type indices ``(N,)``. + Z: Optional atomic numbers ``(N,)``. + molecule_ids: Optional per-atom molecule id strings length ``N``. + """ + pred = pred.detach().long().reshape(-1).cpu() + target = target.detach().long().reshape(-1).cpu() + if pred.shape != target.shape: + raise ValueError(f"pred/target shape mismatch {pred.shape} vs {target.shape}") + self._pred.append(pred) + self._true.append(target) + if Z is not None: + self._Z.append(Z.detach().long().reshape(-1).cpu()) + if molecule_ids is not None: + if len(molecule_ids) != pred.numel(): + raise ValueError("molecule_ids length must match N atoms") + self._mol.extend(list(molecule_ids)) + + def reset(self) -> None: + self._pred.clear() + self._true.clear() + self._Z.clear() + self._mol.clear() + + def compute(self) -> TypingRecoveryReport: + if not self._pred: + empty = torch.zeros(self.num_types, self.num_types, dtype=torch.int64) + return TypingRecoveryReport( + overall_accuracy=0.0, + per_element_accuracy={}, + rare_type_accuracy=None, + confusion=empty, + molecule_error_counts={}, + n_atoms=0, + n_molecules=0, + ) + pred = torch.cat(self._pred) + true = torch.cat(self._true) + n = int(pred.numel()) + correct = pred == true + overall = float(correct.float().mean().item()) if n else 0.0 + + confusion = torch.zeros(self.num_types, self.num_types, dtype=torch.int64) + for t, p in zip(true.tolist(), pred.tolist(), strict=True): + if 0 <= t < self.num_types and 0 <= p < self.num_types: + confusion[t, p] += 1 + + per_el: dict[int, float] = {} + if self._Z: + Z = torch.cat(self._Z) + for z in sorted(set(Z.tolist())): + mask = Z == z + if mask.any(): + per_el[int(z)] = float(correct[mask].float().mean().item()) + + # Rare types by true-label support + rare_acc: float | None = None + type_counts = torch.bincount(true, minlength=self.num_types) + rare_mask = torch.zeros(n, dtype=torch.bool) + for t in range(self.num_types): + if 0 < int(type_counts[t]) <= self.rare_max_count: + rare_mask |= true == t + if rare_mask.any(): + rare_acc = float(correct[rare_mask].float().mean().item()) + + mol_err: dict[str, int] = {} + if self._mol: + for mid, ok in zip(self._mol, correct.tolist(), strict=True): + if not ok: + mol_err[mid] = mol_err.get(mid, 0) + 1 + n_mol = len(set(self._mol)) + else: + n_mol = 0 + + return TypingRecoveryReport( + overall_accuracy=overall, + per_element_accuracy=per_el, + rare_type_accuracy=rare_acc, + confusion=confusion, + molecule_error_counts=mol_err, + n_atoms=n, + n_molecules=n_mol, + ) diff --git a/src/molrep/chem/typing_probe.py b/src/molrep/chem/typing_probe.py new file mode 100644 index 0000000..e4ce82a --- /dev/null +++ b/src/molrep/chem/typing_probe.py @@ -0,0 +1,78 @@ +"""Eval-only GAFF atom-type readout on continuous ChemEncoder embeddings. + +Hard isolation rules +-------------------- +* Temporary :class:`~molrep.heads.TypeHead` sits **on top of** continuous + embeddings — atom-type labels never enter the encoder as inputs. +* Production :class:`~molpot.composition.parameterizer.ClassicalMMParameterizer` + must not import this module. + +References: + Wang et al., GAFF, J. Comput. Chem. 2004. + Espaloma Chem. Sci. 2022 DOI 10.1039/D2SC02739A. +""" + +from __future__ import annotations + +from typing import Any + +import torch +import torch.nn as nn +from tensordict import TensorDict + +from molrep.chem.encoder import ChemEncoder +from molrep.heads.type import TypeHead + +__all__ = ["AtomTypeReadout"] + + +class AtomTypeReadout(nn.Module): + """Compose :class:`ChemEncoder` + :class:`TypeHead` for typing recovery. + + Args: + encoder: Continuous chemical perception encoder. + num_types: Discrete type vocabulary size. + dropout: TypeHead dropout. + """ + + def __init__( + self, + encoder: ChemEncoder, + *, + num_types: int, + dropout: float = 0.0, + ) -> None: + super().__init__() + self.encoder = encoder + self.head = TypeHead( + hidden_dim=int(encoder.atom_dim), + num_types=num_types, + dropout=dropout, + ) + self.num_types = num_types + + @property + def atom_dim(self) -> int: + return int(self.head.hidden_dim) + + def encode(self, batch: TensorDict) -> torch.Tensor: + """Run encoder and return atom chem features ``(N, D)``. + + Labels under ``atoms.atom_type`` / ``atom_type_id`` are **ignored**. + """ + out = self.encoder(batch) + feats = out["atoms", "chem_features"] + return feats + + def forward(self, batch: TensorDict) -> dict[str, Any]: + """Return logits and argmax type ids from continuous embeddings. + + Args: + batch: Nested TensorDict with atoms.Z and topology. + + Returns: + Dict with ``logits`` ``(N, num_types)`` and ``pred_type_id`` ``(N,)``. + """ + feats = self.encode(batch) + logits = self.head(feats) + return {"logits": logits, "pred_type_id": self.head.decode(logits), "features": feats} diff --git a/src/molrep/condensation/__init__.py b/src/molrep/condensation/__init__.py new file mode 100644 index 0000000..86fcb63 --- /dev/null +++ b/src/molrep/condensation/__init__.py @@ -0,0 +1,52 @@ +"""Physics-aware multi-system chemical class condensation. + +Turn continuous per-interaction MM parameters into a discrete +:class:`TypeSystem` per :class:`InteractionClass` via greedy merge under +:class:`MergeCriterion` budgets. Optional physics gating uses an injected +``physical_eval`` callable — this package never builds molpot energy graphs +and never emits SMARTS/SMIRKS text (sub-spec 07). + +Public surface: + InteractionClass, MergeCriterion, default_criterion, + TypeRecord, TypeSystem, UNMATCHED_TYPE_ID, + ClassAssignment, PhysicalErrorMetrics, + Condenser, CondensationResult, + TypeSystemLabeler +""" + +from molrep.condensation.assignment import ClassAssignment +from molrep.condensation.classes import InteractionClass +from molrep.condensation.condenser import CondensationResult, Condenser +from molrep.condensation.criterion import ( + MergeCriterion, + angle_default_criterion, + bond_default_criterion, + charge_default_criterion, + default_criterion, + improper_default_criterion, + lj_default_criterion, + proper_default_criterion, +) +from molrep.condensation.labeler import TypeSystemLabeler +from molrep.condensation.metrics import PhysicalErrorMetrics +from molrep.condensation.type_system import UNMATCHED_TYPE_ID, TypeRecord, TypeSystem + +__all__ = [ + "InteractionClass", + "MergeCriterion", + "default_criterion", + "bond_default_criterion", + "angle_default_criterion", + "proper_default_criterion", + "improper_default_criterion", + "lj_default_criterion", + "charge_default_criterion", + "TypeRecord", + "TypeSystem", + "UNMATCHED_TYPE_ID", + "ClassAssignment", + "PhysicalErrorMetrics", + "Condenser", + "CondensationResult", + "TypeSystemLabeler", +] diff --git a/src/molrep/condensation/assignment.py b/src/molrep/condensation/assignment.py new file mode 100644 index 0000000..10d7b09 --- /dev/null +++ b/src/molrep/condensation/assignment.py @@ -0,0 +1,80 @@ +"""Class assignment tables produced by condensation. + +Maps every condensed interaction row to a global type id and its source +system. Integer ids only — no SMARTS text. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch + +from molrep.condensation.classes import InteractionClass + +__all__ = ["ClassAssignment"] + + +@dataclass +class ClassAssignment: + """Global type-id assignment for one :class:`InteractionClass`. + + Attributes: + interaction: Interaction family these rows belong to. + type_ids: Global type ids ``(N,)`` long, values in + ``0 .. n_types-1`` (or ``-1`` for unmatched soft-fail). + system_ids: Source system index per row ``(N,)`` long. + row_indices: Original row index within each system ``(N,)`` long. + """ + + interaction: InteractionClass + type_ids: torch.Tensor + system_ids: torch.Tensor + row_indices: torch.Tensor + + def __post_init__(self) -> None: + n = int(self.type_ids.numel()) + if self.system_ids.numel() != n or self.row_indices.numel() != n: + raise ValueError( + "ClassAssignment tensors must share length: " + f"type_ids={n}, system_ids={self.system_ids.numel()}, " + f"row_indices={self.row_indices.numel()}" + ) + self.type_ids = self.type_ids.long().reshape(-1) + self.system_ids = self.system_ids.long().reshape(-1) + self.row_indices = self.row_indices.long().reshape(-1) + + @property + def n_rows(self) -> int: + """Total number of assigned interaction rows.""" + return int(self.type_ids.numel()) + + def for_system(self, system_id: int) -> torch.Tensor: + """Type ids for rows belonging to ``system_id``. + + Args: + system_id: System index used during multi-system merge. + + Returns: + Long tensor of type ids in original row order for that system. + """ + mask = self.system_ids == int(system_id) + ids = self.type_ids[mask] + order = self.row_indices[mask] + if ids.numel() == 0: + return ids + # Restore original within-system order. + sorted_idx = torch.argsort(order) + return ids[sorted_idx] + + def type_id_at(self, system_id: int, row_index: int) -> int: + """Return the type id for one ``(system, row)`` pair. + + Raises: + KeyError: If the pair is not present. + """ + mask = (self.system_ids == int(system_id)) & (self.row_indices == int(row_index)) + hits = self.type_ids[mask] + if hits.numel() == 0: + raise KeyError(f"No assignment for system={system_id}, row={row_index}") + return int(hits[0].item()) diff --git a/src/molrep/condensation/classes.py b/src/molrep/condensation/classes.py new file mode 100644 index 0000000..b81c6d4 --- /dev/null +++ b/src/molrep/condensation/classes.py @@ -0,0 +1,31 @@ +"""InteractionClass — discrete MM interaction kinds for condensation / SMARTS.""" + +from __future__ import annotations + +from enum import Enum + +__all__ = ["InteractionClass"] + + +class InteractionClass(str, Enum): + """Chemical interaction classes that own a discrete type table. + + Values align with valence TensorDict namespaces where applicable + (``bonds`` / ``angles`` / ``propers`` / ``impropers``) plus nonbonded + classes for LJ and partial charges. + + Attributes: + BOND: Pair bonded stretch terms. + ANGLE: Three-body angle terms. + PROPER: Four-body proper torsions. + IMPROPER: Four-body impropers (center-first topology). + LJ: Lennard-Jones (or equivalent) nonbonded types. + CHARGE: Partial-charge classes (often not condensed aggressively). + """ + + BOND = "bond" + ANGLE = "angle" + PROPER = "proper" + IMPROPER = "improper" + LJ = "lj" + CHARGE = "charge" diff --git a/src/molrep/condensation/condenser.py b/src/molrep/condensation/condenser.py new file mode 100644 index 0000000..633ffcd --- /dev/null +++ b/src/molrep/condensation/condenser.py @@ -0,0 +1,282 @@ +"""Physics-aware greedy merge of continuous MM parameters into discrete types. + +:class:`Condenser` never builds energy graphs. Optional physical gating uses an +injected ``physical_eval`` callable (torch kernel, molpy evaluator, or test +double) whose residuals feed :class:`PhysicalErrorMetrics`. +""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass +from typing import Any + +import torch + +from molrep.condensation.assignment import ClassAssignment +from molrep.condensation.classes import InteractionClass +from molrep.condensation.criterion import MergeCriterion, default_criterion +from molrep.condensation.metrics import PhysicalErrorMetrics, parse_physical_eval_result +from molrep.condensation.type_system import TypeSystem + +__all__ = ["Condenser", "CondensationResult", "PhysicalEval"] + +# physical_eval(prototype, candidate, interaction) -> residual +PhysicalEval = Callable[ + [Mapping[str, torch.Tensor], Mapping[str, torch.Tensor], InteractionClass], + Any, +] + + +@dataclass(frozen=True) +class CondensationResult: + """Outputs of one condensation run. + + Attributes: + type_system: Discrete type table for the interaction class. + assignment: Per-row global type ids with system provenance. + metrics: Physical residual aggregates (empty if no ``physical_eval``). + """ + + type_system: TypeSystem + assignment: ClassAssignment + metrics: PhysicalErrorMetrics + + +def _infer_n_rows(params: Mapping[str, torch.Tensor | float] | torch.Tensor) -> int: + if isinstance(params, torch.Tensor): + if params.ndim == 0: + return 1 + return int(params.shape[0]) + if not params: + return 0 + first = next(iter(params.values())) + if isinstance(first, torch.Tensor): + if first.ndim == 0: + return 1 + return int(first.shape[0]) + return 1 + + +def _row_at( + params: Mapping[str, torch.Tensor | float] | torch.Tensor, + index: int, + *, + param_keys: Sequence[str] | None = None, +) -> dict[str, torch.Tensor]: + """Extract one interaction row as a float-tensor mapping.""" + if isinstance(params, torch.Tensor): + # Bare tensor: treat as a single unnamed feature vector → "value" + # or a stacked table (N, D) with optional key names. + t = params.detach().to(dtype=torch.float64) + if t.ndim == 0: + return {"value": t.clone()} + if t.ndim == 1: + # Ambiguous: (D,) one row vs (N,) scalar-per-row. + # Prefer scalar-per-row when param_keys is None and we index. + if param_keys is not None and len(param_keys) == t.numel(): + return { + k: t[i].detach().to(dtype=torch.float64).clone() + for i, k in enumerate(param_keys) + } + return {"value": t[index].detach().to(dtype=torch.float64).clone()} + # (N, D) or (N, ...) + row = t[index] + if param_keys is not None and row.ndim == 1 and len(param_keys) == row.numel(): + return { + k: row[i].detach().to(dtype=torch.float64).clone() for i, k in enumerate(param_keys) + } + return {"value": row.clone()} + + out: dict[str, torch.Tensor] = {} + for key, val in params.items(): + if isinstance(val, torch.Tensor): + t = val.detach().to(dtype=torch.float64) + if t.ndim == 0: + out[key] = t.clone() + else: + out[key] = t[index].clone() + else: + out[key] = torch.as_tensor(val, dtype=torch.float64) + return out + + +def _flatten_systems( + params_by_system: Sequence[Mapping[str, torch.Tensor | float] | torch.Tensor], +) -> list[tuple[int, int, dict[str, torch.Tensor]]]: + """Expand multi-system params to ``(system_id, row_index, row_params)``. + + Sort key (documented, deterministic): ``(system_id ascending, + row_index ascending)`` — insertion order within each system, systems in + the order provided. + """ + rows: list[tuple[int, int, dict[str, torch.Tensor]]] = [] + for sys_id, params in enumerate(params_by_system): + n = _infer_n_rows(params) + for row_i in range(n): + rows.append((sys_id, row_i, _row_at(params, row_i))) + # Explicit sort for determinism even if caller order changes later. + rows.sort(key=lambda item: (item[0], item[1])) + return rows + + +class Condenser: + """Physics-aware greedy merge across one or many chemical systems. + + Algorithm + --------- + 1. Flatten ``params_by_system`` with ``system_id`` tracking. + 2. Sort by documented key ``(system_id, row_index)``. + 3. For each row: assign to the first existing type whose prototype + passes :class:`MergeCriterion.accepts` **and** the optional physics + gate; otherwise spawn a new type. + 4. Update the type centroid as an online mean of its members. + + Multi-system merges produce **global** type ids (not renumbered per + molecule). No SMARTS text is emitted. + + Args: + update_centroid: If True (default), absorb updates the prototype + mean; if False, the first member freezes the prototype. + """ + + def __init__(self, *, update_centroid: bool = True) -> None: + self.update_centroid = update_centroid + + def merge( + self, + params_by_system: Sequence[Mapping[str, torch.Tensor | float] | torch.Tensor], + *, + interaction: InteractionClass, + criterion: MergeCriterion | None = None, + physical_eval: PhysicalEval | None = None, + energy_tol: float | None = None, + force_tol: float | None = None, + ) -> CondensationResult: + """Greedy merge continuous parameters into a discrete type system. + + Args: + params_by_system: Sequence of per-system parameter bags. Each bag + is a mapping of Class-I tensors with leading batch dim + ``(N_sys, ...)`` (or a bare tensor of rows). + interaction: Interaction family for this merge. + criterion: Merge budgets. Defaults to + :func:`~molrep.condensation.criterion.default_criterion`. + physical_eval: Optional injected residual evaluator + ``(prototype, candidate, interaction) -> residual``. + Condensation does **not** import molpot energy graphs; the + caller supplies any physical oracle. + energy_tol: If set with ``physical_eval``, reject merges whose + absolute energy residual exceeds this value. + force_tol: If set with ``physical_eval``, reject merges whose + absolute force residual exceeds this value. + + Returns: + :class:`CondensationResult` with type system, assignment table, + and physical metrics. + """ + crit = criterion if criterion is not None else default_criterion(interaction) + if crit.interaction != interaction: + raise ValueError( + f"criterion.interaction={crit.interaction} does not match interaction={interaction}" + ) + + type_system = TypeSystem(interaction, criterion=crit) + metrics = PhysicalErrorMetrics() + rows = _flatten_systems(params_by_system) + + type_ids: list[int] = [] + system_ids: list[int] = [] + row_indices: list[int] = [] + + for flat_id, (sys_id, row_i, row_params) in enumerate(rows): + assigned = self._try_assign( + type_system, + row_params, + criterion=crit, + physical_eval=physical_eval, + energy_tol=energy_tol, + force_tol=force_tol, + metrics=metrics, + support_id=flat_id, + ) + if assigned is None: + assigned = type_system._spawn(row_params, support_id=flat_id) + type_ids.append(assigned) + system_ids.append(sys_id) + row_indices.append(row_i) + + assignment = ClassAssignment( + interaction=interaction, + type_ids=torch.tensor(type_ids, dtype=torch.long), + system_ids=torch.tensor(system_ids, dtype=torch.long), + row_indices=torch.tensor(row_indices, dtype=torch.long), + ) + return CondensationResult( + type_system=type_system, + assignment=assignment, + metrics=metrics, + ) + + def greedy_merge( + self, + params_by_system: Sequence[Mapping[str, torch.Tensor | float] | torch.Tensor], + *, + interaction: InteractionClass, + criterion: MergeCriterion | None = None, + physical_eval: PhysicalEval | None = None, + energy_tol: float | None = None, + force_tol: float | None = None, + ) -> CondensationResult: + """Alias of :meth:`merge` (spec name ``Condenser.greedy_merge``).""" + return self.merge( + params_by_system, + interaction=interaction, + criterion=criterion, + physical_eval=physical_eval, + energy_tol=energy_tol, + force_tol=force_tol, + ) + + def _try_assign( + self, + type_system: TypeSystem, + row_params: Mapping[str, torch.Tensor], + *, + criterion: MergeCriterion, + physical_eval: PhysicalEval | None, + energy_tol: float | None, + force_tol: float | None, + metrics: PhysicalErrorMetrics, + support_id: int, + ) -> int | None: + """Return type_id if an existing type accepts, else None.""" + for rec in type_system.records(): + proto_t = rec.prototype_tensors() + if not criterion.accepts(proto_t, row_params): + continue + if physical_eval is not None and (energy_tol is not None or force_tol is not None): + raw = physical_eval(proto_t, row_params, type_system.interaction) + e_err, f_err = parse_physical_eval_result(raw) + rejected = False + if energy_tol is not None and e_err is not None and e_err > energy_tol: + rejected = True + if force_tol is not None and f_err is not None and f_err > force_tol: + rejected = True + metrics.record(energy_error=e_err, force_error=f_err, rejected=rejected) + if rejected: + continue + elif physical_eval is not None: + # Evaluate for metrics only (no hard gate without tolerances). + raw = physical_eval(proto_t, row_params, type_system.interaction) + e_err, f_err = parse_physical_eval_result(raw) + metrics.record(energy_error=e_err, force_error=f_err, rejected=False) + + type_system._absorb( + rec.type_id, + row_params, + support_id=support_id, + update_centroid=self.update_centroid, + ) + return rec.type_id + return None diff --git a/src/molrep/condensation/criterion.py b/src/molrep/condensation/criterion.py new file mode 100644 index 0000000..9233d6a --- /dev/null +++ b/src/molrep/condensation/criterion.py @@ -0,0 +1,285 @@ +"""Merge budgets for physics-aware chemical type condensation. + +Units follow CLASS_I_CANONICAL (kcal/mol, Å, e, rad) as used by continuous +MM heads. Budgets are configuration objects — tune per library; defaults +prioritise force-field readability while bounding energetic drift. + +References: + Spec: learnable-classical-ff-06-condensation + OpenMM User Guide §19 "Forces" +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass, field +from typing import Mapping + +import torch + +from molrep.condensation.classes import InteractionClass + +__all__ = [ + "MergeCriterion", + "bond_default_criterion", + "angle_default_criterion", + "proper_default_criterion", + "improper_default_criterion", + "lj_default_criterion", + "charge_default_criterion", + "default_criterion", +] + +# --------------------------------------------------------------------------- +# Default budgets (documented physical rationale) +# --------------------------------------------------------------------------- +# Bond: Δr0 ≤ 0.01 Å keeps geometry tables readable; relative k ≤ 5% (with +# absolute floor) limits force-constant drift for stiff bonds. +_BOND_ABS_R0 = 0.01 # Å +_BOND_REL_K = 0.05 # dimensionless +_BOND_ABS_K_FLOOR = 1.0 # kcal/mol/Ų + +# Angle: Δθ0 ≤ 1° in rad; same relative-k spirit as bonds. +_ANGLE_ABS_THETA0 = math.radians(1.0) # rad +_ANGLE_REL_K = 0.05 +_ANGLE_ABS_K_FLOOR = 1.0 # kcal/mol/rad² + +# Proper / improper: per-term barrier and phase tolerances. +_TORSION_ABS_K = 0.1 # kcal/mol +_TORSION_REL_K = 0.05 +_TORSION_ABS_PHASE = math.radians(15.0) # rad +_TORSION_ABS_CHI0 = math.radians(5.0) # rad (harmonic improper) + +# LJ: Δσ ≤ 0.01 Å; relative ε ≤ 5%. +_LJ_ABS_SIGMA = 0.01 # Å +_LJ_REL_EPSILON = 0.05 + +# Charge: absolute Δq (e); rarely merged aggressively. +_CHARGE_ABS_Q = 0.02 # e + + +def _as_tensor( + value: torch.Tensor | float | tuple[float, ...] | list[float], + *, + like: torch.Tensor | None = None, +) -> torch.Tensor: + if isinstance(value, torch.Tensor): + return value + if isinstance(value, (tuple, list)): + value = list(value) + if like is not None: + return torch.as_tensor(value, dtype=like.dtype, device=like.device) + return torch.as_tensor(value, dtype=torch.float64) + + +def _flatten_row( + params: Mapping[str, torch.Tensor | float | tuple[float, ...] | list[float]], +) -> dict[str, torch.Tensor]: + """Normalize a single interaction's params to tensors (no batch dim).""" + out: dict[str, torch.Tensor] = {} + for key, val in params.items(): + t = _as_tensor(val) + # Squeeze a leading singleton batch dim only: (1,) → scalar, (1, T) → (T,). + if t.ndim >= 1 and t.shape[0] == 1: + t = t.reshape(*t.shape[1:]) if t.ndim > 1 else t.reshape(()) + out[key] = t + return out + + +@dataclass(frozen=True) +class MergeCriterion: + """Per-class merge budgets for continuous → discrete type condensation. + + All absolute length thresholds are in **Å**; angles in **radians**; + energy-related constants in **kcal/mol** (and force-constant dimensions + matching Class-I IR). Relative thresholds are dimensionless fractions. + + Args: + interaction: Owning :class:`InteractionClass`. + abs_tol: Absolute tolerances keyed by parameter name + (e.g. ``{"r0": 0.01}`` for bonds). + rel_tol: Relative tolerances keyed by parameter name + (e.g. ``{"k": 0.05}``). Compared as + ``|a - b| ≤ rel * max(|proto|, |cand|, abs_floor)``. + abs_floor: Denominator floors for relative checks, same units as the + parameter (prevents tiny-k blow-ups). + required_keys: Parameter names that must be present on both sides. + + Notes: + :meth:`accepts` is pure parameter-space. Physics gating (energy / + force residuals) is applied by :class:`Condenser` via an injected + ``physical_eval`` callable — this type never builds energy graphs. + """ + + interaction: InteractionClass + abs_tol: Mapping[str, float] = field(default_factory=dict) + rel_tol: Mapping[str, float] = field(default_factory=dict) + abs_floor: Mapping[str, float] = field(default_factory=dict) + required_keys: tuple[str, ...] = () + + def accepts( + self, + prototype_params: Mapping[str, torch.Tensor | float | tuple[float, ...]], + candidate_params: Mapping[str, torch.Tensor | float | tuple[float, ...]], + ) -> bool: + """Return True if candidate params fall inside budgets of prototype. + + Args: + prototype_params: Type centroid / prototype parameter mapping. + candidate_params: Candidate interaction parameters. + + Returns: + ``True`` when every budgeted key is within absolute and/or + relative tolerance; ``False`` otherwise. + + Raises: + KeyError: If a required key is missing on either side. + """ + proto = _flatten_row(prototype_params) + cand = _flatten_row(candidate_params) + + for key in self.required_keys: + if key not in proto or key not in cand: + raise KeyError( + f"MergeCriterion({self.interaction.value}) requires key " + f"{key!r} on both prototype and candidate" + ) + + budgeted = set(self.abs_tol) | set(self.rel_tol) + if not budgeted: + keys = set(proto) & set(cand) + else: + # Only compare keys present on both sides (improper may emit + # harmonic-only or periodic-only parameter families). + keys = budgeted & set(proto) & set(cand) + if not keys and self.required_keys: + # Required keys already validated; nothing budgeted to compare. + return True + if not keys: + return False + + for key in keys: + a = proto[key] + b = cand[key] + if a.shape != b.shape: + return False + diff = torch.abs(a - b) + + abs_ok = True + if key in self.abs_tol: + abs_ok = bool(torch.all(diff <= self.abs_tol[key]).item()) + + rel_ok = True + if key in self.rel_tol: + floor = float(self.abs_floor.get(key, 0.0)) + scale = torch.maximum(torch.abs(a), torch.abs(b)) + scale = torch.clamp(scale, min=floor) + # When both near zero and floor is 0, require exact match via abs. + if floor == 0.0 and bool(torch.all(scale == 0).item()): + rel_ok = bool(torch.all(diff == 0).item()) + else: + rel_ok = bool(torch.all(diff <= self.rel_tol[key] * scale).item()) + + # Key passes if *any* configured budget for that key holds when both + # are set; if only one family is set, that family decides. + if key in self.abs_tol and key in self.rel_tol: + if not (abs_ok or rel_ok): + return False + elif key in self.abs_tol: + if not abs_ok: + return False + elif key in self.rel_tol: + if not rel_ok: + return False + + return True + + +def bond_default_criterion() -> MergeCriterion: + """Default bond budgets: Δr0 ≤ 0.01 Å; Δk/k ≤ 5% with abs floor 1.0.""" + return MergeCriterion( + interaction=InteractionClass.BOND, + abs_tol={"r0": _BOND_ABS_R0}, + rel_tol={"k": _BOND_REL_K}, + abs_floor={"k": _BOND_ABS_K_FLOOR}, + required_keys=("k", "r0"), + ) + + +def angle_default_criterion() -> MergeCriterion: + """Default angle budgets: Δθ0 ≤ 1° (rad); relative k ≤ 5%.""" + return MergeCriterion( + interaction=InteractionClass.ANGLE, + abs_tol={"theta0": _ANGLE_ABS_THETA0}, + rel_tol={"k": _ANGLE_REL_K}, + abs_floor={"k": _ANGLE_ABS_K_FLOOR}, + required_keys=("k", "theta0"), + ) + + +def proper_default_criterion() -> MergeCriterion: + """Default proper-torsion budgets: per-term |Δk| / rel k and |Δphase|.""" + return MergeCriterion( + interaction=InteractionClass.PROPER, + abs_tol={"k": _TORSION_ABS_K, "phase": _TORSION_ABS_PHASE}, + rel_tol={"k": _TORSION_REL_K}, + abs_floor={"k": _TORSION_ABS_K}, + required_keys=("k", "phase"), + ) + + +def improper_default_criterion() -> MergeCriterion: + """Default improper budgets (harmonic χ0 and/or periodic k/phase).""" + return MergeCriterion( + interaction=InteractionClass.IMPROPER, + abs_tol={ + "k": _TORSION_ABS_K, + "phase": _TORSION_ABS_PHASE, + "chi0": _TORSION_ABS_CHI0, + "k_harmonic": _TORSION_ABS_K, + "k_periodic": _TORSION_ABS_K, + }, + rel_tol={"k": _TORSION_REL_K, "k_harmonic": _TORSION_REL_K, "k_periodic": _TORSION_REL_K}, + abs_floor={"k": _TORSION_ABS_K, "k_harmonic": _TORSION_ABS_K, "k_periodic": _TORSION_ABS_K}, + required_keys=(), + ) + + +def lj_default_criterion() -> MergeCriterion: + """Default LJ budgets: Δσ ≤ 0.01 Å; relative ε ≤ 5%.""" + return MergeCriterion( + interaction=InteractionClass.LJ, + abs_tol={"sigma": _LJ_ABS_SIGMA}, + rel_tol={"epsilon": _LJ_REL_EPSILON}, + abs_floor={"epsilon": 1e-4}, + required_keys=("epsilon", "sigma"), + ) + + +def charge_default_criterion() -> MergeCriterion: + """Default charge budget: |Δq| ≤ 0.02 e (conservative; often unused).""" + return MergeCriterion( + interaction=InteractionClass.CHARGE, + abs_tol={"q": _CHARGE_ABS_Q, "charge": _CHARGE_ABS_Q}, + required_keys=(), + ) + + +def default_criterion(interaction: InteractionClass) -> MergeCriterion: + """Return the documented default :class:`MergeCriterion` for ``interaction``. + + Args: + interaction: Target interaction family. + + Returns: + A frozen criterion with Class-I unit budgets. + """ + table = { + InteractionClass.BOND: bond_default_criterion, + InteractionClass.ANGLE: angle_default_criterion, + InteractionClass.PROPER: proper_default_criterion, + InteractionClass.IMPROPER: improper_default_criterion, + InteractionClass.LJ: lj_default_criterion, + InteractionClass.CHARGE: charge_default_criterion, + } + return table[interaction]() diff --git a/src/molrep/condensation/labeler.py b/src/molrep/condensation/labeler.py new file mode 100644 index 0000000..49fce5e --- /dev/null +++ b/src/molrep/condensation/labeler.py @@ -0,0 +1,90 @@ +"""TypeSystemLabeler — Labeler Protocol for condensed discrete types.""" + +from __future__ import annotations + +from typing import Mapping + +import torch + +from molrep.condensation.criterion import MergeCriterion +from molrep.condensation.type_system import UNMATCHED_TYPE_ID, TypeSystem + +__all__ = ["TypeSystemLabeler"] + + +class TypeSystemLabeler: + """Map continuous parameter rows to condensed type ids. + + Implements the :class:`~molrep.heads.labeler.Labeler` surface + (``num_types``, ``type_map``, ``label(...)``) so callers can treat + condensed ids like atom-type labels. ``label`` accepts a parameter + mapping (preferred) or a tensor of precomputed type ids; it does + **not** emit SMARTS text. + + Args: + type_system: Frozen (or post-merge) discrete type table. + criterion: Optional criterion override for nearest-prototype assign. + type_map: Optional human-readable names; defaults to + ``{i: "_"}``. + """ + + def __init__( + self, + type_system: TypeSystem, + *, + criterion: MergeCriterion | None = None, + type_map: dict[int, str] | None = None, + ) -> None: + self._type_system = type_system + self._criterion = criterion + if type_map is None: + prefix = type_system.interaction.value + type_map = {i: f"{prefix}_{i}" for i in range(type_system.n_types)} + self._type_map = dict(type_map) + + @property + def type_system(self) -> TypeSystem: + """Underlying condensed type table.""" + return self._type_system + + @property + def num_types(self) -> int: + """Number of discrete types (Labeler Protocol).""" + return self._type_system.n_types + + @property + def type_map(self) -> dict[int, str]: + """Mapping from type id to a display name (Labeler Protocol).""" + return self._type_map + + def label( + self, + params: Mapping[str, torch.Tensor] | torch.Tensor, + ) -> torch.Tensor: + """Assign condensed type ids for each parameter row. + + Args: + params: Either + + * a mapping of Class-I parameter tensors with batch dim + ``(N, ...)``, assigned via :meth:`TypeSystem.assign_many`, or + * a long tensor of already-known type ids (pass-through + clamp / validate path). + + Returns: + Long tensor ``(N,)`` with ids in ``0 .. num_types-1``. Unmatched + rows become ``0`` only when ``num_types == 0`` is impossible; + otherwise unmatched rows stay :data:`UNMATCHED_TYPE_ID` (``-1``). + """ + if isinstance(params, torch.Tensor): + ids = params.long().reshape(-1) + # Validate range for known ids; leave -1 soft-fails intact. + if self.num_types > 0: + bad = (ids >= self.num_types) | ((ids < 0) & (ids != UNMATCHED_TYPE_ID)) + if bool(bad.any()): + raise ValueError( + f"type ids out of range for num_types={self.num_types}: {ids[bad].tolist()}" + ) + return ids + + return self._type_system.assign_many(params, criterion=self._criterion) diff --git a/src/molrep/condensation/metrics.py b/src/molrep/condensation/metrics.py new file mode 100644 index 0000000..f79c8a1 --- /dev/null +++ b/src/molrep/condensation/metrics.py @@ -0,0 +1,115 @@ +"""Physical residual metrics for physics-gated condensation. + +:class:`PhysicalErrorMetrics` records energy/force residuals reported by an +**injected** ``physical_eval`` callable. Condensation never builds molpot +energy graphs itself. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, Mapping + +__all__ = ["PhysicalErrorMetrics", "parse_physical_eval_result"] + + +@dataclass +class PhysicalErrorMetrics: + """Aggregated residuals from optional physical evaluation during merge. + + Attributes: + n_compared: Number of prototype–candidate pairs evaluated. + rejected_by_physics: Pairs that passed param budgets but failed the + physical residual gate. + max_abs_energy_error: Max |ΔE| observed (caller units, typically + kcal/mol), or ``None`` if never set. + mean_abs_energy_error: Mean |ΔE| over compared pairs. + max_abs_force_error: Max force residual norm if provided. + mean_abs_force_error: Mean force residual if provided. + energy_errors: Per-comparison energy residuals (append-only log). + force_errors: Per-comparison force residuals (append-only log). + """ + + n_compared: int = 0 + rejected_by_physics: int = 0 + max_abs_energy_error: float | None = None + mean_abs_energy_error: float | None = None + max_abs_force_error: float | None = None + mean_abs_force_error: float | None = None + energy_errors: list[float] = field(default_factory=list) + force_errors: list[float] = field(default_factory=list) + + def record( + self, + *, + energy_error: float | None = None, + force_error: float | None = None, + rejected: bool = False, + ) -> None: + """Append one comparison result and refresh aggregates. + + Args: + energy_error: Absolute energy residual for this pair. + force_error: Absolute force residual for this pair. + rejected: Whether the physics gate rejected the merge. + """ + self.n_compared += 1 + if rejected: + self.rejected_by_physics += 1 + if energy_error is not None: + e = abs(float(energy_error)) + self.energy_errors.append(e) + self.max_abs_energy_error = ( + e if self.max_abs_energy_error is None else max(self.max_abs_energy_error, e) + ) + self.mean_abs_energy_error = sum(self.energy_errors) / len(self.energy_errors) + if force_error is not None: + f = abs(float(force_error)) + self.force_errors.append(f) + self.max_abs_force_error = ( + f if self.max_abs_force_error is None else max(self.max_abs_force_error, f) + ) + self.mean_abs_force_error = sum(self.force_errors) / len(self.force_errors) + + +def parse_physical_eval_result( + result: Any, +) -> tuple[float | None, float | None]: + """Normalize a ``physical_eval`` return value to ``(energy_err, force_err)``. + + Accepted shapes: + + * ``float`` / ``int`` → energy error only + * ``Mapping`` with keys ``energy`` / ``energy_error`` / ``delta_e`` and + optional ``force`` / ``force_error`` / ``delta_f`` + * ``(energy, force)`` 2-tuple + + Args: + result: Raw return from the injected callable. + + Returns: + Pair of optional absolute residuals. + """ + if result is None: + return None, None + if isinstance(result, (int, float)): + return abs(float(result)), None + if isinstance(result, Mapping): + e = None + for key in ("energy_error", "energy", "delta_e", "abs_energy_error"): + if key in result: + e = abs(float(result[key])) + break + f = None + for key in ("force_error", "force", "delta_f", "abs_force_error"): + if key in result: + f = abs(float(result[key])) + break + return e, f + if isinstance(result, (tuple, list)) and len(result) >= 1: + e = abs(float(result[0])) if result[0] is not None else None + f = abs(float(result[1])) if len(result) > 1 and result[1] is not None else None + return e, f + raise TypeError( + f"physical_eval must return float, mapping, or (energy, force) tuple; got {type(result)!r}" + ) diff --git a/src/molrep/condensation/type_system.py b/src/molrep/condensation/type_system.py new file mode 100644 index 0000000..30b22d1 --- /dev/null +++ b/src/molrep/condensation/type_system.py @@ -0,0 +1,351 @@ +"""TypeSystem / TypeRecord — discrete condensed parameter prototypes. + +Integer type ids + numeric prototypes only. No SMARTS / SMIRKS emission +(owned by sub-spec 07). Mutable helpers ``_spawn`` / ``_absorb`` are used by +:class:`~molrep.condensation.condenser.Condenser` during greedy merge; the +public table is otherwise treated as frozen after merge. +""" + +from __future__ import annotations + +from collections.abc import Iterable, Iterator, Mapping, Sequence +from dataclasses import dataclass +from typing import Any + +import torch + +from molrep.condensation.classes import InteractionClass +from molrep.condensation.criterion import MergeCriterion, default_criterion + +__all__ = ["TypeRecord", "TypeSystem", "UNMATCHED_TYPE_ID", "Prototype"] + +# Soft-fail sentinel for :meth:`TypeSystem.assign` when no prototype accepts. +UNMATCHED_TYPE_ID: int = -1 + +Prototype = Mapping[str, float | tuple[float, ...] | torch.Tensor] + + +def _to_tensor(value: float | tuple[float, ...] | torch.Tensor) -> torch.Tensor: + if isinstance(value, torch.Tensor): + return value.detach().to(dtype=torch.float64).clone() + if isinstance(value, (tuple, list)): + return torch.as_tensor(list(value), dtype=torch.float64) + return torch.as_tensor(value, dtype=torch.float64) + + +def _to_python(value: torch.Tensor) -> float | tuple[float, ...]: + t = value.detach().cpu().reshape(-1) + if t.numel() == 1: + return float(t.item()) + return tuple(float(x) for x in t.tolist()) + + +def _prototype_tensors( + prototype: Mapping[str, float | tuple[float, ...] | torch.Tensor], +) -> dict[str, torch.Tensor]: + return {k: _to_tensor(v) for k, v in prototype.items()} + + +def _prototype_python( + prototype: Mapping[str, float | tuple[float, ...] | torch.Tensor], +) -> dict[str, float | tuple[float, ...]]: + out: dict[str, float | tuple[float, ...]] = {} + for k, v in prototype.items(): + if isinstance(v, torch.Tensor): + out[k] = _to_python(v) + elif isinstance(v, (tuple, list)): + out[k] = tuple(float(x) for x in v) + else: + out[k] = float(v) + return out + + +def _running_mean_tensors( + prototype: Mapping[str, torch.Tensor], + candidate: Mapping[str, torch.Tensor], + member_count: int, +) -> dict[str, torch.Tensor]: + """Online mean: mean_new = mean_old + (x - mean_old) / (n+1).""" + updated: dict[str, torch.Tensor] = {} + for key, pval in prototype.items(): + if key not in candidate or pval.shape != candidate[key].shape: + updated[key] = pval.clone() + continue + cval = candidate[key].to(dtype=pval.dtype) + updated[key] = pval + (cval - pval) / float(member_count + 1) + for key, cval in candidate.items(): + if key not in updated: + updated[key] = cval.clone() + return updated + + +@dataclass(frozen=True) +class TypeRecord: + """One discrete chemical type and its parameter prototype. + + Attributes: + type_id: Non-negative integer id within its :class:`TypeSystem`. + prototype: Named parameter values in Class-I units + (e.g. ``{"k": 300.0, "r0": 1.09}``). Tensor values are accepted + at construction and stored as Python floats / tuples. + member_count: How many interactions were merged into this type. + label: Optional human-readable name. + support_ids: Optional member indices into the condensation input stream. + """ + + type_id: int + prototype: Mapping[str, float | tuple[float, ...]] + member_count: int = 0 + label: str | None = None + support_ids: tuple[int, ...] = () + + def __post_init__(self) -> None: + if self.type_id < 0: + raise ValueError(f"type_id must be non-negative, got {self.type_id}") + if self.member_count < 0: + raise ValueError(f"member_count must be non-negative, got {self.member_count}") + object.__setattr__(self, "prototype", _prototype_python(self.prototype)) + object.__setattr__(self, "support_ids", tuple(self.support_ids)) + + def prototype_tensors(self) -> dict[str, torch.Tensor]: + """Prototype as float64 tensors (for criterion / assign).""" + return _prototype_tensors(self.prototype) + + +class TypeSystem: + """Ordered discrete type table for one :class:`InteractionClass`. + + Built by :class:`~molrep.condensation.condenser.Condenser`. At inference, + :meth:`assign` maps new continuous parameters to an existing type when a + :class:`MergeCriterion` accepts the prototype; otherwise returns + :data:`UNMATCHED_TYPE_ID` (soft reject — does **not** spawn types). + + Args: + interaction: Owning interaction family. + records: Ordered type records (ids should be unique). + criterion: Optional criterion used by :meth:`assign`. Defaults to + the documented class default when omitted. + """ + + def __init__( + self, + interaction: InteractionClass, + records: Sequence[TypeRecord] | None = None, + *, + criterion: MergeCriterion | None = None, + ) -> None: + self._interaction = interaction + recs = list(records or ()) + seen: set[int] = set() + for rec in recs: + if rec.type_id in seen: + raise ValueError( + f"duplicate type_id {rec.type_id} in TypeSystem for {interaction.value}" + ) + seen.add(rec.type_id) + self._records: list[TypeRecord] = recs + self._by_id: dict[int, TypeRecord] = {r.type_id: r for r in recs} + self._criterion = criterion if criterion is not None else default_criterion(interaction) + + @property + def interaction(self) -> InteractionClass: + """Interaction class this type table covers.""" + return self._interaction + + @property + def n_types(self) -> int: + """Number of discrete types.""" + return len(self._records) + + @property + def criterion(self) -> MergeCriterion: + """Merge criterion used for :meth:`assign`.""" + return self._criterion + + def records(self) -> list[TypeRecord]: + """Return ordered type records (copy).""" + return list(self._records) + + def get(self, type_id: int) -> TypeRecord: + """Look up a type by id. + + Args: + type_id: Discrete type identifier. + + Returns: + The matching :class:`TypeRecord`. + + Raises: + KeyError: If ``type_id`` is unknown. + """ + try: + return self._by_id[type_id] + except KeyError as exc: + raise KeyError( + f"type_id {type_id} not in TypeSystem for {self._interaction.value}" + ) from exc + + def prototypes_table(self) -> list[dict[str, Any]]: + """Return prototypes as a list of plain dicts ordered by table order.""" + return [dict(r.prototype) for r in self._records] + + def prototypes_tensor_table(self) -> dict[str, torch.Tensor]: + """Stack shared prototype keys into tensors ``(n_types, ...)``. + + Keys missing on any record are omitted. + """ + if not self._records: + return {} + keys = set(self._records[0].prototype) + for rec in self._records[1:]: + keys &= set(rec.prototype) + table: dict[str, torch.Tensor] = {} + for key in sorted(keys): + table[key] = torch.stack( + [_to_tensor(rec.prototype[key]) for rec in self._records], + dim=0, + ) + return table + + def assign( + self, + params: Mapping[str, torch.Tensor | float | tuple[float, ...]], + *, + criterion: MergeCriterion | None = None, + ) -> int: + """Map params to the first acceptable existing type id. + + Soft-fail contract: if no prototype accepts ``params``, return + :data:`UNMATCHED_TYPE_ID` (``-1``). Never creates a new type. + + Args: + params: Continuous parameters for one interaction row. + criterion: Optional override; defaults to the system criterion. + + Returns: + Existing ``type_id`` or :data:`UNMATCHED_TYPE_ID`. + """ + crit = criterion if criterion is not None else self._criterion + row = _prototype_tensors(params) + for rec in self._records: + if crit.accepts(rec.prototype_tensors(), row): + return rec.type_id + return UNMATCHED_TYPE_ID + + def assign_many( + self, + params: Mapping[str, torch.Tensor], + *, + criterion: MergeCriterion | None = None, + ) -> torch.Tensor: + """Vectorized :meth:`assign` over a batch of parameter rows. + + Args: + params: Mapping of tensors with leading batch dim ``(N, ...)``. + criterion: Optional criterion override. + + Returns: + Long tensor ``(N,)`` of type ids (``-1`` for unmatched). + """ + if not params: + return torch.zeros(0, dtype=torch.long) + first = next(iter(params.values())) + n = int(first.shape[0]) + out = torch.full((n,), UNMATCHED_TYPE_ID, dtype=torch.long) + for i in range(n): + row = {k: v[i] for k, v in params.items()} + out[i] = self.assign(row, criterion=criterion) + return out + + def __iter__(self) -> Iterator[TypeRecord]: + return iter(self._records) + + def __len__(self) -> int: + return len(self._records) + + def __repr__(self) -> str: + return f"TypeSystem(interaction={self._interaction!r}, n_types={self.n_types})" + + @classmethod + def from_prototypes( + cls, + interaction: InteractionClass, + prototypes: Iterable[Prototype], + *, + labels: Sequence[str | None] | None = None, + criterion: MergeCriterion | None = None, + ) -> TypeSystem: + """Build a type system with sequential type ids from prototypes. + + Args: + interaction: Interaction class. + prototypes: One prototype mapping per type (order defines type_id). + labels: Optional labels parallel to ``prototypes``. + criterion: Optional assign criterion. + + Returns: + A new :class:`TypeSystem`. + """ + protos = list(prototypes) + labs = list(labels) if labels is not None else [None] * len(protos) + if len(labs) != len(protos): + raise ValueError("labels length must match prototypes length") + records = [ + TypeRecord(type_id=i, prototype=p, member_count=0, label=lab) + for i, (p, lab) in enumerate(zip(protos, labs, strict=True)) + ] + return cls(interaction, records, criterion=criterion) + + # --- mutation used by Condenser (package-internal) --------------------- + + def _spawn( + self, + params: Mapping[str, torch.Tensor | float | tuple[float, ...]], + *, + support_id: int | None = None, + label: str | None = None, + ) -> int: + """Append a new type from ``params``; return its type_id.""" + type_id = self.n_types + support = (support_id,) if support_id is not None else () + rec = TypeRecord( + type_id=type_id, + prototype=_prototype_python(params), + member_count=1, + label=label, + support_ids=support, + ) + self._records.append(rec) + self._by_id[type_id] = rec + return type_id + + def _absorb( + self, + type_id: int, + params: Mapping[str, torch.Tensor | float | tuple[float, ...]], + *, + support_id: int | None = None, + update_centroid: bool = True, + ) -> None: + """Merge ``params`` into an existing type, updating the centroid.""" + rec = self._by_id[type_id] + proto_t = rec.prototype_tensors() + cand_t = _prototype_tensors(params) + if update_centroid: + new_proto = _running_mean_tensors(proto_t, cand_t, rec.member_count) + else: + new_proto = proto_t + support = rec.support_ids + ((support_id,) if support_id is not None else ()) + updated = TypeRecord( + type_id=rec.type_id, + prototype=_prototype_python(new_proto), + member_count=rec.member_count + 1, + label=rec.label, + support_ids=support, + ) + # Replace in ordered list + for i, existing in enumerate(self._records): + if existing.type_id == type_id: + self._records[i] = updated + break + self._by_id[type_id] = updated diff --git a/src/molrep/embedding/__init__.py b/src/molrep/embedding/__init__.py index deca007..ec2238b 100644 --- a/src/molrep/embedding/__init__.py +++ b/src/molrep/embedding/__init__.py @@ -5,15 +5,29 @@ - SphericalHarmonics: Equivariant angular basis functions - BesselRBF / GaussianBasis / PolynomialBasis: Radial basis functions - CosineCutoff / TanhCutoff / HalfCosineCutoff / PolynomialCutoff: Cutoff envelopes +- ChemicalSupportIndex: L2 kNN chemical-support bank (provenance surfaces) """ from .angular import SphericalHarmonics +from .covalent import covalent_radii from .cutoff import CosineCutoff, HalfCosineCutoff, PolynomialCutoff, TanhCutoff +from .mlp import MomentNormalizedMLP, normalize2mom from .node import JointEmbedding, JointFeatureEmbedding, JointFeatureSpec -from .radial import BesselRBF, GaussianBasis, PolynomialBasis +from .radial import ( + AgnesiTransform, + BesselRBF, + GaussianBasis, + PolynomialBasis, +) +from .support import ChemicalSupportIndex __all__ = [ + "AgnesiTransform", "BesselRBF", + "ChemicalSupportIndex", + "covalent_radii", + "MomentNormalizedMLP", + "normalize2mom", "CosineCutoff", "GaussianBasis", "HalfCosineCutoff", diff --git a/src/molrep/embedding/covalent.py b/src/molrep/embedding/covalent.py new file mode 100644 index 0000000..ae663ee --- /dev/null +++ b/src/molrep/embedding/covalent.py @@ -0,0 +1,159 @@ +"""Z-indexed covalent-radius lookup table. + +A single free function because the table has no natural owning type: it is +published numerical data materialised as a dense ``Z``-indexed tensor so radial +transforms and pair-repulsion terms can index it with raw atomic numbers. Two +consumers today — :class:`molrep.embedding.radial.AgnesiTransform` and +:class:`molpot.potentials.repulsion.ZBLRepulsion`. + +The values are inlined rather than read from ``molpy.Element`` deliberately. +This is a constant of nature, not a molpy domain type, and it sits on the import +path of a core ``molrep`` block: sourcing it from molpy would make every +equivariant encoder depend on the compiled molrs extension for 119 floats — and +molpy stores radii in single precision, which is a fidelity loss against the +float64 reference table the MACE foundation checkpoints were fitted with. + +Reference: + Cordero et al. "Covalent radii revisited" Dalton Trans. 2008, 2832-2838. + https://doi.org/10.1039/B801115J +""" + +from __future__ import annotations + +import torch + +from molix import config + +#: Radius used where the source tabulates none — index 0 (not an element) and +#: everything past curium. Matches the reference table byte for byte, which is +#: what a MACE checkpoint's own ``covalent_radii`` buffer carries. +_MISSING_RADIUS = 0.2 + +# Cordero (2008) covalent radii in Angstrom for Z = 1..96, the full span the +# paper tabulates. Heavier elements fall back to _MISSING_RADIUS. +# fmt: off — 10-per-row layout mirrors the paper's table; one-per-line kills readability +_CORDERO_RADII: tuple[float, ...] = ( + 0.31, + 0.28, + 1.28, + 0.96, + 0.84, + 0.76, + 0.71, + 0.66, + 0.57, + 0.58, + 1.66, + 1.41, + 1.21, + 1.11, + 1.07, + 1.05, + 1.02, + 1.06, + 2.03, + 1.76, + 1.70, + 1.60, + 1.53, + 1.39, + 1.39, + 1.32, + 1.26, + 1.24, + 1.32, + 1.22, + 1.22, + 1.20, + 1.19, + 1.20, + 1.20, + 1.16, + 2.20, + 1.95, + 1.90, + 1.75, + 1.64, + 1.54, + 1.47, + 1.46, + 1.42, + 1.39, + 1.45, + 1.44, + 1.42, + 1.39, + 1.39, + 1.38, + 1.39, + 1.40, + 2.44, + 2.15, + 2.07, + 2.04, + 2.03, + 2.01, + 1.99, + 1.98, + 1.98, + 1.96, + 1.94, + 1.92, + 1.92, + 1.89, + 1.90, + 1.87, + 1.87, + 1.75, + 1.70, + 1.62, + 1.51, + 1.44, + 1.41, + 1.36, + 1.36, + 1.32, + 1.45, + 1.46, + 1.48, + 1.40, + 1.50, + 1.50, + 2.60, + 2.21, + 2.15, + 2.06, + 2.00, + 1.96, + 1.90, + 1.87, + 1.80, + 1.69, +) +# fmt: on + +#: Largest atomic number a table may be built for. +MAX_Z = 118 + +_COVALENT_RADII: tuple[float, ...] = ( + (_MISSING_RADIUS,) + _CORDERO_RADII + (_MISSING_RADIUS,) * (MAX_Z - len(_CORDERO_RADII)) +) + + +def covalent_radii(max_z: int = MAX_Z) -> torch.Tensor: + """Build a dense ``Z``-indexed covalent-radius table. + + Args: + max_z: Largest atomic number to include; the table has ``max_z + 1`` + rows so ``table[Z]`` is valid for ``Z`` in ``[0, max_z]``. + + Returns: + Covalent radii in Angstrom ``(max_z + 1,)``; index 0 holds a + dummy-atom placeholder. + + Raises: + ValueError: If ``max_z`` is below 1 or beyond the tabulated range. + """ + if not 1 <= max_z <= MAX_Z: + raise ValueError(f"max_z must be in [1, {MAX_Z}], got {max_z}") + return torch.tensor(_COVALENT_RADII[: max_z + 1], dtype=config.ftype) diff --git a/src/molrep/embedding/cutoff.py b/src/molrep/embedding/cutoff.py index 17ed89c..ec25f59 100644 --- a/src/molrep/embedding/cutoff.py +++ b/src/molrep/embedding/cutoff.py @@ -53,7 +53,7 @@ def __init__(self, *, r_cut: float): ) # Register buffers with type annotations - r_cut_tensor = torch.tensor(float(self.config.r_cut)) + r_cut_tensor = torch.tensor(float(self.config.r_cut), dtype=config.ftype) self.register_buffer("r_cut", r_cut_tensor, persistent=False) self.r_cut: torch.Tensor @@ -131,25 +131,31 @@ def __init__(self, *, r_cut: float, exponent: int = 6): exponent=exponent, ) - r_cut_tensor = torch.tensor(float(self.config.r_cut)) + r_cut_tensor = torch.tensor(float(self.config.r_cut), dtype=config.ftype) self.register_buffer("r_cut", r_cut_tensor, persistent=False) self.r_cut: torch.Tensor self.exponent = int(self.config.exponent) - def forward(self, r: torch.Tensor) -> torch.Tensor: - """Apply polynomial cutoff to distances. + @staticmethod + def envelope(r: torch.Tensor, r_cut: torch.Tensor | float, exponent: int) -> torch.Tensor: + """Evaluate the envelope for a per-element (or scalar) cutoff radius. + + Split out from :meth:`forward` because pair-repulsion terms (ZBL) use a + **per-edge** cutoff derived from the two atoms' covalent radii, whereas + the module itself carries one fixed ``r_cut``. Args: - r: Input distances. + r: Distances, any shape. + r_cut: Cutoff radius; scalar or broadcastable to ``r``. + exponent: Polynomial exponent ``p``. Returns: - Cutoff values. Values range from 1.0 (at r=0) to 0.0 (at r>=r_cut). + Envelope values, same shape as ``r``: 1.0 at ``r = 0``, 0.0 for + ``r >= r_cut``. """ - x = r / self.r_cut - mask = x < 1.0 - - p = float(self.exponent) + x = r / r_cut + p = float(exponent) c_p = (p + 1.0) * (p + 2.0) / 2.0 c_p1 = p * (p + 2.0) c_p2 = p * (p + 1.0) / 2.0 @@ -159,7 +165,18 @@ def forward(self, r: torch.Tensor) -> torch.Tensor: x_p2 = x_p1 * x c = 1.0 - c_p * x_p + c_p1 * x_p1 - c_p2 * x_p2 - return torch.where(mask, c, torch.zeros_like(c)) + return torch.where(x < 1.0, c, torch.zeros_like(c)) + + def forward(self, r: torch.Tensor) -> torch.Tensor: + """Apply polynomial cutoff to distances. + + Args: + r: Input distances. + + Returns: + Cutoff values. Values range from 1.0 (at r=0) to 0.0 (at r>=r_cut). + """ + return self.envelope(r, self.r_cut, self.exponent) class TanhCutoffSpec(BaseModel): diff --git a/src/molrep/embedding/mace.py b/src/molrep/embedding/mace.py new file mode 100644 index 0000000..dcf2448 --- /dev/null +++ b/src/molrep/embedding/mace.py @@ -0,0 +1,155 @@ +"""MACE-only embedding block. + +:class:`EmbeddingBlock` composes the embedding-layer primitives the MACE +encoders need — ``JointEmbedding`` for node attributes, ``BesselRBF`` for the +radial basis, ``SphericalHarmonics`` for edge directions and ``CosineCutoff`` +for the envelope — into the single ``(node_feats, edge_attrs, edge_feats)`` +producer that feeds the interaction stack. + +Relocated verbatim from ``molzoo/mace.py``, which continues to re-export both +names for backwards compatibility. ``molrep.embedding.__init__`` deliberately +does *not* re-export them: ``EmbeddingBlock`` / ``EmbeddingSpec`` are too +generic a name at package level, so consumers import from this module path. + +Reference: + Batatia et al. "MACE: Higher Order Equivariant Message Passing Neural + Networks for Fast and Accurate Force Fields" NeurIPS 2022. + https://arxiv.org/abs/2206.07697 +""" + +from __future__ import annotations + +import torch +import torch.nn as nn +from pydantic import BaseModel, ConfigDict, Field + +from molrep.embedding.angular import SphericalHarmonics +from molrep.embedding.cutoff import CosineCutoff +from molrep.embedding.node import ( + ContinuousEmbeddingSpec, + DiscreteEmbeddingSpec, + JointEmbedding, +) +from molrep.embedding.radial import BesselRBF + + +class EmbeddingSpec(BaseModel): + """Configuration for the embedding block. + + Attributes: + node_attr_specs: Embedding specifications for node attributes + (e.g. atomic number Z, charge). + num_features: Number of feature channels (scalar multiplicity at l=0). + r_max: Radial cutoff distance in Angstroms. + num_bessel: Number of Bessel radial basis functions. + l_max: Maximum angular momentum order. + """ + + model_config = ConfigDict(arbitrary_types_allowed=True) + + node_attr_specs: list[DiscreteEmbeddingSpec | ContinuousEmbeddingSpec] = Field( + ..., min_length=1 + ) + num_features: int = Field(..., gt=0) + r_max: float = Field(..., gt=0.0) + num_bessel: int = Field(8, gt=0) + l_max: int = Field(2, ge=0) + + +class EmbeddingBlock(nn.Module): + """Node and edge embedding block. + + Computes initial node features via ``JointEmbedding`` and edge features + via Bessel radial basis, spherical harmonics, and a cosine cutoff envelope. + + Attributes: + node_embedding: Joint embedding for node attributes. + radial_embedding: Bessel radial basis functions. + spherical_harmonics: Spherical harmonics for edge directions. + cutoff_fn: Cosine cutoff envelope. + """ + + def __init__( + self, + *, + node_attr_specs: list[DiscreteEmbeddingSpec | ContinuousEmbeddingSpec], + num_features: int, + r_max: float, + num_bessel: int = 8, + l_max: int = 2, + ): + """Initialize embedding block. + + Args: + node_attr_specs: Embedding specs for node attributes. + num_features: Scalar channel multiplicity (l=0 count). + r_max: Radial cutoff in Angstroms. + num_bessel: Number of Bessel basis functions. + l_max: Maximum angular momentum order. + """ + super().__init__() + + self.config = EmbeddingSpec( + node_attr_specs=node_attr_specs, + num_features=num_features, + r_max=r_max, + num_bessel=num_bessel, + l_max=l_max, + ) + + # Node embedding + self.node_embedding = JointEmbedding( + embedding_specs=node_attr_specs, + out_dim=num_features, + ) + + # Edge radial basis + self.radial_embedding = BesselRBF( + r_cut=r_max, + num_radial=num_bessel, + ) + + # Spherical harmonics + self.spherical_harmonics = SphericalHarmonics( + l_max=l_max, + ) + + # Cutoff envelope + self.cutoff_fn = CosineCutoff( + r_cut=r_max, + ) + + def forward( + self, + Z: torch.Tensor, + edge_dist: torch.Tensor, + edge_diff: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Compute initial node and edge features. + + Args: + Z: Atomic numbers (n_nodes,). + edge_dist: Bond distances (n_edges,). + edge_diff: Bond vectors (target - source) (n_edges, 3). + + Returns: + tuple of: + - node_feats: Node features (n_nodes, num_features). + - edge_attrs: Spherical harmonics (n_edges, sh_dim). + - edge_feats: Radial basis features (n_edges, num_bessel). + """ + # Node features + node_feats = self.node_embedding(Z=Z) + + # Edge direction + edge_dir = edge_diff / (edge_dist.unsqueeze(-1) + 1e-8) + + # Spherical harmonics + edge_attrs = self.spherical_harmonics(edge_dir) + + # Radial basis * cutoff → edge_feats + edge_radial = self.radial_embedding(edge_dist) + edge_cutoff = self.cutoff_fn(edge_dist) + edge_feats = edge_radial * edge_cutoff.unsqueeze(-1) + + return node_feats, edge_attrs, edge_feats diff --git a/src/molrep/embedding/mlp.py b/src/molrep/embedding/mlp.py index 1f714f0..0e268e5 100644 --- a/src/molrep/embedding/mlp.py +++ b/src/molrep/embedding/mlp.py @@ -16,10 +16,12 @@ from __future__ import annotations import math +from collections.abc import Callable from typing import Optional import torch import torch.nn as nn +import torch.nn.functional as F from molix import config @@ -130,3 +132,78 @@ def __init__( def forward(self, x: torch.Tensor) -> torch.Tensor: return self.mlp(x) + + +def normalize2mom(act: Callable[[torch.Tensor], torch.Tensor]) -> float: + """Return the constant that normalises ``act`` to unit second moment. + + ``c`` such that ``E[(c·act(z))²] = 1`` for ``z ~ N(0, 1)``, estimated on a + fixed one-million-sample draw (seed 0) so the value is reproducible across + processes and matches e3nn's ``e3nn.math.normalize2mom``. + + Args: + act: Elementwise activation. + + Returns: + The scaling constant. + """ + gen = torch.Generator(device="cpu").manual_seed(0) + z = torch.randn(1_000_000, generator=gen, dtype=torch.float64) + return act(z).pow(2).mean().pow(-0.5).item() + + +class MomentNormalizedMLP(nn.Module): + """Scalar MLP with ``1/√fan_in`` weight scaling and moment-normalised SiLU. + + Port of e3nn's ``e3nn.nn.FullyConnectedNet`` (with ``act=silu``), which is + what MACE uses for its radial weight generator and edge-density head. Each + layer computes ``x @ (W / √fan_in)``; every layer but the last then applies + ``c·SiLU`` with ``c`` from :func:`normalize2mom`. There are no biases. + + Sub-modules are named ``layer0 … layerN`` so official MACE weights transfer + by direct copy. + + Args: + channels: Widths ``[in, hidden…, out]`` (at least two entries). + + Reference: + Batatia et al. "MACE: Higher Order Equivariant Message Passing Neural + Networks for Fast and Accurate Force Fields" NeurIPS 2022. + https://arxiv.org/abs/2206.07697 + """ + + class _Layer(nn.Module): + """One ``x @ (W / √fan_in)`` layer; ``e3nn.nn._fc._Layer`` weight layout.""" + + def __init__(self, in_features: int, out_features: int) -> None: + super().__init__() + self.weight = nn.Parameter(torch.randn(in_features, out_features, dtype=config.ftype)) + self.alpha = 1.0 / math.sqrt(in_features) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x @ (self.weight * self.alpha) + + def __init__(self, channels: list[int]) -> None: + super().__init__() + if len(channels) < 2: + raise ValueError(f"channels must have at least [in, out], got {channels}") + self.channels = list(channels) + self._act_cst = normalize2mom(F.silu) + for i, (h_in, h_out) in enumerate(zip(channels, channels[1:])): + self.add_module(f"layer{i}", self._Layer(h_in, h_out)) + self._num_layers = len(channels) - 1 + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Apply the MLP. + + Args: + x: Input ``(..., channels[0])``. + + Returns: + ``(..., channels[-1])``. + """ + for i in range(self._num_layers): + x = getattr(self, f"layer{i}")(x) + if i < self._num_layers - 1: + x = F.silu(x) * self._act_cst + return x diff --git a/src/molrep/embedding/node.py b/src/molrep/embedding/node.py index 551bfc3..50ae94e 100644 --- a/src/molrep/embedding/node.py +++ b/src/molrep/embedding/node.py @@ -10,6 +10,8 @@ import torch.nn as nn from pydantic import BaseModel, Field, model_validator +from molix import config + class DiscreteEmbeddingSpec(BaseModel): """Specification for a single discrete feature embedding. @@ -136,15 +138,17 @@ def __init__( output_key="_unused", # Placeholder since we are removing keys ) + ftype = config.ftype + # Store embedders in a list to match spec order self.embedders = nn.ModuleList() for spec in self.config.specs: if isinstance(spec, DiscreteEmbeddingSpec): - self.embedders.append(nn.Embedding(spec.num_classes, spec.emb_dim)) + self.embedders.append(nn.Embedding(spec.num_classes, spec.emb_dim, dtype=ftype)) else: self.embedders.append( nn.Sequential( - nn.Linear(spec.in_dim, spec.emb_dim, bias=spec.use_bias), + nn.Linear(spec.in_dim, spec.emb_dim, bias=spec.use_bias, dtype=ftype), ) ) @@ -157,6 +161,7 @@ def __init__( cue.Irreps("O3", f"{total_dim}x0e"), cue.Irreps("O3", f"{out_dim}x0e"), layout=cue.ir_mul, + dtype=ftype, ), ) @@ -256,24 +261,26 @@ def __init__( self.specs = feature_specs self.out_dim = int(out_dim) + ftype = config.ftype + self.embedders = nn.ModuleDict() for s in feature_specs: if s.kind == "categorical": if s.num_classes is None: raise ValueError(f"categorical feature {s.name!r} needs num_classes.") - self.embedders[s.name] = nn.Embedding(s.num_classes, s.emb_dim) + self.embedders[s.name] = nn.Embedding(s.num_classes, s.emb_dim, dtype=ftype) else: if s.in_dim is None: raise ValueError(f"continuous feature {s.name!r} needs in_dim.") self.embedders[s.name] = nn.Sequential( - nn.Linear(s.in_dim, s.emb_dim, bias=s.use_bias), + nn.Linear(s.in_dim, s.emb_dim, bias=s.use_bias, dtype=ftype), nn.SiLU(), - nn.Linear(s.emb_dim, s.emb_dim, bias=s.use_bias), + nn.Linear(s.emb_dim, s.emb_dim, bias=s.use_bias, dtype=ftype), ) total_dim = sum(s.emb_dim for s in feature_specs) self.project = nn.Sequential( - nn.Linear(total_dim, self.out_dim, bias=False), + nn.Linear(total_dim, self.out_dim, bias=False, dtype=ftype), nn.SiLU(), ) diff --git a/src/molrep/embedding/radial.py b/src/molrep/embedding/radial.py index 7a2c5f0..9c9a8e1 100644 --- a/src/molrep/embedding/radial.py +++ b/src/molrep/embedding/radial.py @@ -8,6 +8,8 @@ from molix import config +from .covalent import covalent_radii + class BesselRBFSpec(BaseModel): """Specification for Bessel radial basis function. @@ -275,3 +277,82 @@ def forward(self, r: torch.Tensor, *, fc: torch.Tensor | None = None) -> torch.T if fc is not None: basis = basis * fc.unsqueeze(-1) return basis + + +class AgnesiTransform(nn.Module): + """Element-pair distance transform applied before the radial basis. + + Compresses the radial coordinate onto a bounded interval before the Bessel + expansion, using a pair-specific length scale ``r_0 = (R_i + R_j) / 2`` built + from the two atoms' covalent radii: + + .. math:: + + u(r) = \\left[1 + \\frac{a\\,(r/r_0)^q}{1 + (r/r_0)^{q-p}}\\right]^{-1} + + This concentrates basis resolution where the pair distribution is dense, + which is what makes a single basis transfer across the periodic table. + + Args: + q: Numerator exponent. + p: Denominator exponent (``p > q`` gives the long-range decay). + a: Overall amplitude. + max_z: Largest atomic number in the covalent-radius table. + trainable: If ``True``, ``a`` / ``q`` / ``p`` become learnable. + + Reference: + Witt et al. "ACEpotentials.jl: A Julia implementation of the atomic + cluster expansion" J. Chem. Phys. 159, 164101 (2023), § Radial + transformations. https://doi.org/10.1063/5.0158783 + """ + + def __init__( + self, + *, + q: float = 0.9183, + p: float = 4.5791, + a: float = 1.0805, + max_z: int = 118, + trainable: bool = False, + ) -> None: + super().__init__() + ftype = config.ftype + values = { + "q": torch.tensor(float(q), dtype=ftype), + "p": torch.tensor(float(p), dtype=ftype), + "a": torch.tensor(float(a), dtype=ftype), + } + for name, value in values.items(): + if trainable: + setattr(self, name, nn.Parameter(value)) + else: + self.register_buffer(name, value) + self.q: torch.Tensor + self.p: torch.Tensor + self.a: torch.Tensor + + self.register_buffer("covalent_radii", covalent_radii(max_z)) + self.covalent_radii: torch.Tensor + + def forward( + self, + r: torch.Tensor, + z_source: torch.Tensor, + z_target: torch.Tensor, + ) -> torch.Tensor: + """Transform edge distances. + + Args: + r: Edge distances ``(E,)``. + z_source: Atomic number of each edge's source atom ``(E,)``. + z_target: Atomic number of each edge's target atom ``(E,)``. + + Returns: + Transformed distances ``(E,)``. + """ + radii = self.covalent_radii + r_0 = 0.5 * (radii[z_source.long()] + radii[z_target.long()]) + x = r / r_0 + return torch.reciprocal( + 1.0 + self.a * torch.pow(x, self.q) / (1.0 + torch.pow(x, self.q - self.p)) + ) diff --git a/src/molrep/embedding/support.py b/src/molrep/embedding/support.py new file mode 100644 index 0000000..15611d4 --- /dev/null +++ b/src/molrep/embedding/support.py @@ -0,0 +1,207 @@ +"""Chemical support indexing via an L2 k-nearest-neighbour bank. + +Thin perception-side bookkeeping for learnable classical force fields: +store training embeddings (or discrete type-id points) and query whether +new vectors fall inside a support radius. No molpot imports; no active- +learning loop. + +Reference: + Spec: learnable-classical-ff-09-provenance +""" + +from __future__ import annotations + +from collections.abc import Iterable + +import torch + +__all__ = ["ChemicalSupportIndex"] + + +class ChemicalSupportIndex: + """L2 kNN bank of chemical-support reference vectors. + + Built from continuous training embeddings / fingerprints, or from a + discrete type-id set via :meth:`from_type_ids` (1-D points). Membership + is ``min L2 distance <= radius``. + + Args: + bank: Support vectors ``(n_support, dim)``. + radius: Inclusive L2 radius for :meth:`contains`. + k: Number of nearest neighbours returned by :meth:`knn`. + + Raises: + ValueError: If ``bank`` is not 2-D, empty, or ``k`` / ``radius`` + are invalid. + """ + + def __init__( + self, + bank: torch.Tensor, + *, + radius: float, + k: int = 1, + ) -> None: + if bank.ndim != 2: + raise ValueError( + f"ChemicalSupportIndex bank must be 2-D (n, dim); got shape {tuple(bank.shape)}" + ) + if bank.shape[0] == 0: + raise ValueError("ChemicalSupportIndex bank must contain at least one vector") + if k < 1: + raise ValueError(f"k must be >= 1; got {k}") + if radius < 0: + raise ValueError(f"radius must be >= 0; got {radius}") + # Detach + clone so the bank is a pure lookup table (no grad graph). + self._bank = bank.detach().clone() + self._radius = float(radius) + self._k = min(int(k), int(self._bank.shape[0])) + + @classmethod + def from_type_ids( + cls, + type_ids: Iterable[int], + *, + radius: float = 0.0, + k: int = 1, + ) -> ChemicalSupportIndex: + """Build a 1-D L2 bank from discrete type identifiers. + + Each type id becomes a point ``[float(id)]``. With ``radius=0`` this + is exact set membership under L2. + + Args: + type_ids: Iterable of integer type ids in the known support. + radius: Inclusive L2 radius (default exact match). + k: Neighbours for :meth:`knn`. + + Returns: + Index over unique type-id points. + """ + unique = sorted({int(t) for t in type_ids}) + if not unique: + raise ValueError("from_type_ids requires at least one type id") + bank = torch.tensor(unique, dtype=torch.float64).unsqueeze(-1) + return cls(bank, radius=radius, k=k) + + @property + def bank(self) -> torch.Tensor: + """Support vectors ``(n_support, dim)``.""" + return self._bank + + @property + def radius(self) -> float: + """Inclusive L2 support radius.""" + return self._radius + + @property + def k(self) -> int: + """Number of neighbours used by :meth:`knn`.""" + return self._k + + @property + def n_support(self) -> int: + """Number of bank vectors.""" + return int(self._bank.shape[0]) + + def knn(self, query: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Return L2 distances and indices of the ``k`` nearest bank members. + + Args: + query: Query vectors ``(m, dim)`` matching bank feature dim. + + Returns: + ``(distances, indices)`` each of shape ``(m, k)``. Distances are + non-negative L2 norms; indices index into :attr:`bank`. + """ + q = self._validate_query(query) + if q.shape[0] == 0: + return ( + torch.empty(0, self._k, dtype=self._bank.dtype, device=q.device), + torch.empty(0, self._k, dtype=torch.long, device=q.device), + ) + # (m, n) pairwise L2 + # cdist is the primitive for kNN L2 banks; keeps this module free of + # third-party ANN deps. + dists = torch.cdist(q, self._bank.to(device=q.device, dtype=q.dtype), p=2) + k = self._k + values, indices = torch.topk(dists, k=k, dim=-1, largest=False, sorted=True) + return values, indices + + def min_distance(self, query: torch.Tensor) -> torch.Tensor: + """L2 distance to the nearest bank member for each query row. + + Args: + query: Query vectors ``(m, dim)``. + + Returns: + Distances ``(m,)``. + """ + dist, _ = self.knn(query) + if dist.numel() == 0: + return dist.reshape(0) + return dist[:, 0] + + def contains(self, query: torch.Tensor) -> torch.Tensor: + """Whether each query lies inside the support radius. + + Args: + query: Query vectors ``(m, dim)``. + + Returns: + Boolean mask ``(m,)`` — ``True`` where ``min L2 <= radius``. + """ + return self.min_distance(query) <= self._radius + + def coverage_fraction(self, query: torch.Tensor) -> float: + """Fraction of query rows inside the support radius. + + Empty queries are defined as fully covered (``1.0``) so callers can + treat "no predictions" as no coverage gap. + + Args: + query: Query vectors ``(m, dim)``. + + Returns: + Scalar in ``[0, 1]``. + """ + if query.shape[0] == 0: + return 1.0 + mask = self.contains(query) + return float(mask.float().mean().item()) + + def contains_type_ids(self, type_ids: torch.Tensor) -> torch.Tensor: + """Membership for integer type ids against a 1-D type-id bank. + + Args: + type_ids: Integer ids ``(m,)``. + + Returns: + Boolean mask ``(m,)``. + """ + if type_ids.ndim != 1: + raise ValueError(f"type_ids must be 1-D; got shape {tuple(type_ids.shape)}") + query = type_ids.to(dtype=torch.float64).unsqueeze(-1) + # Compare in bank dtype/device without forcing caller's long tensor. + return self.contains(query.to(device=self._bank.device)) + + def coverage_fraction_type_ids(self, type_ids: torch.Tensor) -> float: + """Coverage fraction for integer type-id queries. + + Args: + type_ids: Integer ids ``(m,)``. + + Returns: + Scalar in ``[0, 1]``. + """ + if type_ids.numel() == 0: + return 1.0 + mask = self.contains_type_ids(type_ids) + return float(mask.float().mean().item()) + + def _validate_query(self, query: torch.Tensor) -> torch.Tensor: + if query.ndim != 2: + raise ValueError(f"query must be 2-D (m, dim); got shape {tuple(query.shape)}") + if query.shape[-1] != self._bank.shape[-1]: + raise ValueError(f"query dim {query.shape[-1]} != bank dim {self._bank.shape[-1]}") + return query diff --git a/src/molrep/heads/__init__.py b/src/molrep/heads/__init__.py index 9859682..af5beb3 100644 --- a/src/molrep/heads/__init__.py +++ b/src/molrep/heads/__init__.py @@ -5,8 +5,15 @@ scale-shift) live in :mod:`molpot.heads` — do not import those from here. """ -from .labeler import Labeler, ProxyLabeler +from .labeler import Labeler, ProxyLabeler, TypeSystemLabeler from .scalar import ScalarHead -from .type import TypeHead +from .type import MultiTypeHead, TypeHead -__all__ = ["TypeHead", "Labeler", "ProxyLabeler", "ScalarHead"] +__all__ = [ + "TypeHead", + "MultiTypeHead", + "Labeler", + "ProxyLabeler", + "TypeSystemLabeler", + "ScalarHead", +] diff --git a/src/molrep/heads/labeler.py b/src/molrep/heads/labeler.py index 7f72190..42b07f9 100644 --- a/src/molrep/heads/labeler.py +++ b/src/molrep/heads/labeler.py @@ -1,19 +1,28 @@ -"""Labeler protocol and implementations.""" +"""Labeler protocol and implementations. + +:class:`TypeSystemLabeler` (condensed MM types) lives in +:mod:`molrep.condensation.labeler` and is re-exported here so perception-side +callers import labelers from one place. +""" from typing import Protocol, runtime_checkable import torch +from molrep.condensation.labeler import TypeSystemLabeler + +__all__ = ["Labeler", "ProxyLabeler", "TypeSystemLabeler"] + @runtime_checkable class Labeler(Protocol): - """Protocol for atom type labelers.""" + """Protocol for type labelers (atom proxy or condensed MM types).""" num_types: int type_map: dict[int, str] def label(self, z: torch.Tensor) -> torch.Tensor: - """Generate type labels for atoms in batch.""" + """Generate type labels for items in a batch.""" ... diff --git a/src/molrep/heads/scalar.py b/src/molrep/heads/scalar.py index 9071bfc..c2f3e58 100644 --- a/src/molrep/heads/scalar.py +++ b/src/molrep/heads/scalar.py @@ -1,6 +1,8 @@ import torch import torch.nn as nn +from molix import config + class ScalarHead(nn.Module): """Pool per-atom representations and predict scalar property. @@ -35,12 +37,13 @@ def __init__( raise ValueError(f"Unknown pooling: {pooling}. Use 'mean', 'sum', or 'max'") # Simple MLP: d_model -> hidden_dim -> 1 + ftype = config.ftype self.mlp = nn.Sequential( - nn.Linear(d_model, hidden_dim), + nn.Linear(d_model, hidden_dim, dtype=ftype), nn.SiLU(), - nn.Linear(hidden_dim, hidden_dim), + nn.Linear(hidden_dim, hidden_dim, dtype=ftype), nn.SiLU(), - nn.Linear(hidden_dim, 1), + nn.Linear(hidden_dim, 1, dtype=ftype), ) def forward(self, h: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: diff --git a/src/molrep/heads/type.py b/src/molrep/heads/type.py index 5daeb75..92ff770 100644 --- a/src/molrep/heads/type.py +++ b/src/molrep/heads/type.py @@ -1,32 +1,85 @@ -"""Type Head for discrete atom type classification.""" +"""Type heads for discrete classification (atom types and multi-class systems). + +``TypeHead`` remains the single-system atom/interaction classifier. +``MultiTypeHead`` holds one :class:`TypeHead` per named interaction class +without a ``method=`` switch — composition is the caller's job. +""" + +from __future__ import annotations + +from collections.abc import Mapping +from typing import TYPE_CHECKING import torch import torch.nn as nn +from molix import config + +if TYPE_CHECKING: + from molrep.condensation.type_system import TypeSystem + +__all__ = ["TypeHead", "MultiTypeHead"] + class TypeHead(nn.Module): - """Classification head for discrete atom type prediction. + """Classification head for discrete type prediction. Args: - hidden_dim: Dimension of input embeddings - num_types: Number of type classes - dropout: Dropout rate + hidden_dim: Dimension of input embeddings. + num_types: Number of type classes. + dropout: Dropout rate. """ def __init__(self, hidden_dim: int, num_types: int, dropout: float = 0.0): super().__init__() + if num_types < 1: + raise ValueError(f"num_types must be >= 1, got {num_types}") self.hidden_dim = hidden_dim self.num_types = num_types + ftype = config.ftype self.classifier = nn.Sequential( - nn.Linear(hidden_dim, hidden_dim), + nn.Linear(hidden_dim, hidden_dim, dtype=ftype), nn.SiLU(), nn.Dropout(dropout), - nn.Linear(hidden_dim, num_types), + nn.Linear(hidden_dim, num_types, dtype=ftype), ) + @classmethod + def from_type_system( + cls, + hidden_dim: int, + type_system: TypeSystem, + *, + dropout: float = 0.0, + ) -> TypeHead: + """Build a head whose ``num_types`` matches a condensed type system. + + Args: + hidden_dim: Input embedding dimension. + type_system: Frozen :class:`~molrep.condensation.TypeSystem`. + dropout: Dropout rate. + + Returns: + A :class:`TypeHead` with ``num_types == type_system.n_types``. + + Raises: + ValueError: If the type system is empty. + """ + n = type_system.n_types + if n < 1: + raise ValueError("TypeHead.from_type_system requires n_types >= 1") + return cls(hidden_dim=hidden_dim, num_types=n, dropout=dropout) + def forward(self, embeddings: torch.Tensor) -> torch.Tensor: - """Compute type logits from embeddings.""" + """Compute type logits from embeddings. + + Args: + embeddings: Features ``(N, hidden_dim)``. + + Returns: + Logits ``(N, num_types)``. + """ return self.classifier(embeddings) def decode(self, logits: torch.Tensor) -> torch.Tensor: @@ -46,7 +99,93 @@ def decode_with_confidence( self, logits: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: - """Decode with confidence scores.""" + """Decode with confidence scores. + + Returns: + ``(indices, confidence)`` from softmax max. + """ probs = torch.softmax(logits, dim=-1) confidence, indices = probs.max(dim=-1) return indices, confidence + + +class MultiTypeHead(nn.Module): + """Named collection of :class:`TypeHead` modules (multi-system / multi-class). + + Does not use a ``method=`` switch — each head is an independent classifier + keyed by interaction name (typically :class:`InteractionClass` value). + + Args: + heads: Mapping from class key (str) to a configured :class:`TypeHead`. + """ + + def __init__(self, heads: Mapping[str, TypeHead]) -> None: + super().__init__() + if not heads: + raise ValueError("MultiTypeHead requires at least one TypeHead") + self.heads = nn.ModuleDict({str(k): v for k, v in heads.items()}) + + @classmethod + def from_type_systems( + cls, + hidden_dims: Mapping[str, int] | int, + type_systems: Mapping[str, TypeSystem], + *, + dropout: float = 0.0, + ) -> MultiTypeHead: + """Build one head per condensed type system. + + Args: + hidden_dims: Per-key hidden dim, or a single int shared by all. + type_systems: Mapping of class key → :class:`TypeSystem`. + dropout: Dropout rate for every head. + + Returns: + A :class:`MultiTypeHead` covering every non-empty type system. + """ + heads: dict[str, TypeHead] = {} + for key, ts in type_systems.items(): + if ts.n_types < 1: + continue + dim = hidden_dims if isinstance(hidden_dims, int) else hidden_dims[key] + heads[str(key)] = TypeHead.from_type_system(dim, ts, dropout=dropout) + return cls(heads) + + @property + def num_types(self) -> dict[str, int]: + """Per-key ``num_types`` for each contained head.""" + return {k: h.num_types for k, h in self.heads.items()} + + def forward( + self, + embeddings: Mapping[str, torch.Tensor], + ) -> dict[str, torch.Tensor]: + """Compute logits for each provided embedding key that has a head. + + Args: + embeddings: Mapping of class key → features ``(N_k, D_k)``. + + Returns: + Mapping of class key → logits ``(N_k, num_types_k)``. + """ + out: dict[str, torch.Tensor] = {} + for key, emb in embeddings.items(): + k = str(key) + if k in self.heads: + out[k] = self.heads[k](emb) + return out + + def decode(self, logits: Mapping[str, torch.Tensor]) -> dict[str, torch.Tensor]: + """Argmax decode per class key.""" + return {k: self.heads[k].decode(v) for k, v in logits.items() if k in self.heads} + + def decode_with_confidence( + self, + logits: Mapping[str, torch.Tensor], + ) -> dict[str, tuple[torch.Tensor, torch.Tensor]]: + """Softmax-max decode with confidence per class key.""" + return { + k: self.heads[k].decode_with_confidence(v) + for k, v in logits.items() + if k in self.heads + } diff --git a/src/molrep/interaction/__init__.py b/src/molrep/interaction/__init__.py index 7a830bc..7355050 100644 --- a/src/molrep/interaction/__init__.py +++ b/src/molrep/interaction/__init__.py @@ -8,9 +8,9 @@ from .element import ElementUpdate, ElementUpdateSpec from .gate import GatedNonlinearity from .linear import EquivariantLinear +from .mace.conv import ConvTP, ConvTPSpec +from .mace.density import DensityInteraction, DensityResidualInteraction from .product import ( - ConvTP, - ConvTPSpec, irreps_from_l_max, sh_irreps_from_l_max, ) @@ -19,6 +19,8 @@ from .residual import ResidualInteraction __all__ = [ + "DensityInteraction", + "DensityResidualInteraction", "GatedNonlinearity", "ResidualInteraction", "EquivariantProductBasis", diff --git a/src/molrep/interaction/density.py b/src/molrep/interaction/density.py new file mode 100644 index 0000000..76162c8 --- /dev/null +++ b/src/molrep/interaction/density.py @@ -0,0 +1,11 @@ +"""Deprecated location — moved to :mod:`molrep.interaction.mace.density`. + +Kept as a re-export shim for chain step mace-subpackage-restructure-01; +removed in 06-wire. +""" + +from molrep.interaction.mace.density import ( # noqa: F401 + SKIP_TP_METHOD, + DensityInteraction, + DensityResidualInteraction, +) diff --git a/src/molrep/interaction/element.py b/src/molrep/interaction/element.py index 16f4965..b8ccc91 100644 --- a/src/molrep/interaction/element.py +++ b/src/molrep/interaction/element.py @@ -51,9 +51,11 @@ class ElementUpdate(nn.Module): where each element type $z \\in [0, \\text{num_species}]$ has its own $(\\text{hidden_dim} \\times \\text{hidden_dim})$ weight matrix $W_z$. - Uses cuEquivariance's **indexed_linear** backend for 8-11x speedup: - - **indexed_linear**: Hardware-optimized kernel for sorted species indices - - **naive**: Fallback for general cases + Uses cuEquivariance's ``naive`` indexed-weights backend on every device. + The ``indexed_linear`` CUDA kernel requires *sorted* species indices, and + the argsort + un-permute round-trip that requirement forces was measured + **2.8x slower** than the naive path at production shapes (N=672, H=128, + GH200) — the kernel's own advantage never survives the reordering. Physical Interpretation: "I take my current state and add element-specific weighted information @@ -94,20 +96,10 @@ def __init__( # Create cuEquivariance irreps (scalars only for hidden features) irreps = cue.Irreps("O3", f"{hidden_dim}x0e") - # ``indexed_linear`` is CUDA-only and asserts sorted indices. The - # ``naive`` path works on both CUDA and CPU. Keep two layers that - # share the same weight parameter and dispatch at forward time based - # on the actual tensor device (torch.cuda.is_available() can be true - # while the model runs on CPU, which is how most unit tests run). - self._linear_indexed = cuet.Linear( - irreps_in=irreps, - irreps_out=irreps, - internal_weights=False, - weight_classes=num_species, - layout=cue.ir_mul, - method="indexed_linear", - dtype=config.ftype, - ) + # ``naive`` on every device: the alternative ``indexed_linear`` CUDA + # kernel asserts sorted indices, and the argsort + un-permute round + # trip that costs measured 2.8x slower than naive at production + # shapes (0.974 ms vs 0.343 ms, N=672 H=128, GH200). self._linear_naive = cuet.Linear( irreps_in=irreps, irreps_out=irreps, @@ -144,35 +136,16 @@ def forward( Updated features $(n\\_{nodes}, \\text{hidden_dim})$ via: $$h\\_new[i] = h\\_prev[i] + W_{Z[i]} \\otimes m\\_curr[i]$$ - Implementation (cuEquivariance indexed_linear): - Uses hardware-optimized kernel that: - 1. Performs indexed lookup of weight matrices (W[species]) - 2. Applies element-specific linear transformation in single kernel - 3. Aggregates with h_prev for residual connection - Performance: - On CUDA the ``indexed_linear`` kernel is 8-11× faster than - the naive loop but is CUDA-only and asserts - ``weight_indices`` is non-decreasing. Real atom_types from a - collated batch are arbitrary, so we sort before and - un-permute after. On CPU we use the naive path and skip the - sort (the kernel assertion doesn't apply and ``argsort`` - would just be overhead). + The ``naive`` indexed-weights path runs on any device with + arbitrary index order. cuEq's ``indexed_linear`` kernel needs + sorted indices; the argsort + un-permute round-trip that forces + was measured 2.8x slower than this path at production shapes + (GH200), so it is deliberately not used. """ - if m_curr.is_cuda: - perm = torch.argsort(atom_types, stable=True) - inv_perm = torch.empty_like(perm) - inv_perm[perm] = torch.arange(perm.numel(), device=perm.device) - m_transformed = self._linear_indexed( - m_curr[perm], - weight=self.weight, - weight_indices=atom_types[perm], - )[inv_perm] - else: - m_transformed = self._linear_naive( - m_curr, - weight=self.weight, - weight_indices=atom_types, - ) - + m_transformed = self._linear_naive( + m_curr, + weight=self.weight, + weight_indices=atom_types, + ) return h_prev + m_transformed diff --git a/src/molrep/interaction/mace/__init__.py b/src/molrep/interaction/mace/__init__.py new file mode 100644 index 0000000..3d56a88 --- /dev/null +++ b/src/molrep/interaction/mace/__init__.py @@ -0,0 +1,36 @@ +"""MACE-only interaction blocks. + +Home for the interaction-layer building blocks that exist solely to serve the +MACE family of encoders (plain MACE, MACE-MP / MatPES, MACE-OMOL). These are +**pure geometric-convolution blocks**: they consume node/edge features and emit +node/edge features, and never compute an energy or a force — energy heads and +force derivation stay in ``molpot``. + +This mirrors the role of :mod:`molrep.interaction.pinet` for the PiNet family: +one namespace per model family, with the genuinely shared blocks +(``contraction`` / ``radial`` / ``gate`` / ``linear`` / ``aggregation`` / +``product``) remaining at the :mod:`molrep.interaction` top level. + +Reference: + Batatia et al. "MACE: Higher Order Equivariant Message Passing Neural + Networks for Fast and Accurate Force Fields" NeurIPS 2022. + https://arxiv.org/abs/2206.07697 +""" + +from molrep.interaction.mace.block import InteractionBlock, InteractionSpec +from molrep.interaction.mace.conv import ConvTP, ConvTPSpec +from molrep.interaction.mace.density import ( + SKIP_TP_METHOD, + DensityInteraction, + DensityResidualInteraction, +) + +__all__ = [ + "ConvTP", + "ConvTPSpec", + "InteractionBlock", + "InteractionSpec", + "DensityInteraction", + "DensityResidualInteraction", + "SKIP_TP_METHOD", +] diff --git a/src/molrep/interaction/mace/block.py b/src/molrep/interaction/mace/block.py new file mode 100644 index 0000000..da24b69 --- /dev/null +++ b/src/molrep/interaction/mace/block.py @@ -0,0 +1,194 @@ +"""MACE interaction block — one equivariant message-passing layer. + +:class:`InteractionBlock` wires the pre-convolution equivariant linear, the +radial weight MLP, :class:`~molrep.interaction.mace.conv.ConvTP` and the +post-convolution linear into the ``avg_num_neighbors``-normalised layer used by +the plain MACE encoder. The density-normalised variants (MACE-MP / MatPES) live +in :mod:`molrep.interaction.mace.density`. + +Relocated verbatim from ``molzoo/mace.py``, which continues to re-export both +names for backwards compatibility. + +Reference: + Batatia et al. "MACE: Higher Order Equivariant Message Passing Neural + Networks for Fast and Accurate Force Fields" NeurIPS 2022. + https://arxiv.org/abs/2206.07697 + https://docs.nvidia.com/cuda/cuequivariance/tutorials/pytorch/MACE.html +""" + +from __future__ import annotations + +import cuequivariance as cue +import cuequivariance_torch as cuet +import torch +import torch.nn as nn +from pydantic import BaseModel, ConfigDict, Field + +from molix import config +from molrep.interaction.mace.conv import ConvTP +from molrep.interaction.product import irreps_from_l_max, sh_irreps_from_l_max +from molrep.interaction.radial import RadialWeightMLP + + +class InteractionSpec(BaseModel): + """Configuration for a single interaction block. + + Attributes: + num_features: Scalar channel multiplicity. + num_bessel: Number of Bessel radial basis functions. + l_max: Maximum angular momentum order. + avg_num_neighbors: Average number of neighbors for normalization. + """ + + model_config = ConfigDict(arbitrary_types_allowed=True) + + num_features: int = Field(..., gt=0) + num_bessel: int = Field(8, gt=0) + l_max: int = Field(2, ge=0) + avg_num_neighbors: float = Field(1.0, gt=0.0) + use_fallback: bool = True + + +class InteractionBlock(nn.Module): + """Equivariant message passing with tensor product convolution. + + Performs geometric message passing via cuEquivariance-accelerated tensor products, + returning updated node features and skip connection for residual updates. + + Architecture: + node_feats → node_linear → tensor_product(edge_attrs, tp_weights) + → aggregate → linear → (node_feats_out, skip_connection) + + Attributes: + conv_tp: Tensor product convolution (cuEquivariance ChannelWiseTensorProduct). + node_linear: Pre-convolution equivariant linear transformation. + radial_mlp: MLP generating tensor product weights from edge features. + linear: Post-convolution equivariant linear projection. + avg_num_neighbors: Message normalization constant. + + Reference: + https://docs.nvidia.com/cuda/cuequivariance/tutorials/pytorch/MACE.html + """ + + def __init__( + self, + *, + num_features: int, + num_bessel: int = 8, + l_max: int = 2, + avg_num_neighbors: float = 1.0, + use_fallback: bool = True, + ): + """Initialize interaction block. + + Args: + num_features: Scalar channel multiplicity. + num_bessel: Number of Bessel basis functions. + l_max: Maximum angular momentum order. + avg_num_neighbors: Average neighbor count for message normalization. + use_fallback: Pure-torch cuEq path for the tensor product (default + ``True``, functorch-safe); ``False`` selects the fused kernels + for autograd-backed force paths. + """ + super().__init__() + + self.config = InteractionSpec( + num_features=num_features, + num_bessel=num_bessel, + l_max=l_max, + avg_num_neighbors=avg_num_neighbors, + use_fallback=use_fallback, + ) + + # Node *state* is pure scalar (l=0); the mixed-l message irreps live only + # transiently in the tensor-product output, where they are contracted + # back to invariant scalars by the downstream ProductHead. Keeping the + # node state scalar makes every node-state operation (node_linear, + # ElementUpdate, projections) equivariant by construction — fabricating + # l>0 node components from scalars via a plain/dense linear is exactly + # what breaks rotation invariance. + node_irreps_str = f"{num_features}x0e" + irreps_str = irreps_from_l_max(l_max, num_features) # mixed-l message irreps + sh_irreps_str = sh_irreps_from_l_max(l_max) + + # 1. Tensor product convolution (define first to get weight_numel): + # scalar node features ⊗ Y_l(r̂) -> mixed-l equivariant messages. + self.conv_tp = ConvTP( + in_irreps=node_irreps_str, + out_irreps=irreps_str, + sh_irreps=sh_irreps_str, + use_fallback=use_fallback, + ) + + # Actual TP output irreps (may differ from requested out_irreps) + tp_out_irreps = str(self.conv_tp.cue_tp.irreps_out) + + # 2. Pre-convolution equivariant linear (scalar -> scalar) + self.node_linear = cuet.Linear( + irreps_in=cue.Irreps("O3", node_irreps_str), + irreps_out=cue.Irreps("O3", node_irreps_str), + layout=cue.ir_mul, + dtype=config.ftype, + ) + + # 3. Radial MLP for TP weights + self.radial_mlp = RadialWeightMLP( + in_dim=num_bessel, + hidden_dim=num_features, + out_dim=self.conv_tp.weight_numel, + num_layers=2, + ) + + # 4. Post-convolution equivariant linear + self.linear = cuet.Linear( + irreps_in=cue.Irreps("O3", tp_out_irreps), + irreps_out=cue.Irreps("O3", irreps_str), + layout=cue.ir_mul, + dtype=config.ftype, + ) + + self.avg_num_neighbors = avg_num_neighbors + + def forward( + self, + node_feats: torch.Tensor, + edge_attrs: torch.Tensor, + edge_feats: torch.Tensor, + edge_index: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Run one interaction layer. + + Args: + node_feats: Node features ``(n_nodes, irreps_dim)``. + edge_attrs: Spherical harmonics ``(n_edges, sh_dim)``. + edge_feats: Radial basis features ``(n_edges, num_bessel)``. + edge_index: Edge indices ``(n_edges, 2)``. + + Returns: + tuple of: + - ``node_feats``: Updated node features ``(n_nodes, irreps_dim)``. + - ``sc``: Skip connection (original input) ``(n_nodes, irreps_dim)``. + """ + sc = node_feats # skip connection for EquivariantProductBasisBlock + + # Pre-convolution linear + node_feats_up = self.node_linear(node_feats) + + # TP weights from radial basis + tp_weights = self.radial_mlp(edge_feats) + + # Tensor product convolution with neighbor aggregation + messages = self.conv_tp( + node_features=node_feats_up, + edge_angular=edge_attrs, + edge_index=edge_index, + tp_weights=tp_weights, + ) + + # Normalize by average number of neighbors + messages = messages / self.avg_num_neighbors + + # Post-convolution linear + node_feats = self.linear(messages) + + return node_feats, sc diff --git a/src/molrep/interaction/mace/conv.py b/src/molrep/interaction/mace/conv.py new file mode 100644 index 0000000..8de60a6 --- /dev/null +++ b/src/molrep/interaction/mace/conv.py @@ -0,0 +1,134 @@ +"""Channel-wise tensor product convolution for MACE-style message passing. + +:class:`ConvTP` is a thin wrapper around ``cuet.ChannelWiseTensorProduct`` +(subscripts ``"uv,iu,jv,kuv+ijk"``) with gather/scatter folded in via +``indices_1``/``indices_out``/``size_out``, mirroring the signature of +``cuet.ChannelWiseTensorProduct``. + +Relocated verbatim from :mod:`molrep.interaction.product`, which keeps the +generic :class:`~molrep.interaction.product.EquivariantPolynomialTP` wrapper +and the ``irreps_from_l_max`` / ``sh_irreps_from_l_max`` helpers (shared with +non-MACE encoders) and re-exports the two names below for backwards +compatibility. + +Reference: + Batatia et al. "MACE: Higher Order Equivariant Message Passing Neural + Networks for Fast and Accurate Force Fields" NeurIPS 2022. + https://arxiv.org/abs/2206.07697 +""" + +from __future__ import annotations + +import cuequivariance as cue +import cuequivariance_torch as cuet +import torch +import torch.nn as nn +from pydantic import BaseModel + + +class ConvTPSpec(BaseModel): + r"""Specification for tensor product convolution layer. + + One-particle basis: + $\phi_{ij} = \sum_{l_1,l_2,m_1,m_2} c_{l_3 m_3}^{l_1 m_1, l_2 m_2} + R(r_{ij}) Y_{l_1}^{m_1}(\hat{r}_{ij}) h_j^{l_2 m_2}$ + + Attributes: + in_irreps: Input irreps. + out_irreps: Output irreps. + sh_irreps: Spherical harmonics irreps. + """ + + in_irreps: str + out_irreps: str + sh_irreps: str + + +class ConvTP(nn.Module): + r"""Channelwise tensor product for equivariant message passing. + + Computes messages via tensor product: + $$\phi_{ij} = \sum_{l_1,l_2,m_1,m_2} c_{l_3 m_3}^{l_1 m_1, l_2 m_2} + R(r_{ij}) Y_{l_1}^{m_1}(\hat{r}_{ij}) h_j^{l_2 m_2}$$ + + Attributes: + config: ConvTPSpec configuration. + cue_tp: ChannelWiseTensorProduct layer. + weight_numel: Number of elements in TP weights. + """ + + def __init__( + self, + *, + in_irreps: str, + out_irreps: str, + sh_irreps: str, + use_fallback: bool = True, + ): + """Initialize channelwise tensor product layer. + + Args: + in_irreps: Input irreps for node features. + out_irreps: Output irreps for messages. + sh_irreps: Irreps for spherical harmonics. + use_fallback: If ``True`` (default), pure-torch cuEq path so + ``ForceDerivation(method="functorch")`` can trace. Set + ``False`` for fused kernels when forces use + ``method="autograd"`` (e.g. MACE-OMOL). + """ + super().__init__() + + self.config = ConvTPSpec( + in_irreps=in_irreps, + out_irreps=out_irreps, + sh_irreps=sh_irreps, + ) + self.use_fallback = use_fallback + + irreps_in = cue.Irreps("O3", in_irreps) + irreps_sh = cue.Irreps("O3", sh_irreps) + irreps_out = cue.Irreps("O3", out_irreps) + + self.cue_tp = cuet.ChannelWiseTensorProduct( + irreps_in, + irreps_sh, + irreps_out, + layout=cue.ir_mul, + shared_weights=False, + internal_weights=False, + use_fallback=use_fallback, + ) + + self.weight_numel = self.cue_tp.weight_numel + + def forward( + self, + node_features: torch.Tensor, + edge_angular: torch.Tensor, + edge_index: torch.Tensor, + tp_weights: torch.Tensor, + ) -> torch.Tensor: + """Compute tensor product messages with integrated gather/scatter. + + Args: + node_features: Node features. + edge_angular: Spherical harmonics. + edge_index: Edge indices ``(E, 2)``. + tp_weights: TP weights. + + Returns: + Computed messages (n_edges, out_irreps_dim). + """ + indices_1 = edge_index[:, 0] + indices_out = edge_index[:, 1] + + messages = self.cue_tp( + node_features, + edge_angular, + tp_weights, + indices_1=indices_1, + indices_out=indices_out, + size_out=node_features.shape[0], + ) + + return messages diff --git a/src/molrep/interaction/mace/density.py b/src/molrep/interaction/mace/density.py new file mode 100644 index 0000000..ca7159f --- /dev/null +++ b/src/molrep/interaction/mace/density.py @@ -0,0 +1,292 @@ +"""Density-normalised equivariant interaction blocks (MACE-MP / MatPES variant). + +Faithful port of MACE's ``RealAgnosticDensityInteractionBlock`` and +``RealAgnosticDensityResidualInteractionBlock`` on the cuEquivariance backend. +One edge→node message-passing layer with: + +* a channel-wise tensor product ``node ⊗ Y_l(r̂)`` with neighbour scatter, +* a **learned density** normalisation ``message / (ρ + 1)`` where + ``ρ_i = Σ_j tanh(MLP(edge_feats_ij)²)·u(r_ij)`` — this replaces the fixed + ``avg_num_neighbors`` divisor and is what lets one model span the + density range of the periodic table, +* an element-selecting ``skip_tp`` (fully-connected tensor product against the + one-hot node attributes). + +The two blocks differ only in where ``skip_tp`` acts: the first-layer variant +applies it to the outgoing message (there is no residual to carry), while the +residual variant applies it to the incoming node features and hands the result +to the downstream product block as the skip connection. + +Built entirely from ``cuequivariance`` primitives in the ``cue.ir_mul`` layout; +all sub-layer names mirror MACE so the official weights transfer by direct copy. + +Reference: + Batatia et al. "MACE: Higher Order Equivariant Message Passing Neural + Networks for Fast and Accurate Force Fields" NeurIPS 2022. + https://arxiv.org/abs/2206.07697 + Batatia et al. "A foundation model for atomistic materials chemistry" + (MACE-MP-0). https://arxiv.org/abs/2401.00096 +""" + +from __future__ import annotations + +import cuequivariance as cue +import cuequivariance_torch as cuet +import torch +import torch.nn as nn + +from molix import config +from molix.F.scatter import scatter_sum_compile_safe as _scatter_sum +from molrep.embedding.mlp import MomentNormalizedMLP + +#: cuEquivariance execution method for ``skip_tp``. +#: +#: ``skip_tp`` contracts the node features against a **one-hot** element vector +#: (``89x0e`` for a MatPES model), a shape cuEq's ``fused_tp`` kernel handles +#: badly: on GH200 it measures 33 ms against 1.7 ms for ``naive`` on a 193-atom +#: graph — a 20x penalty for bit-identical output (1.4e-15). ``naive`` is also +#: what MACE's own converted cueq model runs, so this default keeps us on the +#: reference's execution path as well as its math. +SKIP_TP_METHOD = "naive" + + +class _DensityInteractionBase(nn.Module): + """Shared wiring of the two density-normalised interaction blocks. + + Subclasses own only ``skip_tp``'s irreps and the placement of the skip + connection in :meth:`forward`. + + Args: + node_attrs_irreps: One-hot atomic-number irreps, e.g. ``"89x0e"``. + node_feats_irreps: Incoming node feature irreps. + edge_attrs_irreps: Spherical-harmonics irreps, e.g. ``"1x0e+1x1o+1x2e+1x3o"``. + edge_feats_irreps: Radial basis irreps, e.g. ``"10x0e"``. + edge_irreps: ``linear_up`` output irreps. + target_irreps: Tensor-product / block output irreps. + radial_mlp: Hidden widths of the radial weight MLP, e.g. ``[64, 64, 64]``. + use_fallback: Pure-torch cuEq path (default ``True``, functorch-safe). + Set ``False`` for fused kernels when forces use autograd. + skip_tp_method: cuEquivariance execution method for ``skip_tp``; see + :data:`SKIP_TP_METHOD`. + """ + + def __init__( + self, + *, + node_attrs_irreps: str, + node_feats_irreps: str, + edge_attrs_irreps: str, + edge_feats_irreps: str, + edge_irreps: str, + target_irreps: str, + radial_mlp: list[int], + use_fallback: bool = True, + skip_tp_method: str = SKIP_TP_METHOD, + ) -> None: + super().__init__() + ftype = config.ftype + self.use_fallback = use_fallback + self.skip_tp_method = skip_tp_method + + node_feats = cue.Irreps("O3", node_feats_irreps) + edge_attrs = cue.Irreps("O3", edge_attrs_irreps) + edge_feats = cue.Irreps("O3", edge_feats_irreps) + edge_ir = cue.Irreps("O3", edge_irreps) + target = cue.Irreps("O3", target_irreps) + self.node_attrs_irreps = cue.Irreps("O3", node_attrs_irreps) + self.irreps_out = target + + self.linear_up = cuet.Linear(node_feats, edge_ir, layout=cue.ir_mul, dtype=ftype) + + self.conv_tp = cuet.ChannelWiseTensorProduct( + edge_ir, + edge_attrs, + target, + layout=cue.ir_mul, + shared_weights=False, + internal_weights=False, + dtype=ftype, + use_fallback=use_fallback, + ) + + num_radial = edge_feats.dim + self.conv_tp_weights = MomentNormalizedMLP( + [num_radial] + list(radial_mlp) + [self.conv_tp.weight_numel] + ) + self.density_fn = MomentNormalizedMLP([num_radial, 1]) + + self.linear = cuet.Linear(self.conv_tp.irreps_out, target, layout=cue.ir_mul, dtype=ftype) + + # reshape_irreps (ir_mul): per-irrep (N, mul*d) -> (N, d, mul), cat on dim -2 + self._reshape_dims = [(mi.mul, mi.ir.dim) for mi in target] + + def _build_skip_tp(self, dtype: torch.dtype) -> cuet.FullyConnectedTensorProduct: + """Construct ``skip_tp`` at ``dtype``; subclasses choose the irreps.""" + raise NotImplementedError + + def _apply(self, *args, **kwargs): + """Rebuild ``skip_tp`` when the module's dtype changes. + + ``cuet.FullyConnectedTensorProduct`` compiles its contraction graph at + construction and bakes the working precision into it, so + ``nn.Module.double()`` converts the parameters but leaves the graph in + float32 — the next forward then raises a bare dtype mismatch deep inside + cuEquivariance. Rebuilding at the new dtype and re-adopting the (already + converted) weight makes ``Model(...).double()`` behave like it does for + every other module here, instead of only supporting construction under + ``molix.config.set_precision("fp64")``. + """ + module = super()._apply(*args, **kwargs) + weight = module.skip_tp.weight + if weight.dtype != module._skip_tp_dtype: + rebuilt = module._build_skip_tp(weight.dtype).to(weight.device) + with torch.no_grad(): + rebuilt.weight.copy_(weight) + rebuilt.weight.requires_grad_(weight.requires_grad) + module.skip_tp = rebuilt + module._skip_tp_dtype = weight.dtype + return module + + def _reshape(self, tensor: torch.Tensor) -> torch.Tensor: + ix, out, batch = 0, [], tensor.shape[0] + for mul, d in self._reshape_dims: + field = tensor[:, ix : ix + mul * d].reshape(batch, d, mul) + ix += mul * d + out.append(field) + return torch.cat(out, dim=-2) + + def _message( + self, + node_feats: torch.Tensor, + edge_attrs: torch.Tensor, + edge_feats: torch.Tensor, + edge_index: torch.Tensor, + cutoff: torch.Tensor | None, + ) -> torch.Tensor: + """Density-normalised aggregated message ``(N, target.dim)``.""" + num_nodes = node_feats.shape[0] + source, target = edge_index[:, 0], edge_index[:, 1] + node_feats = self.linear_up(node_feats) + + tp_weights = self.conv_tp_weights(edge_feats) + edge_density = torch.tanh(self.density_fn(edge_feats) ** 2) + if cutoff is not None: + tp_weights = tp_weights * cutoff + edge_density = edge_density * cutoff + density = _scatter_sum(edge_density, target, num_nodes) + + mji = self.conv_tp(node_feats[source], edge_attrs, tp_weights) + message = _scatter_sum(mji, target, num_nodes) + return self.linear(message) / (density + 1.0) + + +class DensityInteraction(_DensityInteractionBase): + """First-layer density-normalised interaction (MACE ``RealAgnosticDensity``). + + ``skip_tp`` mixes the outgoing message with the one-hot element attributes; + there is no residual to pass on, so the returned skip connection is ``None``. + + Args: + See :class:`_DensityInteractionBase`. + """ + + def __init__(self, **kwargs) -> None: + super().__init__(**kwargs) + self.skip_tp = self._build_skip_tp(config.ftype) + self._skip_tp_dtype = config.ftype + + def _build_skip_tp(self, dtype: torch.dtype) -> cuet.FullyConnectedTensorProduct: + """``message ⊗ one_hot(Z) → message``: element-selective output mixing.""" + return cuet.FullyConnectedTensorProduct( + self.irreps_out, + self.node_attrs_irreps, + self.irreps_out, + layout=cue.ir_mul, + dtype=dtype, + method=self.skip_tp_method, + ) + + def forward( + self, + node_attrs: torch.Tensor, + node_feats: torch.Tensor, + edge_attrs: torch.Tensor, + edge_feats: torch.Tensor, + edge_index: torch.Tensor, + cutoff: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, None]: + """Run one first-layer interaction. + + Args: + node_attrs: One-hot atomic numbers ``(N, n_elements)``. + node_feats: Node features ``(N, node_feats_irreps.dim)``. + edge_attrs: Spherical harmonics ``(E, edge_attrs_irreps.dim)``. + edge_feats: Radial basis features ``(E, num_radial)``. + edge_index: ``(E, 2)`` with ``[:, 0]`` = source, ``[:, 1]`` = target + (the repo-wide edge convention). + cutoff: Optional per-edge cutoff envelope ``(E, 1)``. + + Returns: + ``(reshaped_message (N, ir_dim, mul), None)``. + """ + message = self._message(node_feats, edge_attrs, edge_feats, edge_index, cutoff) + message = self.skip_tp(message, node_attrs) + return self._reshape(message), None + + +class DensityResidualInteraction(_DensityInteractionBase): + """Residual density-normalised interaction (MACE ``RealAgnosticDensityResidual``). + + ``skip_tp`` maps the *incoming* node features (mixed with the one-hot element + attributes) onto ``hidden_irreps``; the result is the skip connection the + downstream product block adds back. + + Args: + hidden_irreps: ``skip_tp`` output irreps (consumed by the product block). + Other arguments: see :class:`_DensityInteractionBase`. + """ + + def __init__(self, *, hidden_irreps: str, node_feats_irreps: str, **kwargs) -> None: + super().__init__(node_feats_irreps=node_feats_irreps, **kwargs) + self._skip_tp_in = cue.Irreps("O3", node_feats_irreps) + self._skip_tp_out = cue.Irreps("O3", hidden_irreps) + self.skip_tp = self._build_skip_tp(config.ftype) + self._skip_tp_dtype = config.ftype + + def _build_skip_tp(self, dtype: torch.dtype) -> cuet.FullyConnectedTensorProduct: + """``node_feats ⊗ one_hot(Z) → hidden``: the residual the product adds back.""" + return cuet.FullyConnectedTensorProduct( + self._skip_tp_in, + self.node_attrs_irreps, + self._skip_tp_out, + layout=cue.ir_mul, + dtype=dtype, + method=self.skip_tp_method, + ) + + def forward( + self, + node_attrs: torch.Tensor, + node_feats: torch.Tensor, + edge_attrs: torch.Tensor, + edge_feats: torch.Tensor, + edge_index: torch.Tensor, + cutoff: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Run one residual interaction. + + Args: + node_attrs: One-hot atomic numbers ``(N, n_elements)``. + node_feats: Node features ``(N, node_feats_irreps.dim)``. + edge_attrs: Spherical harmonics ``(E, edge_attrs_irreps.dim)``. + edge_feats: Radial basis features ``(E, num_radial)``. + edge_index: ``(E, 2)`` with ``[:, 0]`` = source, ``[:, 1]`` = target + (the repo-wide edge convention). + cutoff: Optional per-edge cutoff envelope ``(E, 1)``. + + Returns: + ``(reshaped_message (N, ir_dim, mul), skip (N, hidden_irreps.dim))``. + """ + sc = self.skip_tp(node_feats, node_attrs) + message = self._message(node_feats, edge_attrs, edge_feats, edge_index, cutoff) + return self._reshape(message), sc diff --git a/src/molrep/interaction/product.py b/src/molrep/interaction/product.py index 045eb56..df84600 100644 --- a/src/molrep/interaction/product.py +++ b/src/molrep/interaction/product.py @@ -1,16 +1,18 @@ -"""Tensor product convolution for equivariant message passing. +"""Generic tensor products and irreps helpers for equivariant message passing. -Two tensor-product wrappers are exposed here: +Exposed here: -- :class:`ConvTP`: thin wrapper around ``cuet.ChannelWiseTensorProduct`` - (subscripts ``"uv,iu,jv,kuv+ijk"``). Used by MACE-style encoders. - :class:`EquivariantPolynomialTP`: general wrapper around an arbitrary ``cue.EquivariantPolynomial``. Lets callers build custom descriptors via - ``cue.SegmentedTensorProduct.from_subscripts(...)``. + ``cue.SegmentedTensorProduct.from_subscripts(...)``. It supports + gather/scatter via ``indices_1``/``indices_2``/``indices_out``/``size_out``, + mirroring the signature of ``cuet.ChannelWiseTensorProduct``. +- ``irreps_from_l_max`` / ``sh_irreps_from_l_max``: irreps-string builders + shared by MACE, Allegro and the MACE readouts. -Both wrappers support gather/scatter via ``indices_1``/``indices_2``/ -``indices_out``/``size_out``, mirroring the signature of -``cuet.ChannelWiseTensorProduct``. +The MACE-specific :class:`~molrep.interaction.mace.conv.ConvTP` wrapper moved to +:mod:`molrep.interaction.mace.conv`; it is re-exported at the bottom of this +module so ``from molrep.interaction.product import ConvTP`` keeps working. Model-specific descriptor builders (e.g. Allegro's per-channel ``"u,iu,ju,ku+ijk"`` kernel) live in the corresponding ``molzoo`` module, not here. @@ -24,115 +26,6 @@ import cuequivariance_torch as cuet import torch import torch.nn as nn -from pydantic import BaseModel - - -class ConvTPSpec(BaseModel): - r"""Specification for tensor product convolution layer. - - One-particle basis: - $\phi_{ij} = \sum_{l_1,l_2,m_1,m_2} c_{l_3 m_3}^{l_1 m_1, l_2 m_2} - R(r_{ij}) Y_{l_1}^{m_1}(\hat{r}_{ij}) h_j^{l_2 m_2}$ - - Attributes: - in_irreps: Input irreps. - out_irreps: Output irreps. - sh_irreps: Spherical harmonics irreps. - """ - - in_irreps: str - out_irreps: str - sh_irreps: str - - -class ConvTP(nn.Module): - r"""Channelwise tensor product for equivariant message passing. - - Computes messages via tensor product: - $$\phi_{ij} = \sum_{l_1,l_2,m_1,m_2} c_{l_3 m_3}^{l_1 m_1, l_2 m_2} - R(r_{ij}) Y_{l_1}^{m_1}(\hat{r}_{ij}) h_j^{l_2 m_2}$$ - - Attributes: - config: ConvTPSpec configuration. - cue_tp: ChannelWiseTensorProduct layer. - weight_numel: Number of elements in TP weights. - """ - - def __init__( - self, - *, - in_irreps: str, - out_irreps: str, - sh_irreps: str, - use_fallback: bool = True, - ): - """Initialize channelwise tensor product layer. - - Args: - in_irreps: Input irreps for node features. - out_irreps: Output irreps for messages. - sh_irreps: Irreps for spherical harmonics. - use_fallback: If ``True`` (default), pure-torch cuEq path so - ``ForceDerivation(method="functorch")`` can trace. Set - ``False`` for fused kernels when forces use - ``method="autograd"`` (e.g. MACE-OMOL). - """ - super().__init__() - - self.config = ConvTPSpec( - in_irreps=in_irreps, - out_irreps=out_irreps, - sh_irreps=sh_irreps, - ) - self.use_fallback = use_fallback - - irreps_in = cue.Irreps("O3", in_irreps) - irreps_sh = cue.Irreps("O3", sh_irreps) - irreps_out = cue.Irreps("O3", out_irreps) - - self.cue_tp = cuet.ChannelWiseTensorProduct( # type: ignore - irreps_in, - irreps_sh, - irreps_out, - layout=cue.ir_mul, - shared_weights=False, - internal_weights=False, - use_fallback=use_fallback, - ) - - self.weight_numel = self.cue_tp.weight_numel - - def forward( - self, - node_features: torch.Tensor, - edge_angular: torch.Tensor, - edge_index: torch.Tensor, - tp_weights: torch.Tensor, - ) -> torch.Tensor: - """Compute tensor product messages with integrated gather/scatter. - - Args: - node_features: Node features. - edge_angular: Spherical harmonics. - edge_index: Edge indices ``(E, 2)``. - tp_weights: TP weights. - - Returns: - Computed messages (n_edges, out_irreps_dim). - """ - indices_1 = edge_index[:, 0] - indices_out = edge_index[:, 1] - - messages = self.cue_tp( - node_features, - edge_angular, - tp_weights, - indices_1=indices_1, - indices_out=indices_out, - size_out=node_features.shape[0], - ) - - return messages def irreps_from_l_max(l_max: int, hidden_dim: int) -> str: @@ -299,3 +192,11 @@ def forward( output_shapes=sizes_out, output_indices=output_indices, )[0] + + +# Back-compat re-export of the MACE convolution, moved to +# ``molrep.interaction.mace.conv`` by mace-subpackage-restructure-01. Kept at the +# module tail on purpose: ``mace.block`` imports the irreps helpers above from +# here, so this module must be fully populated before the mace package is pulled +# in. Removed in 06-wire. +from molrep.interaction.mace.conv import ConvTP, ConvTPSpec # noqa: E402,F401 diff --git a/src/molrep/interaction/radial.py b/src/molrep/interaction/radial.py index fc7ac54..5a0c304 100644 --- a/src/molrep/interaction/radial.py +++ b/src/molrep/interaction/radial.py @@ -73,13 +73,14 @@ def __init__( num_layers=num_layers, ) + ftype = config.ftype layers: list[nn.Module] = [] current_dim = in_dim for _ in range(num_layers): - layers.append(nn.Linear(current_dim, hidden_dim)) + layers.append(nn.Linear(current_dim, hidden_dim, dtype=ftype)) layers.append(nn.SiLU()) current_dim = hidden_dim - layers.append(nn.Linear(current_dim, out_dim)) + layers.append(nn.Linear(current_dim, out_dim, dtype=ftype)) self.mlp = nn.Sequential(*layers) diff --git a/src/molrep/interaction/residual.py b/src/molrep/interaction/residual.py index 0b3d465..7746ae7 100644 --- a/src/molrep/interaction/residual.py +++ b/src/molrep/interaction/residual.py @@ -163,29 +163,31 @@ def forward( node_feats: Node features ``(N, node_feats_irreps.dim)``. edge_attrs: Spherical harmonics ``(E, edge_attrs_irreps.dim)``. edge_feats: Radial basis features ``(E, num_bessel)``. - edge_index: ``(2, E)`` with row 0 = sender, row 1 = receiver. + edge_index: ``(E, 2)`` with ``[:, 0]`` = source, ``[:, 1]`` = target + (the repo-wide edge convention). cutoff: Optional per-edge cutoff envelope ``(E, 1)``. Returns: ``(reshaped_message (N, ir_dim, mul), skip (N, hidden_dim))``. """ num_nodes = node_feats.shape[0] + source, target = edge_index[:, 0], edge_index[:, 1] sc = self.skip_tp(node_feats) node_feats = self.linear_up(node_feats) node_feats_res = self.linear_res(node_feats) - source = self.source_embedding(node_attrs) - target = self.target_embedding(node_attrs) - edge_feats = torch.cat([edge_feats, source[edge_index[0]], target[edge_index[1]]], dim=-1) + source_emb = self.source_embedding(node_attrs) + target_emb = self.target_embedding(node_attrs) + edge_feats = torch.cat([edge_feats, source_emb[source], target_emb[target]], dim=-1) tp_weights = self.conv_tp_weights(edge_feats) edge_density = torch.tanh(self.density_fn(edge_feats) ** 2) if cutoff is not None: tp_weights = tp_weights * cutoff edge_density = edge_density * cutoff - density = _scatter_sum(edge_density, edge_index[1], num_nodes) + density = _scatter_sum(edge_density, target, num_nodes) - mji = self.conv_tp(node_feats[edge_index[0]], edge_attrs, tp_weights) - message = _scatter_sum(mji, edge_index[1], num_nodes) + mji = self.conv_tp(node_feats[source], edge_attrs, tp_weights) + message = _scatter_sum(mji, target, num_nodes) message = self.linear_1(message) / (density * self.beta + self.alpha) message = message + node_feats_res diff --git a/src/molrep/perception/__init__.py b/src/molrep/perception/__init__.py new file mode 100644 index 0000000..68eb753 --- /dev/null +++ b/src/molrep/perception/__init__.py @@ -0,0 +1,28 @@ +"""molrep.perception — symbolic SMARTS/SMIRKS chemical-perception interface. + +Binds discrete condensed classes (06) to SMARTS/SMIRKS patterns, matches them +via a :class:`SmartsMatcher` Protocol, and exposes :class:`SymbolicForceField` +for export/human inspection. Matching is pure perception — no energy evaluation. + +Spec: learnable-classical-ff-07-smarts +""" + +from molrep.perception.forcefield import SymbolicForceField +from molrep.perception.matcher import ( + FakeSmartsMatcher, + MolpySmartsMatcher, + SmartsMatcher, +) +from molrep.perception.patterns import SymbolicPattern +from molrep.perception.records import DiscreteClassRecord +from molrep.perception.registry import ClassPatternRegistry + +__all__ = [ + "ClassPatternRegistry", + "DiscreteClassRecord", + "FakeSmartsMatcher", + "MolpySmartsMatcher", + "SmartsMatcher", + "SymbolicForceField", + "SymbolicPattern", +] diff --git a/src/molrep/perception/forcefield.py b/src/molrep/perception/forcefield.py new file mode 100644 index 0000000..42856c6 --- /dev/null +++ b/src/molrep/perception/forcefield.py @@ -0,0 +1,162 @@ +"""SymbolicForceField — discrete classes + patterns + prototypes (no energy).""" + +from __future__ import annotations + +from collections.abc import Mapping +from typing import Any + +import torch + +from molrep.condensation.classes import InteractionClass +from molrep.condensation.type_system import TypeSystem +from molrep.perception.matcher import SmartsMatcher +from molrep.perception.records import DiscreteClassRecord +from molrep.perception.registry import ClassPatternRegistry + +__all__ = ["SymbolicForceField"] + + +class SymbolicForceField: + """Discrete classes + patterns + prototypes for human/export consumption. + + Pairs condensed :class:`TypeSystem` tables with a + :class:`ClassPatternRegistry`. Matching is pure perception — this module + never evaluates energy or imports molpot kernels. + + Args: + type_systems: Map of interaction class → condensed type table. + registry: Pattern bindings for discrete type ids. + + Spec: + learnable-classical-ff-07-smarts + """ + + def __init__( + self, + type_systems: Mapping[InteractionClass, TypeSystem], + registry: ClassPatternRegistry, + ) -> None: + if not isinstance(registry, ClassPatternRegistry): + raise TypeError(f"registry must be ClassPatternRegistry, got {type(registry)!r}") + systems: dict[InteractionClass, TypeSystem] = {} + for interaction, ts in type_systems.items(): + if not isinstance(interaction, InteractionClass): + raise TypeError( + f"type_systems keys must be InteractionClass, got {type(interaction)!r}" + ) + if not isinstance(ts, TypeSystem): + raise TypeError(f"type_systems values must be TypeSystem, got {type(ts)!r}") + if ts.interaction != interaction: + raise ValueError( + f"TypeSystem.interaction is {ts.interaction.value!r} but " + f"mapped under {interaction.value!r}" + ) + systems[interaction] = ts + self._type_systems = systems + self._registry = registry + + @property + def type_systems(self) -> Mapping[InteractionClass, TypeSystem]: + """Frozen view of interaction → type system.""" + return self._type_systems + + @property + def registry(self) -> ClassPatternRegistry: + """Pattern registry.""" + return self._registry + + def records(self) -> list[DiscreteClassRecord]: + """Assemble discrete class records with bound SMARTS/SMIRKS when present. + + Returns: + One :class:`DiscreteClassRecord` per type across all type systems, + ordered by interaction enum order then table order. Unbound types + have ``smarts=None`` and ``smirks=None``. + """ + out: list[DiscreteClassRecord] = [] + for interaction in InteractionClass: + ts = self._type_systems.get(interaction) + if ts is None: + continue + for rec in ts.records(): + pat = self._registry.get_optional(interaction, rec.type_id) + smarts: str | None = None + smirks: str | None = None + if pat is not None: + if pat.kind == "smirks": + smirks = pat.pattern + else: + smarts = pat.pattern + out.append( + DiscreteClassRecord( + interaction=interaction, + type_id=rec.type_id, + prototype=dict(rec.prototype), + smarts=smarts, + smirks=smirks, + label=rec.label, + ) + ) + return out + + def match_molecule( + self, + mol: Any, + matcher: SmartsMatcher, + ) -> dict[InteractionClass, dict[str, torch.Tensor]]: + """Assign type ids by matching bound patterns against ``mol``. + + For each bound ``(interaction, type_id)`` pattern, run ``matcher`` and + collect hits. **No energy evaluation.** + + Args: + mol: Molecule graph accepted by ``matcher``. + matcher: Object implementing :class:`SmartsMatcher`. + + Returns: + Mapping ``interaction → {"matches": LongTensor [arity, K], + "type_ids": LongTensor [K]}`` for interactions with at least one + hit. Interactions with zero hits are omitted. + """ + buckets: dict[InteractionClass, dict[str, list[torch.Tensor]]] = {} + + for (interaction, type_id), pattern in self._registry.items(): + hits = matcher.match(mol, pattern) + if hits.ndim != 2: + raise ValueError( + f"matcher returned shape {tuple(hits.shape)}; expected [arity, n_hits]" + ) + if hits.shape[1] == 0: + continue + type_ids = torch.full((hits.shape[1],), int(type_id), dtype=torch.long) + bucket = buckets.setdefault(interaction, {"matches": [], "type_ids": []}) + bucket["matches"].append(hits.long()) + bucket["type_ids"].append(type_ids) + + result: dict[InteractionClass, dict[str, torch.Tensor]] = {} + for interaction, bucket in buckets.items(): + # Patterns for the same interaction may have different arities + # (e.g. atom LJ vs bond); group by leading dim. + by_arity: dict[int, list[torch.Tensor]] = {} + type_by_arity: dict[int, list[torch.Tensor]] = {} + for m, t in zip(bucket["matches"], bucket["type_ids"], strict=True): + a = int(m.shape[0]) + by_arity.setdefault(a, []).append(m) + type_by_arity.setdefault(a, []).append(t) + + if len(by_arity) == 1: + arity = next(iter(by_arity)) + result[interaction] = { + "matches": torch.cat(by_arity[arity], dim=1), + "type_ids": torch.cat(type_by_arity[arity], dim=0), + } + else: + # Mixed arities under one InteractionClass — return a flat + # concatenation only if all share arity; otherwise keep the + # dominant path documented as single-arity per class. + # Spec assumes one arity per InteractionClass (bond=2, …). + raise ValueError( + f"mixed arities under {interaction.value}: " + f"{sorted(by_arity)}; bind patterns of one arity per class" + ) + return result diff --git a/src/molrep/perception/matcher.py b/src/molrep/perception/matcher.py new file mode 100644 index 0000000..87b2f75 --- /dev/null +++ b/src/molrep/perception/matcher.py @@ -0,0 +1,173 @@ +"""SmartsMatcher Protocol + Fake / molpy-backed implementations. + +Matching returns LongTensor ``[arity, n_hits]`` (COO-style columns) to align +with valence topology conventions. Energy evaluation is out of scope. + +molpy only — never bare ``molrs``. +""" + +from __future__ import annotations + +from collections.abc import Mapping +from typing import Any, Protocol, runtime_checkable + +import torch + +from molrep.perception.patterns import SymbolicPattern + +__all__ = [ + "FakeSmartsMatcher", + "MolpySmartsMatcher", + "SmartsMatcher", +] + + +@runtime_checkable +class SmartsMatcher(Protocol): + """Protocol for SMARTS/SMIRKS graph matching. + + Implementations return atom-index hits as a LongTensor shaped + ``[arity, n_hits]`` (column per hit). Empty result is ``[arity, 0]``. + """ + + def match(self, mol: Any, pattern: SymbolicPattern) -> torch.Tensor: + """Match ``pattern`` against ``mol``. + + Args: + mol: Molecule graph accepted by the backend (opaque to callers + that use :class:`FakeSmartsMatcher`; molpy ``Atomistic`` for + :class:`MolpySmartsMatcher`). + pattern: Validated symbolic pattern with arity. + + Returns: + LongTensor of shape ``[pattern.arity, n_hits]``. + """ + ... + + +class FakeSmartsMatcher: + """Deterministic test double: ``pattern_str → matches`` table. + + No molpy required. Used in unit tests that do not need real chemistry. + + Args: + hits: Mapping from pattern string to a LongTensor of shape + ``[arity, K]``. Unknown patterns yield empty ``[arity, 0]``. + + Notes: + Configured tensors must be integer dtype. When a configured tensor's + leading dimension differs from ``pattern.arity``, :meth:`match` raises + ``ValueError``. + """ + + def __init__(self, hits: Mapping[str, torch.Tensor] | None = None) -> None: + self._hits: dict[str, torch.Tensor] = {} + if hits: + for key, tensor in hits.items(): + t = torch.as_tensor(tensor) + if t.ndim != 2: + raise ValueError( + f"FakeSmartsMatcher hit for {key!r} must be 2-D " + f"[arity, K]; got shape {tuple(t.shape)}" + ) + if t.dtype not in ( + torch.int8, + torch.int16, + torch.int32, + torch.int64, + torch.long, + ): + t = t.long() + self._hits[key] = t.long().contiguous() + + def match(self, mol: Any, pattern: SymbolicPattern) -> torch.Tensor: + """Return preconfigured hits for ``pattern.pattern``, or empty. + + Args: + mol: Ignored (present for Protocol compatibility). + pattern: Pattern whose string key is looked up. + + Returns: + LongTensor ``[arity, K]``. + + Raises: + ValueError: Configured hit arity disagrees with ``pattern.arity``. + """ + del mol # unused — deterministic table lookup + hit = self._hits.get(pattern.pattern) + if hit is None: + return torch.empty((pattern.arity, 0), dtype=torch.long) + if hit.shape[0] != pattern.arity: + raise ValueError( + f"configured hit arity {hit.shape[0]} != pattern.arity " + f"{pattern.arity} for {pattern.pattern!r}" + ) + return hit + + +class MolpySmartsMatcher: + """SMARTS matching via molpy's public ``SmartsPattern`` surface. + + Soft-imports :class:`molpy.SmartsPattern` only (never bare ``molrs``). + Hits are converted to LongTensor ``[arity, n_hits]``. + + Raises: + ImportError: If ``molpy.SmartsPattern`` is unavailable at construction. + """ + + def __init__(self) -> None: + # Soft import — keep the class importable even if SmartsPattern is + # missing from a stripped molpy build; construction fails loudly. + try: + from molpy import SmartsPattern as _SmartsPattern + except ImportError as exc: # pragma: no cover - env-dependent + raise ImportError( + "MolpySmartsMatcher requires molpy.SmartsPattern on the public " + "API. Install molcrafts-molpy>=0.13; do not import molrs from " + "molnex." + ) from exc + self._SmartsPattern = _SmartsPattern + + def match(self, mol: Any, pattern: SymbolicPattern) -> torch.Tensor: + """Match ``pattern`` against a molpy molecule graph. + + Args: + mol: Object accepted by ``SmartsPattern.find_matches`` (typically + molpy ``Atomistic``). + pattern: Symbolic pattern; ``pattern.pattern`` is compiled. + + Returns: + LongTensor ``[arity, n_hits]``. Each column is the first + ``arity`` atom indices of a SMARTS embedding. + + Raises: + ValueError: A hit has fewer than ``arity`` atoms. + """ + compiled = self._SmartsPattern(pattern.pattern) + raw_hits = compiled.find_matches(mol) + if not raw_hits: + return torch.empty((pattern.arity, 0), dtype=torch.long) + + columns: list[list[int]] = [] + for hit in raw_hits: + atoms = _hit_atom_list(hit) + if len(atoms) < pattern.arity: + raise ValueError( + f"SMARTS hit has {len(atoms)} atoms but pattern arity is " + f"{pattern.arity} ({pattern.pattern!r})" + ) + columns.append([int(a) for a in atoms[: pattern.arity]]) + + # Stack as [arity, K] + return torch.tensor(columns, dtype=torch.long).T.contiguous() + + +def _hit_atom_list(hit: Any) -> list[int]: + """Normalize a molpy/molrs SmartsMatch (or list) to atom indices.""" + if isinstance(hit, (list, tuple)): + return [int(x) for x in hit] + if hasattr(hit, "as_list"): + return [int(x) for x in hit.as_list()] + if hasattr(hit, "atoms"): + return [int(x) for x in hit.atoms] + raise TypeError(f"unsupported SMARTS hit type: {type(hit)!r}") diff --git a/src/molrep/perception/patterns.py b/src/molrep/perception/patterns.py new file mode 100644 index 0000000..e500b35 --- /dev/null +++ b/src/molrep/perception/patterns.py @@ -0,0 +1,48 @@ +"""SymbolicPattern — validated SMARTS/SMIRKS string + arity.""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Literal + +__all__ = ["SymbolicPattern"] + +_VALID_ARITIES = frozenset({1, 2, 3, 4}) +PatternKind = Literal["smarts", "smirks"] + + +@dataclass(frozen=True) +class SymbolicPattern: + """A symbolic chemical pattern with fixed interaction arity. + + Args: + pattern: Non-empty SMARTS or SMIRKS string. + arity: Number of atoms in a hit (1 atom, 2 bond, 3 angle, 4 torsion). + Must be in ``{1, 2, 3, 4}``. + atom_maps: Optional documentation of Daylight atom-map labels + (query atom index → map name). Not used for matching constraints. + kind: Whether ``pattern`` is SMARTS or SMIRKS. Default ``"smarts"``. + + Raises: + ValueError: Empty pattern string or arity outside ``{1,2,3,4}``. + """ + + pattern: str + arity: int + atom_maps: Mapping[int, str] | None = None + kind: PatternKind = "smarts" + + def __post_init__(self) -> None: + if not isinstance(self.pattern, str) or not self.pattern.strip(): + raise ValueError( + f"SymbolicPattern.pattern must be a non-empty string; got {self.pattern!r}" + ) + if self.arity not in _VALID_ARITIES: + raise ValueError(f"SymbolicPattern.arity must be in {{1,2,3,4}}; got {self.arity}") + if self.kind not in ("smarts", "smirks"): + raise ValueError( + f"SymbolicPattern.kind must be 'smarts' or 'smirks'; got {self.kind!r}" + ) + if self.atom_maps is not None: + object.__setattr__(self, "atom_maps", dict(self.atom_maps)) diff --git a/src/molrep/perception/records.py b/src/molrep/perception/records.py new file mode 100644 index 0000000..44fe7ec --- /dev/null +++ b/src/molrep/perception/records.py @@ -0,0 +1,40 @@ +"""DiscreteClassRecord — condensed type + optional symbolic pattern binding.""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass + +from molrep.condensation.classes import InteractionClass + +__all__ = ["DiscreteClassRecord"] + + +@dataclass(frozen=True) +class DiscreteClassRecord: + """One discrete chemical class with optional SMARTS/SMIRKS text. + + Produced by :meth:`SymbolicForceField.records` for human inspection and + export. ``smarts`` / ``smirks`` are ``None`` until a pattern is bound in + the registry. + + Attributes: + interaction: Interaction class this type belongs to. + type_id: Discrete id within the interaction's :class:`TypeSystem`. + prototype: Named parameter prototype values. + smarts: Optional SMARTS pattern string when bound. + smirks: Optional SMIRKS string for parameter-bearing transforms. + label: Optional human-readable type label. + """ + + interaction: InteractionClass + type_id: int + prototype: Mapping[str, float | tuple[float, ...]] + smarts: str | None = None + smirks: str | None = None + label: str | None = None + + def __post_init__(self) -> None: + if self.type_id < 0: + raise ValueError(f"type_id must be non-negative, got {self.type_id}") + object.__setattr__(self, "prototype", dict(self.prototype)) diff --git a/src/molrep/perception/registry.py b/src/molrep/perception/registry.py new file mode 100644 index 0000000..092ce05 --- /dev/null +++ b/src/molrep/perception/registry.py @@ -0,0 +1,147 @@ +"""ClassPatternRegistry — bind discrete types to SymbolicPattern.""" + +from __future__ import annotations + +from collections.abc import Iterator + +from molrep.condensation.classes import InteractionClass +from molrep.perception.patterns import SymbolicPattern + +__all__ = ["ClassPatternRegistry"] + +_Key = tuple[InteractionClass, int] + + +class ClassPatternRegistry: + """Map ``(InteractionClass, type_id) → SymbolicPattern`` with reverse lookup. + + Binding is idempotent for the exact same pattern. Binding a *different* + pattern to an already-bound key, or reusing a pattern string for a different + key, raises ``ValueError``. + + Examples: + >>> from molrep.condensation import InteractionClass + >>> from molrep.perception import ClassPatternRegistry, SymbolicPattern + >>> reg = ClassPatternRegistry() + >>> reg.bind(InteractionClass.BOND, 0, SymbolicPattern("[#6]-[#6]", 2)) + >>> reg.get(InteractionClass.BOND, 0).pattern + '[#6]-[#6]' + """ + + def __init__(self) -> None: + self._by_key: dict[_Key, SymbolicPattern] = {} + self._by_pattern: dict[str, _Key] = {} + + def bind( + self, + interaction: InteractionClass, + type_id: int, + pattern: SymbolicPattern, + ) -> None: + """Bind a pattern to ``(interaction, type_id)``. + + Args: + interaction: Interaction class. + type_id: Discrete type id. + pattern: Symbolic pattern to store. + + Raises: + ValueError: Conflicting bind (different pattern for same key, or + same pattern string already bound to another key); negative + ``type_id``. + TypeError: ``pattern`` is not a :class:`SymbolicPattern`. + """ + if type_id < 0: + raise ValueError(f"type_id must be non-negative, got {type_id}") + if not isinstance(pattern, SymbolicPattern): + raise TypeError(f"pattern must be SymbolicPattern, got {type(pattern)!r}") + + key: _Key = (interaction, type_id) + existing = self._by_key.get(key) + if existing is not None: + if existing == pattern: + return # idempotent + raise ValueError( + f"conflicting bind for ({interaction.value}, type_id={type_id}): " + f"already bound to {existing.pattern!r}, " + f"refusing {pattern.pattern!r}" + ) + + owner = self._by_pattern.get(pattern.pattern) + if owner is not None and owner != key: + oi, ot = owner + raise ValueError( + f"pattern {pattern.pattern!r} already bound to " + f"({oi.value}, type_id={ot}); cannot rebind to " + f"({interaction.value}, type_id={type_id})" + ) + + self._by_key[key] = pattern + self._by_pattern[pattern.pattern] = key + + def get(self, interaction: InteractionClass, type_id: int) -> SymbolicPattern: + """Return the bound pattern. + + Args: + interaction: Interaction class. + type_id: Discrete type id. + + Returns: + Bound :class:`SymbolicPattern`. + + Raises: + KeyError: No pattern bound for this key. + """ + key: _Key = (interaction, type_id) + try: + return self._by_key[key] + except KeyError as exc: + raise KeyError( + f"no pattern bound for ({interaction.value}, type_id={type_id})" + ) from exc + + def get_optional(self, interaction: InteractionClass, type_id: int) -> SymbolicPattern | None: + """Return bound pattern or ``None`` if unbound.""" + return self._by_key.get((interaction, type_id)) + + def reverse_lookup(self, pattern: str | SymbolicPattern) -> tuple[InteractionClass, int]: + """Look up ``(interaction, type_id)`` for a pattern string. + + Args: + pattern: Pattern string or :class:`SymbolicPattern`. + + Returns: + Bound ``(InteractionClass, type_id)``. + + Raises: + KeyError: Pattern string is not registered. + """ + key_str = pattern.pattern if isinstance(pattern, SymbolicPattern) else pattern + try: + return self._by_pattern[key_str] + except KeyError as exc: + raise KeyError(f"no type bound for pattern {key_str!r}") from exc + + def items(self) -> Iterator[tuple[_Key, SymbolicPattern]]: + """Iterate ``((interaction, type_id), pattern)`` pairs.""" + return iter(self._by_key.items()) + + def __contains__(self, key: object) -> bool: + if ( + isinstance(key, tuple) + and len(key) == 2 + and isinstance(key[0], InteractionClass) + and isinstance(key[1], int) + ): + return key in self._by_key + if isinstance(key, str): + return key in self._by_pattern + if isinstance(key, SymbolicPattern): + return key.pattern in self._by_pattern + return False + + def __len__(self) -> int: + return len(self._by_key) + + def __repr__(self) -> str: + return f"ClassPatternRegistry(n_bindings={len(self._by_key)})" diff --git a/src/molrep/readout/__init__.py b/src/molrep/readout/__init__.py index ebc9934..206e328 100644 --- a/src/molrep/readout/__init__.py +++ b/src/molrep/readout/__init__.py @@ -4,8 +4,8 @@ features. Physical quantity heads (energy, force, stress) are in molpot. """ +from molrep.readout.mace import ProductHead, ProductHeadSpec from molrep.readout.pooling import masked_mean_pooling, masked_sum_pooling -from molrep.readout.product import ProductHead, ProductHeadSpec from molrep.readout.projection import BasisProjection, BasisProjectionSpec __all__ = [ @@ -16,4 +16,8 @@ "ProductHead", "ProductHeadSpec", ] -from .scalar import NonLinearBiasReadout # noqa: E402,F401 +from .mace import ( # noqa: E402,F401 + LinearReadout, + NonLinearBiasReadout, + NonLinearReadout, +) diff --git a/src/molrep/readout/mace.py b/src/molrep/readout/mace.py new file mode 100644 index 0000000..a13d4b1 --- /dev/null +++ b/src/molrep/readout/mace.py @@ -0,0 +1,350 @@ +"""MACE-only readout heads. + +Two readout families, both used exclusively by the MACE encoders: + +**Scalar readouts** (relocated from :mod:`molrep.readout.scalar`) map per-atom +equivariant features to a scalar (energy): + +- :class:`LinearReadout` — MACE's ``LinearReadoutBlock``: one ``cuet.Linear`` + onto ``1x0e``, used on every interaction layer except the last. +- :class:`NonLinearReadout` — bias-free ``Linear → c·SiLU → Linear`` + (``NonLinearReadoutBlock``); the last-layer readout of MACE-MP / MatPES. +- :class:`NonLinearBiasReadout` — the OMOL variant + ``Linear → SiLU → o3.Linear(+bias) → SiLU → o3.Linear(+bias)``. The first + linear is an equivariant ``cuet.Linear`` projecting to ``MLP_irreps`` scalars; + the two subsequent biased linears act on scalars only and reproduce e3nn's + ``o3.Linear`` normalisation (``out = (x @ W.reshape(in,out)) / sqrt(in) + b``). + The SiLU activations carry e3nn's ``normalize2mom`` scaling. + +**Product head** (relocated from :mod:`molrep.readout.product`): +:class:`ProductHead` combines symmetric basis contraction + projection + linear +readout into a single-responsibility prediction head. + +Sub-layer names mirror MACE (``linear_1``, ``linear_mid``, ``linear_2``) so the +official weights transfer by direct copy. + +Reference: + Batatia et al. "MACE: Higher Order Equivariant Message Passing Neural + Networks for Fast and Accurate Force Fields" NeurIPS 2022. + https://arxiv.org/abs/2206.07697 +""" + +from __future__ import annotations + +import math + +import cuequivariance as cue +import cuequivariance_torch as cuet +import torch +import torch.nn as nn +import torch.nn.functional as F +from pydantic import BaseModel, ConfigDict, Field + +from molix import config +from molrep.embedding.mlp import normalize2mom +from molrep.interaction.contraction import SymmetricContraction +from molrep.interaction.product import irreps_from_l_max +from molrep.readout.projection import BasisProjection + +Key = str | tuple[str, ...] + + +class _ScalarO3Linear(nn.Module): + """Scalar-only e3nn ``o3.Linear`` with bias: ``(x @ W/sqrt(in)) + b``. + + Initialisation follows e3nn's ``o3.Linear`` convention: the weight is drawn + from the **global** RNG as standard normal ``N(0, 1)`` and the path + normalisation ``1/sqrt(in_mul)`` is applied in :meth:`forward` + (``self._alpha``) rather than folded into the init, so the stored weight + stays unit-variance. The bias is zero-initialised. + + A zero weight (the previous init) makes the layer emit a constant per-atom + energy on an untrained model, which silently zeroes the forces and voids + every downstream force / parity assertion. Loading a checkpoint overwrites + the random init exactly, so weight-transfer paths are unaffected. + """ + + def __init__(self, in_mul: int, out_mul: int) -> None: + super().__init__() + self.in_mul, self.out_mul = in_mul, out_mul + self.weight = nn.Parameter(torch.randn(in_mul * out_mul, dtype=config.ftype)) + self.bias = nn.Parameter(torch.zeros(out_mul, dtype=config.ftype)) + self._alpha = 1.0 / math.sqrt(in_mul) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + w = self.weight.reshape(self.in_mul, self.out_mul) + return self._alpha * (x @ w) + self.bias + + +class NonLinearBiasReadout(nn.Module): + """Gated non-linear scalar readout to per-atom energy (single head). + + Args: + irreps_in: Input node feature irreps (last-layer product output). + mlp_dim: Hidden scalar width (``MLP_irreps`` count), e.g. 16. + """ + + def __init__(self, *, irreps_in: str, mlp_dim: int = 16) -> None: + super().__init__() + self.linear_1 = cuet.Linear( + cue.Irreps("O3", irreps_in), + cue.Irreps("O3", f"{mlp_dim}x0e"), + layout=cue.ir_mul, + dtype=config.ftype, + ) + self._act_cst = normalize2mom(F.silu) + self.linear_mid = _ScalarO3Linear(mlp_dim, mlp_dim) + self.linear_2 = _ScalarO3Linear(mlp_dim, 1) + + def _act(self, x: torch.Tensor) -> torch.Tensor: + return F.silu(x) * self._act_cst + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Compute per-atom scalar energy. + + Args: + x: Node features ``(N, irreps_in.dim)``. + + Returns: + Per-atom energy ``(N, 1)``. + """ + x = self._act(self.linear_1(x)) + x = self._act(self.linear_mid(x)) + return self.linear_2(x) + + +class LinearReadout(nn.Module): + """Equivariant linear projection of node features to a per-atom scalar. + + MACE's ``LinearReadoutBlock``: one ``cuet.Linear`` onto ``1x0e``. Used for + every interaction layer except the last, where the non-linear readout + (:class:`NonLinearReadout`) takes over. + + Args: + irreps_in: Input node feature irreps, e.g. ``"128x0e+128x1o"``. + + Reference: + Batatia et al. "MACE: Higher Order Equivariant Message Passing Neural + Networks for Fast and Accurate Force Fields" NeurIPS 2022. + https://arxiv.org/abs/2206.07697 + """ + + def __init__(self, *, irreps_in: str) -> None: + super().__init__() + self.linear = cuet.Linear( + cue.Irreps("O3", irreps_in), + cue.Irreps("O3", "1x0e"), + layout=cue.ir_mul, + dtype=config.ftype, + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Project node features to a per-atom scalar. + + Args: + x: Node features ``(N, irreps_in.dim)``. + + Returns: + Per-atom scalar ``(N, 1)``. + """ + return self.linear(x) + + +class NonLinearReadout(nn.Module): + """Bias-free gated non-linear scalar readout (MACE ``NonLinearReadoutBlock``). + + ``Linear → c·SiLU → Linear``, both equivariant ``cuet.Linear`` layers with no + bias; ``c`` is e3nn's moment normalisation constant. This is the readout on + the **last** interaction layer of the MACE-MP / MatPES foundation models. + + Distinct from :class:`NonLinearBiasReadout`, which is the OMOL variant with + biases and an extra middle layer — do not substitute one for the other, the + weight layouts differ. + + Args: + irreps_in: Input node feature irreps (last-layer product output). + mlp_dim: Hidden scalar width (MACE's ``MLP_irreps``), e.g. 16. + + Reference: + Batatia et al. "MACE: Higher Order Equivariant Message Passing Neural + Networks for Fast and Accurate Force Fields" NeurIPS 2022. + https://arxiv.org/abs/2206.07697 + """ + + def __init__(self, *, irreps_in: str, mlp_dim: int = 16) -> None: + super().__init__() + self.linear_1 = cuet.Linear( + cue.Irreps("O3", irreps_in), + cue.Irreps("O3", f"{mlp_dim}x0e"), + layout=cue.ir_mul, + dtype=config.ftype, + ) + self.linear_2 = cuet.Linear( + cue.Irreps("O3", f"{mlp_dim}x0e"), + cue.Irreps("O3", "1x0e"), + layout=cue.ir_mul, + dtype=config.ftype, + ) + self._act_cst = normalize2mom(F.silu) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Compute per-atom scalar energy. + + Args: + x: Node features ``(N, irreps_in.dim)``. + + Returns: + Per-atom energy ``(N, 1)``. + """ + return self.linear_2(F.silu(self.linear_1(x)) * self._act_cst) + + +class ProductHeadSpec(BaseModel): + """Configuration for product prediction head. + + Combines multi-body basis construction (via SymmetricContraction), + optional basis projection, and linear readout to scalars. + + Attributes: + hidden_dim: Dimension of input node features. + out_dim: Dimension of output predictions (1 for scalar energy). + num_radial: Number of radial basis functions. + l_max: Maximum angular momentum. + max_body_order: Maximum body order for multi-body expansion. + num_species: Number of atomic species. + """ + + model_config = ConfigDict(arbitrary_types_allowed=True) + + hidden_dim: int = Field(..., gt=0) + out_dim: int = Field(..., gt=0) + num_radial: int = Field(8, gt=0) + l_max: int = Field(2, ge=0) + max_body_order: int = Field(2, ge=1, le=3) + num_species: int = Field(118, gt=0) + use_fallback: bool = True + + +class ProductHead(nn.Module): + """Product layer head for multi-body-aware scalar predictions. + + Single-responsibility module that: + 1. Constructs symmetric multi-body basis (SymmetricContraction) + 2. Projects basis features (BasisProjection) + 3. Applies linear transformation to output dimension + + Does NOT apply pooling - that is the responsibility of a separate + pooling module. Returns node-level predictions only. + + Architecture: + node_features (n_nodes, hidden_dim) + atom_types (n_nodes,) + ↓ + [SymmetricContraction] + ↓ + basis (n_nodes, hidden_dim) + ↓ + [BasisProjection] + ↓ + features (n_nodes, hidden_dim) + ↓ + [Linear(hidden_dim → out_dim)] + ↓ + predictions (n_nodes, out_dim) + """ + + def __init__( + self, + *, + hidden_dim: int, + out_dim: int, + num_radial: int = 8, + l_max: int = 2, + max_body_order: int = 2, + num_species: int = 118, + use_fallback: bool = True, + ): + """Initialize product head. + + Args: + hidden_dim: Dimension of node features. + out_dim: Dimension of output predictions. + num_radial: Number of radial basis functions. + l_max: Maximum angular momentum. + max_body_order: Maximum body order (1-3). + num_species: Number of atomic species. + use_fallback: Pure-torch cuEq path for the symmetric contraction + (default ``True``, functorch-safe). Set ``False`` for the + fused kernels when forces use the autograd backend. + """ + super().__init__() + + self.config = ProductHeadSpec( + hidden_dim=hidden_dim, + out_dim=out_dim, + num_radial=num_radial, + l_max=l_max, + max_body_order=max_body_order, + num_species=num_species, + use_fallback=use_fallback, + ) + + # ``hidden_dim`` is the *full* mixed-l feature dim emitted by the + # interaction block (e.g. 144 = 16x0e+16x1o+16x2e for l_max=2). Recover + # the per-l multiplicity (scalar channel count) so the contraction can + # be told the *real* irreps. Declaring this mixed-l tensor as pure + # scalars is the rotation-invariance bug this head exists to avoid. + per_channel_dim = (l_max + 1) ** 2 + if hidden_dim % per_channel_dim != 0: + raise ValueError( + f"hidden_dim={hidden_dim} is not a multiple of (l_max+1)^2=" + f"{per_channel_dim}; cannot infer the mixed-l irreps multiplicity." + ) + num_features = hidden_dim // per_channel_dim + irreps_in = irreps_from_l_max(l_max, num_features) + irreps_out = f"{num_features}x0e" # invariant scalar output + + # Single-responsibility sub-modules + self.symmetric_contraction = SymmetricContraction( + hidden_dim=hidden_dim, + num_species=num_species, + max_body_order=max_body_order, + irreps_in=irreps_in, + irreps_out=irreps_out, + use_fallback=use_fallback, + ) + + self.basis_projection = BasisProjection( + hidden_dim=hidden_dim, + num_radial=num_radial, + l_max=l_max, + max_body_order=max_body_order, + ) + + # The contraction emits ``num_features`` invariant scalars; the readout + # linear maps those scalars (not the full mixed-l dim) to ``out_dim``. + self.linear = nn.Linear(num_features, out_dim, dtype=config.ftype) + + def forward( + self, + node_features: torch.Tensor, + atom_types: torch.Tensor, + ) -> torch.Tensor: + """Compute node-level predictions from features. + + Args: + node_features: Node features (n_nodes, hidden_dim) + atom_types: Atomic numbers (n_nodes,) + + Returns: + Predictions (n_nodes, out_dim). + """ + # Step 1: Symmetric multi-body basis via cuEquivariance + basis = self.symmetric_contraction(node_features, atom_types) + + # Step 2: Project basis features (currently passthrough with cuEquivariance) + features = self.basis_projection(basis) + + # Step 3: Linear transformation to output dimension + predictions = self.linear(features) + + return predictions diff --git a/src/molrep/readout/product.py b/src/molrep/readout/product.py index 474f338..e562c3a 100644 --- a/src/molrep/readout/product.py +++ b/src/molrep/readout/product.py @@ -1,162 +1,7 @@ -"""Product layer head for final scalar predictions. +"""Deprecated location — moved to :mod:`molrep.readout.mace`. -Combines symmetric basis contraction + projection + linear readout into -a single-responsibility prediction head. +Kept as a re-export shim for chain step mace-subpackage-restructure-01; +removed in 06-wire. """ -from __future__ import annotations - -import torch -import torch.nn as nn -from pydantic import BaseModel, ConfigDict, Field - -from molix import config -from molrep.interaction.contraction import SymmetricContraction -from molrep.interaction.product import irreps_from_l_max -from molrep.readout.projection import BasisProjection - -Key = str | tuple[str, ...] - - -class ProductHeadSpec(BaseModel): - """Configuration for product prediction head. - - Combines multi-body basis construction (via SymmetricContraction), - optional basis projection, and linear readout to scalars. - - Attributes: - hidden_dim: Dimension of input node features. - out_dim: Dimension of output predictions (1 for scalar energy). - num_radial: Number of radial basis functions. - l_max: Maximum angular momentum. - max_body_order: Maximum body order for multi-body expansion. - num_species: Number of atomic species. - """ - - model_config = ConfigDict(arbitrary_types_allowed=True) - - hidden_dim: int = Field(..., gt=0) - out_dim: int = Field(..., gt=0) - num_radial: int = Field(8, gt=0) - l_max: int = Field(2, ge=0) - max_body_order: int = Field(2, ge=1, le=3) - num_species: int = Field(118, gt=0) - - -class ProductHead(nn.Module): - """Product layer head for multi-body-aware scalar predictions. - - Single-responsibility module that: - 1. Constructs symmetric multi-body basis (SymmetricContraction) - 2. Projects basis features (BasisProjection) - 3. Applies linear transformation to output dimension - - Does NOT apply pooling - that is the responsibility of a separate - pooling module. Returns node-level predictions only. - - Architecture: - node_features (n_nodes, hidden_dim) + atom_types (n_nodes,) - ↓ - [SymmetricContraction] - ↓ - basis (n_nodes, hidden_dim) - ↓ - [BasisProjection] - ↓ - features (n_nodes, hidden_dim) - ↓ - [Linear(hidden_dim → out_dim)] - ↓ - predictions (n_nodes, out_dim) - """ - - def __init__( - self, - *, - hidden_dim: int, - out_dim: int, - num_radial: int = 8, - l_max: int = 2, - max_body_order: int = 2, - num_species: int = 118, - ): - """Initialize product head. - - Args: - hidden_dim: Dimension of node features. - out_dim: Dimension of output predictions. - num_radial: Number of radial basis functions. - l_max: Maximum angular momentum. - max_body_order: Maximum body order (1-3). - num_species: Number of atomic species. - """ - super().__init__() - - self.config = ProductHeadSpec( - hidden_dim=hidden_dim, - out_dim=out_dim, - num_radial=num_radial, - l_max=l_max, - max_body_order=max_body_order, - num_species=num_species, - ) - - # ``hidden_dim`` is the *full* mixed-l feature dim emitted by the - # interaction block (e.g. 144 = 16x0e+16x1o+16x2e for l_max=2). Recover - # the per-l multiplicity (scalar channel count) so the contraction can - # be told the *real* irreps. Declaring this mixed-l tensor as pure - # scalars is the rotation-invariance bug this head exists to avoid. - per_channel_dim = (l_max + 1) ** 2 - if hidden_dim % per_channel_dim != 0: - raise ValueError( - f"hidden_dim={hidden_dim} is not a multiple of (l_max+1)^2=" - f"{per_channel_dim}; cannot infer the mixed-l irreps multiplicity." - ) - num_features = hidden_dim // per_channel_dim - irreps_in = irreps_from_l_max(l_max, num_features) - irreps_out = f"{num_features}x0e" # invariant scalar output - - # Single-responsibility sub-modules - self.symmetric_contraction = SymmetricContraction( - hidden_dim=hidden_dim, - num_species=num_species, - max_body_order=max_body_order, - irreps_in=irreps_in, - irreps_out=irreps_out, - ) - - self.basis_projection = BasisProjection( - hidden_dim=hidden_dim, - num_radial=num_radial, - l_max=l_max, - max_body_order=max_body_order, - ) - - # The contraction emits ``num_features`` invariant scalars; the readout - # linear maps those scalars (not the full mixed-l dim) to ``out_dim``. - self.linear = nn.Linear(num_features, out_dim, dtype=config.ftype) - - def forward( - self, - node_features: torch.Tensor, - atom_types: torch.Tensor, - ) -> torch.Tensor: - """Compute node-level predictions from features. - - Args: - node_features: Node features (n_nodes, hidden_dim) - atom_types: Atomic numbers (n_nodes,) - - Returns: - Predictions (n_nodes, out_dim). - """ - # Step 1: Symmetric multi-body basis via cuEquivariance - basis = self.symmetric_contraction(node_features, atom_types) - - # Step 2: Project basis features (currently passthrough with cuEquivariance) - features = self.basis_projection(basis) - - # Step 3: Linear transformation to output dimension - predictions = self.linear(features) - - return predictions +from molrep.readout.mace import ProductHead, ProductHeadSpec # noqa: F401 diff --git a/src/molrep/readout/scalar.py b/src/molrep/readout/scalar.py index f93d46a..5231e6b 100644 --- a/src/molrep/readout/scalar.py +++ b/src/molrep/readout/scalar.py @@ -1,87 +1,12 @@ -"""Non-linear scalar readout head (MACE ``NonLinearBiasReadoutBlock``). +"""Deprecated location — moved to :mod:`molrep.readout.mace`. -Maps per-atom equivariant features to a scalar (energy) via: -``Linear → SiLU → o3.Linear(+bias) → SiLU → o3.Linear(+bias)``. The first -linear is an equivariant ``cuet.Linear`` projecting to ``MLP_irreps`` scalars; -the two subsequent biased linears act on scalars only and reproduce e3nn's -``o3.Linear`` normalisation (``out = (x @ W.reshape(in,out)) / sqrt(in) + b``). -The SiLU activations carry e3nn's ``normalize2mom`` scaling. - -Sub-layer names mirror MACE (``linear_1``, ``linear_mid``, ``linear_2``) so the -official weights transfer by direct copy. - -Reference: - Batatia et al. "MACE: Higher Order Equivariant Message Passing Neural - Networks for Fast and Accurate Force Fields" NeurIPS 2022. - https://arxiv.org/abs/2206.07697 +Kept as a re-export shim for chain step mace-subpackage-restructure-01; +removed in 06-wire. """ -from __future__ import annotations - -import math - -import cuequivariance as cue -import cuequivariance_torch as cuet -import torch -import torch.nn as nn -import torch.nn.functional as F - -from molix import config - - -def _normalize2mom(fn) -> float: - gen = torch.Generator(device="cpu").manual_seed(0) - z = torch.randn(1_000_000, generator=gen, dtype=torch.float64) - return fn(z).pow(2).mean().pow(-0.5).item() - - -class _ScalarO3Linear(nn.Module): - """Scalar-only e3nn ``o3.Linear`` with bias: ``(x @ W/sqrt(in)) + b``.""" - - def __init__(self, in_mul: int, out_mul: int) -> None: - super().__init__() - self.in_mul, self.out_mul = in_mul, out_mul - self.weight = nn.Parameter(torch.zeros(in_mul * out_mul, dtype=config.ftype)) - self.bias = nn.Parameter(torch.zeros(out_mul, dtype=config.ftype)) - self._alpha = 1.0 / math.sqrt(in_mul) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - w = self.weight.reshape(self.in_mul, self.out_mul) - return self._alpha * (x @ w) + self.bias - - -class NonLinearBiasReadout(nn.Module): - """Gated non-linear scalar readout to per-atom energy (single head). - - Args: - irreps_in: Input node feature irreps (last-layer product output). - mlp_dim: Hidden scalar width (``MLP_irreps`` count), e.g. 16. - """ - - def __init__(self, *, irreps_in: str, mlp_dim: int = 16) -> None: - super().__init__() - self.linear_1 = cuet.Linear( - cue.Irreps("O3", irreps_in), - cue.Irreps("O3", f"{mlp_dim}x0e"), - layout=cue.ir_mul, - dtype=config.ftype, - ) - self._act_cst = _normalize2mom(F.silu) - self.linear_mid = _ScalarO3Linear(mlp_dim, mlp_dim) - self.linear_2 = _ScalarO3Linear(mlp_dim, 1) - - def _act(self, x: torch.Tensor) -> torch.Tensor: - return F.silu(x) * self._act_cst - - def forward(self, x: torch.Tensor) -> torch.Tensor: - """Compute per-atom scalar energy. - - Args: - x: Node features ``(N, irreps_in.dim)``. - - Returns: - Per-atom energy ``(N, 1)``. - """ - x = self._act(self.linear_1(x)) - x = self._act(self.linear_mid(x)) - return self.linear_2(x) +from molrep.readout.mace import ( # noqa: F401 + LinearReadout, + NonLinearBiasReadout, + NonLinearReadout, + _ScalarO3Linear, +) diff --git a/src/molrep/utils/geometry.py b/src/molrep/utils/geometry.py index 093e3f9..c428ac4 100644 --- a/src/molrep/utils/geometry.py +++ b/src/molrep/utils/geometry.py @@ -5,6 +5,8 @@ import torch import torch.nn as nn +from molix import config + class NeighborGraphBuilder: """Build neighbor graphs from atomic positions. @@ -132,8 +134,8 @@ class GaussianRBF(nn.Module): def __init__(self, num_rbf: int = 50, cutoff: float = 5.0, trainable: bool = False): super().__init__() - centers = torch.linspace(0, cutoff, num_rbf) - widths = torch.full((num_rbf,), cutoff / num_rbf) + centers = torch.linspace(0, cutoff, num_rbf, dtype=config.ftype) + widths = torch.full((num_rbf,), cutoff / num_rbf, dtype=config.ftype) if trainable: self.centers = nn.Parameter(centers) diff --git a/src/molzoo/README.md b/src/molzoo/README.md index 9bfc17c..73b70fc 100644 --- a/src/molzoo/README.md +++ b/src/molzoo/README.md @@ -15,6 +15,31 @@ PiNet is a **package** (`molzoo/pinet/`), not a single file: Public imports stay stable: `from molzoo.pinet import PiNet, PiNetPotential`. The Sonata model lives in `molpot.composition`, not here. +MACE is a **package** too (`molzoo/mace/`), since +mace-subpackage-restructure-02-core: + +| Module | Role | +|--------|------| +| `spec` | `MACESpec` base + `MACEMatpesSpec` / `MACEOMolSpec` (torch-free) | +| `geometry` | `edge_vectors` / `edge_lengths` — additive PBC shifts `S = n·h` | +| `encoder` | `MACEEncoder` — one configuration-driven foundation backbone | +| `potential` | `MACEPotential` — per-graph energy `(B,)` in eV + per-atom forces `(N, 3)` in eV/Å | +| `checkpoint` | `CheckpointRemap` + the `MATPES_REMAP` / `OMOL_REMAP` presets — official-weight key remap | +| `variants` | `MACEMatpes` / `MACEOMol` named foundation models + their `load_*_state_dict` aliases | +| `research` | the freely-configurable research `MACE` encoder (former `molzoo/mace.py`) | + +`MACESpec` changed meaning with that move: it is now the **shared +foundation-variant** configuration base in `molzoo.mace.spec` — the thing +`MACEMatpesSpec` / `MACEOMolSpec` derive from and `MACEPotential` is built +out of. The research encoder is configured by keyword and keeps its own, +unrelated `MACEResearchSpec` (`molzoo.mace.research`). + +The flat modules `molzoo/mace_matpes.py` and `molzoo/mace_omol.py` were +deleted in mace-subpackage-restructure-06-wire; the **top-level** names are +unchanged, so `from molzoo import MACE, MACEMatpes, MACEOMol, +load_matpes_state_dict, load_omol_state_dict` keeps working (lazily — see +`molzoo/__init__.py`), as does `from molzoo.mace import MACE`. + ## Model Specifications Each model in this package ships with **one** spec artifact in `specs/`: @@ -46,38 +71,57 @@ One skill + one agent keep `.md` aligned with code and paper: ## Input Conventions -Both encoders accept keyword tensors: +The encoders take a post-collate batch `TensorDict` and read: -- `Z`: Atomic numbers `(N,)` -- `edge_dist`: Edge distances `(E,)` -- `edge_diff`: Edge vectors `(E, 3)` -- `edge_index`: Edge indices `(E, 2)` +- `atoms.Z`: Atomic numbers `(N,)` +- `atoms.pos`: Positions `(N, 3)` +- `edges.edge_index`: Edge indices `(E, 2)` — `[:,0]` source, `[:,1]` target +- `edges.edge_diff`: Edge vectors `(E, 3)` — `pos[target] - pos[source]` +- `edges.edge_dist`: Edge distances `(E,)` -Output: `(N, num_layers, feature_dim)` — per-atom, per-layer features. +`MACE` writes `atoms.node_features` `(N, num_layers, feature_dim)` back into +the same batch. ## Usage +Research encoder — configured by keyword, features per layer: + ```python -import torch -from molzoo import MACE, MACESpec +from molzoo import MACE from molrep.embedding.node import DiscreteEmbeddingSpec -from molpot import LayerPooling, PotentialComposer, LJParameterHead, LJ126 -encoder = MACE(MACESpec( +encoder = MACE( node_attr_specs=[DiscreteEmbeddingSpec(input_key="Z", num_classes=119, emb_dim=64)], num_elements=119, num_features=64, r_max=5.0, -)) - -Z = torch.randint(0, 10, (20,)) -features = encoder( - Z=Z, - edge_dist=torch.rand(80), - edge_diff=torch.randn(80, 3), - edge_index=torch.randint(0, 20, (80, 2)), ) +batch = encoder(batch) # writes atoms.node_features +features = batch["atoms", "node_features"] # (N, num_layers, 64) +``` + +Foundation backbone — configured by a validated spec, composed by the caller: -pool = LayerPooling("mean") -node_features = pool(features) # (20, 64) +```python +from molzoo.mace.encoder import MACEEncoder +from molzoo.mace.geometry import edge_lengths, edge_vectors +from molzoo.mace.spec import MACEMatpesSpec + +spec = MACEMatpesSpec(atomic_numbers=[1, 6, 8], atomic_energies=[-13.6, -1029.0, -2041.0]) +encoder = MACEEncoder(spec) + +vectors = edge_vectors(pos, edge_index) # optional shifts=... for PBC +node_attrs = encoder.node_attrs(Z, pos.dtype) +edge_feats, cutoff = encoder.radial_features(edge_lengths(vectors), Z, edge_index) +per_layer = encoder.layer_features( + node_feats=encoder.initial_node_features(node_attrs), + node_attrs=node_attrs, + edge_attrs=encoder.angular_features(vectors), + edge_feats=edge_feats, + edge_index=edge_index, + cutoff=cutoff, +) ``` + +Swap `MACEMatpesSpec` for `MACEOMolSpec` to build the charge/spin-conditioned +OMOL stack from the same backbone. diff --git a/src/molzoo/__init__.py b/src/molzoo/__init__.py index fb29292..e5509aa 100644 --- a/src/molzoo/__init__.py +++ b/src/molzoo/__init__.py @@ -1,36 +1,115 @@ -"""MolZoo: molecular model zoo. +"""MolZoo: molecular model zoo — encoder architectures and foundation models. -This package provides encoder architectures and potential models. +An *encoder* here maps a batch of atoms (their chemical element numbers and +Cartesian positions) to a feature vector per atom, which a downstream head +turns into an energy or another property; a *foundation* model is one shipped +with weights already fitted on a large, chemically broad dataset, used +unchanged rather than trained by the caller. + +**Every symbol in :data:`__all__` is exported lazily** — this module imports no +model module at import time, and a name is resolved only when someone asks for +it, through the module-level ``__getattr__`` hook of PEP 562. The reason is +cost: resolving one of them pulls in the cuEquivariance stack (NVIDIA's GPU +library for the equivariant tensor algebra MACE and Allegro are built from), +seconds of import time and a large chunk of memory that a consumer reading a +config, a dataset or a checkpoint manifest should not have to pay. The one name +that would be cheap, ``MACESpec`` (its module, :mod:`molzoo.mace.spec`, imports +nothing but ``typing`` and pydantic), is in the table anyway: a single rule with +no exception list is what keeps this file and ``molzoo/mace/__init__.py`` +honest. + +Direct sub-module imports are unaffected and stay ordinary imports:: + + import molzoo.mace # torch-free: the configurations only + from molzoo.mace import MACE # pulls the equivariance stack + from molzoo import MACEMatpes # the same, through the lazy table """ from typing import TYPE_CHECKING -from molzoo.allegro import Allegro, AllegroSpec -from molzoo.mace import MACE, MACESpec -from molzoo.pinet import PiNet, PiNetSpec - if TYPE_CHECKING: - from molzoo.mace_omol import MACEOMol, load_omol_state_dict + from molzoo.allegro import Allegro, AllegroSpec + from molzoo.chem import ChemPerception, ChemPerceptionSpec + from molzoo.mace.research import MACE + from molzoo.mace.spec import MACESpec + from molzoo.mace.variants import ( + MACEMatpes, + MACEOMol, + load_matpes_state_dict, + load_omol_state_dict, + ) + from molzoo.pinet import PiNet, PiNetSpec -# Lazily exported (PEP 562): MACEOMol pulls in the full MACE-OMOL cuEquivariance -# stack + weight-conversion helpers, which not every molzoo consumer needs. -_LAZY = {"MACEOMol", "load_omol_state_dict"} +#: Public name → the module the lazy import goes to. The one lazy-export +#: table: every entry is also in :data:`__all__`, and nothing outside it +#: resolves. For the MACE names the target is the **defining** module +#: (``molzoo.mace.research`` / ``.spec`` / ``.variants``) rather than the +#: ``molzoo.mace`` re-export surface, so resolving one here does not depend on +#: that package's own lazy table as well. The two PiNet names go to +#: ``molzoo.pinet``, whose ``__init__`` re-exports them eagerly and is +#: therefore already the defining surface for an importer. +_LAZY = { + "Allegro": "molzoo.allegro", + "AllegroSpec": "molzoo.allegro", + "ChemPerception": "molzoo.chem", + "ChemPerceptionSpec": "molzoo.chem", + "MACE": "molzoo.mace.research", + "MACESpec": "molzoo.mace.spec", + "MACEMatpes": "molzoo.mace.variants", + "MACEOMol": "molzoo.mace.variants", + "PiNet": "molzoo.pinet", + "PiNetSpec": "molzoo.pinet", + "load_matpes_state_dict": "molzoo.mace.variants", + "load_omol_state_dict": "molzoo.mace.variants", +} __all__ = [ "Allegro", "AllegroSpec", + "ChemPerception", + "ChemPerceptionSpec", "MACE", - "MACESpec", + "MACEMatpes", "MACEOMol", + "MACESpec", "PiNet", "PiNetSpec", + "load_matpes_state_dict", "load_omol_state_dict", ] def __getattr__(name: str): - if name in _LAZY: - from molzoo import mace_omol + """Import a lazily exported name on first attribute access (PEP 562). + + Args: + name: Attribute requested on the ``molzoo`` module. + + Returns: + The object named ``name``, imported from its module in :data:`_LAZY`. + + Raises: + AttributeError: If ``name`` is not a lazily exported symbol — a typo + must fail the way a missing module attribute does, not as an + ``ImportError`` from somewhere inside the package. + """ + module_name = _LAZY.get(name) + if module_name is None: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + import importlib + + return getattr(importlib.import_module(module_name), name) + + +def __dir__() -> list[str]: + """Report the public surface, so completion survives the lazy table. + + Without this, ``dir(molzoo)`` lists the module globals — which, under the + lazy policy, is every name *except* the models. It returns exactly + :data:`__all__`, the same set ``from molzoo import *`` binds, so the two + cannot drift apart. - return getattr(mace_omol, name) - raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + Returns: + The names in :data:`__all__`, sorted. + """ + return sorted(__all__) diff --git a/src/molzoo/chem/__init__.py b/src/molzoo/chem/__init__.py new file mode 100644 index 0000000..0aadab3 --- /dev/null +++ b/src/molzoo/chem/__init__.py @@ -0,0 +1,14 @@ +"""molzoo.chem — continuous chemical perception recipe. + +Public surface:: + + from molzoo.chem import ChemPerception, ChemPerceptionSpec +""" + +from .encoder import ChemPerception +from .spec import ChemPerceptionSpec + +__all__ = [ + "ChemPerception", + "ChemPerceptionSpec", +] diff --git a/src/molzoo/chem/encoder.py b/src/molzoo/chem/encoder.py new file mode 100644 index 0000000..12efaf7 --- /dev/null +++ b/src/molzoo/chem/encoder.py @@ -0,0 +1,81 @@ +"""ChemPerception — thin molzoo recipe over :class:`~molrep.chem.encoder.ChemEncoder`. + +Encoder-only: writes continuous chem features into the batch TensorDict. +No energy head, no molpot import. +""" + +from __future__ import annotations + +from tensordict import TensorDict + +from molrep.chem.encoder import ChemEncoder +from molrep.chem.features import ChemEmbeddings + +from .spec import ChemPerceptionSpec + + +class ChemPerception(ChemEncoder): + """Continuous chemical-perception recipe (molzoo). + + Thin config-driven wrapper around :class:`~molrep.chem.encoder.ChemEncoder`. + Writes ``*.chem_features``; does not predict energy or forces. + + Args: + spec: Optional :class:`ChemPerceptionSpec`. When omitted, keyword + arguments populate a new spec (same fields as the parent + encoder). + **kwargs: Forwarded to :class:`ChemPerceptionSpec` when ``spec`` is + not provided. + + Example:: + + from molzoo.chem import ChemPerception, ChemPerceptionSpec + + model = ChemPerception(atom_dim=16, bond_dim=16) + batch = model(batch) # writes atoms/bonds/… chem_features + """ + + def __init__( + self, + spec: ChemPerceptionSpec | None = None, + **kwargs, + ) -> None: + if spec is None: + spec = ChemPerceptionSpec(**kwargs) + elif kwargs: + raise TypeError( + "ChemPerception accepts either a ChemPerceptionSpec or keyword overrides, not both" + ) + super().__init__( + atom_dim=spec.atom_dim, + bond_dim=spec.bond_dim, + angle_dim=spec.angle_dim, + proper_dim=spec.proper_dim, + improper_dim=spec.improper_dim, + num_elements=spec.num_elements, + hidden_dim=spec.hidden_dim, + num_bond_types=spec.num_bond_types, + ) + self.config = spec + + def forward(self, batch: TensorDict) -> TensorDict: + """Run chemical perception and write features onto ``batch``. + + Args: + batch: Nested TensorDict with ``atoms.Z`` and valence topology. + + Returns: + The same batch with ``*.chem_features`` populated. + """ + return super().forward(batch) + + def embeddings(self, batch: TensorDict) -> ChemEmbeddings: + """View written chem features as :class:`~molrep.chem.features.ChemEmbeddings`. + + Args: + batch: Batch previously processed by :meth:`forward`. + + Returns: + :class:`ChemEmbeddings` viewing the feature fields. + """ + return super().embeddings(batch) diff --git a/src/molzoo/chem/spec.py b/src/molzoo/chem/spec.py new file mode 100644 index 0000000..480fc32 --- /dev/null +++ b/src/molzoo/chem/spec.py @@ -0,0 +1,31 @@ +"""Validated configuration snapshot for :class:`~molzoo.chem.encoder.ChemPerception`.""" + +from __future__ import annotations + +from pydantic import BaseModel, ConfigDict, Field + + +class ChemPerceptionSpec(BaseModel): + """Configuration for the continuous chemical-perception recipe. + + Attributes: + atom_dim: Per-atom feature dimension. + bond_dim: Per-bond feature dimension. + angle_dim: Per-angle feature dimension. + proper_dim: Per-proper-torsion feature dimension. + improper_dim: Per-improper feature dimension. + num_elements: Atomic-number embedding table size. + hidden_dim: Optional shared MLP hidden width for context builders. + num_bond_types: Bond-type table size; ``0`` disables type conditioning. + """ + + model_config = ConfigDict(arbitrary_types_allowed=True) + + atom_dim: int = Field(default=32, gt=0) + bond_dim: int = Field(default=32, gt=0) + angle_dim: int = Field(default=32, gt=0) + proper_dim: int = Field(default=32, gt=0) + improper_dim: int = Field(default=32, gt=0) + num_elements: int = Field(default=119, gt=0) + hidden_dim: int | None = None + num_bond_types: int = Field(default=0, ge=0) diff --git a/src/molzoo/mace.py b/src/molzoo/mace.py deleted file mode 100644 index 25bbfff..0000000 --- a/src/molzoo/mace.py +++ /dev/null @@ -1,590 +0,0 @@ -"""MACE: Multi-Atomic Cluster Expansion encoder. - -Equivariant message-passing encoder that produces per-layer node features. -Downstream readout, classical potential terms, and force derivation are -handled outside this module. - -Example: - >>> from molzoo import MACE - >>> from molrep.embedding.node import DiscreteEmbeddingSpec - >>> encoder = MACE( - ... node_attr_specs=[DiscreteEmbeddingSpec( - ... input_key="Z", num_classes=119, emb_dim=64)], - ... num_elements=118, - ... num_features=128, - ... r_max=5.0, - ... ) - >>> features = encoder( - ... Z=Z, - ... edge_dist=edge_dist, - ... edge_diff=edge_diff, - ... edge_index=edge_index, - ... ) - >>> print(features.shape) # (n_nodes, num_layers, num_features) - -Reference: - Batatia et al. "MACE: Higher Order Equivariant Message Passing Neural - Networks for Fast and Accurate Force Fields" NeurIPS 2022 - https://arxiv.org/abs/2206.07697 -""" - -from __future__ import annotations - -import cuequivariance as cue -import cuequivariance_torch as cuet -import torch -import torch.nn as nn -from cuequivariance import O3, Irreps -from pydantic import BaseModel, ConfigDict, Field -from tensordict import TensorDict -from tensordict.nn import TensorDictModuleBase - -from molix import config -from molrep.embedding.angular import SphericalHarmonics -from molrep.embedding.cutoff import CosineCutoff -from molrep.embedding.node import ( - ContinuousEmbeddingSpec, - DiscreteEmbeddingSpec, - JointEmbedding, -) -from molrep.embedding.radial import BesselRBF -from molrep.interaction.element import ElementUpdate -from molrep.interaction.product import ( - ConvTP, - irreps_from_l_max, - sh_irreps_from_l_max, -) -from molrep.interaction.radial import RadialWeightMLP -from molrep.readout.product import ProductHead - -# =========================================================================== -# Embedding Block -# =========================================================================== - - -class EmbeddingSpec(BaseModel): - """Configuration for the embedding block. - - Attributes: - node_attr_specs: Embedding specifications for node attributes - (e.g. atomic number Z, charge). - num_features: Number of feature channels (scalar multiplicity at l=0). - r_max: Radial cutoff distance in Angstroms. - num_bessel: Number of Bessel radial basis functions. - l_max: Maximum angular momentum order. - """ - - model_config = ConfigDict(arbitrary_types_allowed=True) - - node_attr_specs: list[DiscreteEmbeddingSpec | ContinuousEmbeddingSpec] = Field( - ..., min_length=1 - ) - num_features: int = Field(..., gt=0) - r_max: float = Field(..., gt=0.0) - num_bessel: int = Field(8, gt=0) - l_max: int = Field(2, ge=0) - - -class EmbeddingBlock(nn.Module): - """Node and edge embedding block. - - Computes initial node features via ``JointEmbedding`` and edge features - via Bessel radial basis, spherical harmonics, and a cosine cutoff envelope. - - Attributes: - node_embedding: Joint embedding for node attributes. - radial_embedding: Bessel radial basis functions. - spherical_harmonics: Spherical harmonics for edge directions. - cutoff_fn: Cosine cutoff envelope. - """ - - def __init__( - self, - *, - node_attr_specs: list[DiscreteEmbeddingSpec | ContinuousEmbeddingSpec], - num_features: int, - r_max: float, - num_bessel: int = 8, - l_max: int = 2, - ): - """Initialize embedding block. - - Args: - node_attr_specs: Embedding specs for node attributes. - num_features: Scalar channel multiplicity (l=0 count). - r_max: Radial cutoff in Angstroms. - num_bessel: Number of Bessel basis functions. - l_max: Maximum angular momentum order. - """ - super().__init__() - - self.config = EmbeddingSpec( - node_attr_specs=node_attr_specs, - num_features=num_features, - r_max=r_max, - num_bessel=num_bessel, - l_max=l_max, - ) - - # Node embedding - self.node_embedding = JointEmbedding( - embedding_specs=node_attr_specs, - out_dim=num_features, - ) - - # Edge radial basis - self.radial_embedding = BesselRBF( - r_cut=r_max, - num_radial=num_bessel, - ) - - # Spherical harmonics - self.spherical_harmonics = SphericalHarmonics( - l_max=l_max, - ) - - # Cutoff envelope - self.cutoff_fn = CosineCutoff( - r_cut=r_max, - ) - - def forward( - self, - Z: torch.Tensor, - edge_dist: torch.Tensor, - edge_diff: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Compute initial node and edge features. - - Args: - Z: Atomic numbers (n_nodes,). - edge_dist: Bond distances (n_edges,). - edge_diff: Bond vectors (target - source) (n_edges, 3). - - Returns: - tuple of: - - node_feats: Node features (n_nodes, num_features). - - edge_attrs: Spherical harmonics (n_edges, sh_dim). - - edge_feats: Radial basis features (n_edges, num_bessel). - """ - # Node features - node_feats = self.node_embedding(Z=Z) - - # Edge direction - edge_dir = edge_diff / (edge_dist.unsqueeze(-1) + 1e-8) - - # Spherical harmonics - edge_attrs = self.spherical_harmonics(edge_dir) - - # Radial basis * cutoff → edge_feats - edge_radial = self.radial_embedding(edge_dist) - edge_cutoff = self.cutoff_fn(edge_dist) - edge_feats = edge_radial * edge_cutoff.unsqueeze(-1) - - return node_feats, edge_attrs, edge_feats - - -# =========================================================================== -# Interaction Block -# =========================================================================== - - -class InteractionSpec(BaseModel): - """Configuration for a single interaction block. - - Attributes: - num_features: Scalar channel multiplicity. - num_bessel: Number of Bessel radial basis functions. - l_max: Maximum angular momentum order. - avg_num_neighbors: Average number of neighbors for normalization. - """ - - model_config = ConfigDict(arbitrary_types_allowed=True) - - num_features: int = Field(..., gt=0) - num_bessel: int = Field(8, gt=0) - l_max: int = Field(2, ge=0) - avg_num_neighbors: float = Field(1.0, gt=0.0) - - -class InteractionBlock(nn.Module): - """Equivariant message passing with tensor product convolution. - - Performs geometric message passing via cuEquivariance-accelerated tensor products, - returning updated node features and skip connection for residual updates. - - Architecture: - node_feats → node_linear → tensor_product(edge_attrs, tp_weights) - → aggregate → linear → (node_feats_out, skip_connection) - - Attributes: - conv_tp: Tensor product convolution (cuEquivariance ChannelWiseTensorProduct). - node_linear: Pre-convolution equivariant linear transformation. - radial_mlp: MLP generating tensor product weights from edge features. - linear: Post-convolution equivariant linear projection. - avg_num_neighbors: Message normalization constant. - - Reference: - https://docs.nvidia.com/cuda/cuequivariance/tutorials/pytorch/MACE.html - """ - - def __init__( - self, - *, - num_features: int, - num_bessel: int = 8, - l_max: int = 2, - avg_num_neighbors: float = 1.0, - ): - """Initialize interaction block. - - Args: - num_features: Scalar channel multiplicity. - num_bessel: Number of Bessel basis functions. - l_max: Maximum angular momentum order. - avg_num_neighbors: Average neighbor count for message normalization. - """ - super().__init__() - - self.config = InteractionSpec( - num_features=num_features, - num_bessel=num_bessel, - l_max=l_max, - avg_num_neighbors=avg_num_neighbors, - ) - - # Node *state* is pure scalar (l=0); the mixed-l message irreps live only - # transiently in the tensor-product output, where they are contracted - # back to invariant scalars by the downstream ProductHead. Keeping the - # node state scalar makes every node-state operation (node_linear, - # ElementUpdate, projections) equivariant by construction — fabricating - # l>0 node components from scalars via a plain/dense linear is exactly - # what breaks rotation invariance. - node_irreps_str = f"{num_features}x0e" - irreps_str = irreps_from_l_max(l_max, num_features) # mixed-l message irreps - sh_irreps_str = sh_irreps_from_l_max(l_max) - - # 1. Tensor product convolution (define first to get weight_numel): - # scalar node features ⊗ Y_l(r̂) -> mixed-l equivariant messages. - self.conv_tp = ConvTP( - in_irreps=node_irreps_str, - out_irreps=irreps_str, - sh_irreps=sh_irreps_str, - ) - - # Actual TP output irreps (may differ from requested out_irreps) - tp_out_irreps = str(self.conv_tp.cue_tp.irreps_out) - - # 2. Pre-convolution equivariant linear (scalar -> scalar) - self.node_linear = cuet.Linear( - irreps_in=cue.Irreps("O3", node_irreps_str), - irreps_out=cue.Irreps("O3", node_irreps_str), - layout=cue.ir_mul, - dtype=config.ftype, - ) - - # 3. Radial MLP for TP weights - self.radial_mlp = RadialWeightMLP( - in_dim=num_bessel, - hidden_dim=num_features, - out_dim=self.conv_tp.weight_numel, - num_layers=2, - ) - - # 4. Post-convolution equivariant linear - self.linear = cuet.Linear( - irreps_in=cue.Irreps("O3", tp_out_irreps), - irreps_out=cue.Irreps("O3", irreps_str), - layout=cue.ir_mul, - dtype=config.ftype, - ) - - self.avg_num_neighbors = avg_num_neighbors - - def forward( - self, - node_feats: torch.Tensor, - edge_attrs: torch.Tensor, - edge_feats: torch.Tensor, - edge_index: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor]: - """Run one interaction layer. - - Args: - node_feats: Node features ``(n_nodes, irreps_dim)``. - edge_attrs: Spherical harmonics ``(n_edges, sh_dim)``. - edge_feats: Radial basis features ``(n_edges, num_bessel)``. - edge_index: Edge indices ``(n_edges, 2)``. - - Returns: - tuple of: - - ``node_feats``: Updated node features ``(n_nodes, irreps_dim)``. - - ``sc``: Skip connection (original input) ``(n_nodes, irreps_dim)``. - """ - sc = node_feats # skip connection for EquivariantProductBasisBlock - - # Pre-convolution linear - node_feats_up = self.node_linear(node_feats) - - # TP weights from radial basis - tp_weights = self.radial_mlp(edge_feats) - - # Tensor product convolution with neighbor aggregation - messages = self.conv_tp( - node_features=node_feats_up, - edge_angular=edge_attrs, - edge_index=edge_index, - tp_weights=tp_weights, - ) - - # Normalize by average number of neighbors - messages = messages / self.avg_num_neighbors - - # Post-convolution linear - node_feats = self.linear(messages) - - return node_feats, sc - - -# =========================================================================== -# MACE Encoder (Feature Extractor) -# =========================================================================== - - -class MACE(TensorDictModuleBase): - """MACE equivariant feature encoder. - - Accepts a ``TensorDict`` TensorDict and writes ``node_features`` - into the ``atoms`` sub-dict in place, returning the same - ``TensorDict`` with the new key added. - - Architecture:: - - TensorDict(atoms, edges) - → [Embedding] → node_feats, edge_attrs, edge_feats - → [Interaction₁] → [ProductHead₁] → [ElementUpdate₁] - → ... - → [Interactionₙ] → [ProductHeadₙ] - → atoms.node_features (n_nodes, num_interactions, num_features) - - Reference: - Batatia et al. "MACE: Higher Order Equivariant Message Passing Neural - Networks for Fast and Accurate Force Fields" NeurIPS 2022 - https://arxiv.org/abs/2206.07697 - """ - - in_keys = [ - ("atoms", "Z"), - ("atoms", "pos"), - ("edges", "edge_index"), - ("edges", "edge_diff"), - ("edges", "edge_dist"), - ] - out_keys = [("atoms", "node_features")] - - def __init__( - self, - *, - node_attr_specs: list[DiscreteEmbeddingSpec | ContinuousEmbeddingSpec], - num_elements: int, - num_features: int, - r_max: float, - num_bessel: int = 8, - l_max: int = 2, - num_interactions: int = 2, - correlation: int = 2, - avg_num_neighbors: float = 1.0, - layer_norm: bool = False, - ): - """Initialize MACE feature extractor. - - Args: - node_attr_specs: Embedding specs for node attributes (e.g. Z). - num_elements: Number of atomic element types. - num_features: Scalar channel multiplicity at l=0. - r_max: Radial cutoff in Angstroms. - num_bessel: Number of Bessel radial basis functions. - l_max: Maximum angular momentum order. - num_interactions: Number of interaction-product-update layers. - correlation: Body-order correlation for symmetric contraction. - avg_num_neighbors: Average neighbor count for message normalization. - layer_norm: Whether to apply layer normalization between layers. - """ - super().__init__() - - self.config = MACESpec( - node_attr_specs=node_attr_specs, - num_elements=num_elements, - num_features=num_features, - r_max=r_max, - num_bessel=num_bessel, - l_max=l_max, - num_interactions=num_interactions, - correlation=correlation, - avg_num_neighbors=avg_num_neighbors, - layer_norm=layer_norm, - ) - - # Embedding - self.embedding = EmbeddingBlock( - node_attr_specs=node_attr_specs, - num_features=num_features, - r_max=r_max, - num_bessel=num_bessel, - l_max=l_max, - ) - # Mixed-l message dimension (transient TP output consumed by ProductHead) - irreps_str = irreps_from_l_max(l_max, num_features) - with cue.assume(O3): - irreps_dim = Irreps(irreps_str).dim - - # The node *state* carried between layers is pure scalar (num_features); - # only the per-edge messages are mixed-l. This keeps every node-state - # op equivariant. Initial projection is therefore scalar -> scalar. - self.initial_projection = nn.Linear(num_features, num_features, dtype=config.ftype) - - # Interaction blocks - self.interactions = nn.ModuleList( - [ - InteractionBlock( - num_features=num_features, - num_bessel=num_bessel, - l_max=l_max, - avg_num_neighbors=avg_num_neighbors, - ) - for _ in range(num_interactions) - ] - ) - - # Product heads (from molrep, replaces former ProductBlock) - self.products = nn.ModuleList( - [ - ProductHead( - hidden_dim=irreps_dim, - out_dim=num_features, - num_radial=num_bessel, - l_max=l_max, - max_body_order=correlation, - num_species=num_elements, - ) - for _ in range(num_interactions) - ] - ) - - # Projection of the (scalar) product readout back into the scalar node - # state for the residual path. Scalar -> scalar keeps it equivariant. - self.projections = nn.ModuleList( - [ - nn.Linear(num_features, num_features, dtype=config.ftype) - for _ in range(num_interactions) - ] - ) - - # Element-specific residual updates (all layers except last). Operates on - # the scalar node state, so ElementUpdate's scalar (l=0) treatment is now - # correct rather than silently mixing l>0 components. - self.element_updates = nn.ModuleList( - [ - ElementUpdate(hidden_dim=num_features, num_species=num_elements) - for _ in range(max(num_interactions - 1, 0)) - ] - ) - - # Layer normalization (all layers except last) over the scalar state. - self.layer_norms = nn.ModuleList( - [ - nn.LayerNorm(num_features) if layer_norm else nn.Identity() - for _ in range(max(num_interactions - 1, 0)) - ] - ) - - def forward(self, td: TensorDict) -> TensorDict: - """Extract per-layer geometric features. - - Args: - td: ``TensorDict`` with ``atoms`` and ``edges`` sub-dicts. - - Returns: - Same ``TensorDict`` with ``atoms.node_features`` - ``(n_nodes, num_interactions, num_features)`` added. - """ - Z = td["atoms", "Z"] - edge_dist = td["edges", "edge_dist"] - edge_diff = td["edges", "edge_diff"] - edge_index = td["edges", "edge_index"] - - # ---- Embedding ---- - node_feats_init, edge_attrs, edge_feats = self.embedding( - Z=Z, - edge_dist=edge_dist, - edge_diff=edge_diff, - ) - - # ---- Initial projection: scalar embeddings -> hidden irreps ---- - node_feats = self.initial_projection(node_feats_init) - - # ---- Interaction-Product-Update loop ---- - per_layer_features: list[torch.Tensor] = [] - - for i in range(self.config.num_interactions): - node_feats_msg, sc = self.interactions[i]( - node_feats=node_feats, - edge_attrs=edge_attrs, - edge_feats=edge_feats, - edge_index=edge_index, - ) - - h_product = self.products[i]( - node_features=node_feats_msg, - atom_types=Z, - ) - - per_layer_features.append(h_product) - - h_proj = self.projections[i](h_product) - - is_last = i == (self.config.num_interactions - 1) - if not is_last: - node_feats = self.element_updates[i]( - h_prev=sc, - m_curr=h_proj, - atom_types=Z, - ) - node_feats = self.layer_norms[i](node_feats) - else: - node_feats = h_proj - - td["atoms", "node_features"] = torch.stack(per_layer_features, dim=1) - return td - - -class MACESpec(BaseModel): - """Configuration for the MACE feature extractor. - - Attributes: - node_attr_specs: Embedding specs for node attributes. - num_elements: Number of atomic element types. - num_features: Scalar channel multiplicity. - r_max: Radial cutoff in Angstroms. - num_bessel: Number of Bessel basis functions. - l_max: Maximum angular momentum order. - num_interactions: Number of interaction-product layers. - correlation: Body-order correlation for symmetric contraction. - avg_num_neighbors: Average neighbor count for normalization. - layer_norm: Whether to apply layer normalization. - """ - - model_config = ConfigDict(arbitrary_types_allowed=True) - - node_attr_specs: list[DiscreteEmbeddingSpec | ContinuousEmbeddingSpec] = Field( - ..., min_length=1 - ) - num_elements: int = Field(..., gt=0) - num_features: int = Field(..., gt=0) - r_max: float = Field(..., gt=0.0) - num_bessel: int = Field(8, gt=0) - l_max: int = Field(2, ge=0) - num_interactions: int = Field(2, gt=0) - correlation: int = Field(2, ge=1, le=3) - avg_num_neighbors: float = Field(1.0, gt=0.0) - layer_norm: bool = False diff --git a/src/molzoo/mace/__init__.py b/src/molzoo/mace/__init__.py new file mode 100644 index 0000000..db3f8ca --- /dev/null +++ b/src/molzoo/mace/__init__.py @@ -0,0 +1,169 @@ +"""MACE package — configuration, encoder, energy/force potential, checkpoints. + +MACE is a message-passing neural network for the potential energy of a set of +atoms (Batatia et al., NeurIPS 2022, https://arxiv.org/abs/2206.07697). It is +*equivariant*: rotate the atoms and its internal directional features rotate +with them, so the predicted energy is unchanged and the predicted forces rotate +correctly. A *foundation* variant is one shipped with weights already fitted on +a large, chemically broad dataset. This package holds molnex's native +implementation of both sides — a configurable research encoder, and the two +foundation variants (MatPES, OMol) behind +:class:`~molzoo.mace.potential.MACEPotential`. + +Layout (industrial split) — one module, one responsibility: + +* :mod:`molzoo.mace.spec` — configuration and variant presets (torch-free) +* :mod:`molzoo.mace.geometry` — edge displacement vectors and their lengths, + correct under periodic boundary conditions (PBC — the simulation cell is + taken to repeat for ever in every direction, so an atom's nearest neighbour + may be the copy of an atom across a cell face) +* :mod:`molzoo.mace.encoder` — feature encoder built from molrep blocks +* :mod:`molzoo.mace.potential` — per-graph energy ``(B,)`` in eV for the ``B`` + graphs of a batch (one graph = one molecule or one periodic cell) and + per-atom forces ``(N, 3)`` in eV/Å for its ``N`` atoms, on the post-collate + batch (the nested ``TensorDict`` schema described in CLAUDE.md) +* :mod:`molzoo.mace.checkpoint` — official-weight key remap + (:class:`~molzoo.mace.checkpoint.CheckpointRemap`, the two family tables and + their presets :data:`~molzoo.mace.checkpoint.MATPES_REMAP` / + :data:`~molzoo.mace.checkpoint.OMOL_REMAP`; both strict about unfilled + parameters, and they differ only in what an unhoused checkpoint key means) +* :mod:`molzoo.mace.variants` — named foundation models + (:class:`~molzoo.mace.variants.MACEMatpes` / + :class:`~molzoo.mace.variants.MACEOMol`), kept for the call sites that + construct them by keyword +* :mod:`molzoo.mace.research` — the freely configurable research encoder + +Four names in :data:`__all__` are **not** defined here: ``EmbeddingBlock`` / +``EmbeddingSpec`` come from :mod:`molrep.embedding.mace` and ``InteractionBlock`` +/ ``InteractionSpec`` from :mod:`molrep.interaction.mace.block`, where they were +promoted so molrep owns the reusable blocks. They are re-exported unchanged — +the same class objects, pinned by +``tests/test_molzoo/test_imports.py::TestMolzooMaceReexports`` — so existing +``from molzoo.mace import EmbeddingBlock`` imports keep resolving. + +Public import surface is stable:: + + from molzoo.mace import MACE, MACEMatpesSpec, MACEPotential + +``MACESpec`` names the shared foundation-variant configuration base +(:class:`molzoo.mace.spec.MACESpec`); the research encoder's own configuration +is :class:`molzoo.mace.research.MACEResearchSpec`. + +The configuration models are imported eagerly — they are torch-free, and +reading a config should not cost the cuEquivariance stack (NVIDIA's GPU library +for the equivariant tensor algebra MACE is built from). Everything else is +loaded **lazily**: the module object does not hold the attribute until someone +asks for it, at which point the module-level ``__getattr__`` hook of PEP 562 +imports it. That is the same policy as ``molzoo/__init__.py``, which applies it +to every one of its names. + +Reference: + Batatia et al. "MACE: Higher Order Equivariant Message Passing Neural + Networks for Fast and Accurate Force Fields" NeurIPS 2022. + https://arxiv.org/abs/2206.07697 + Batatia et al. "A foundation model for atomistic materials chemistry" + (MACE-MP-0). https://arxiv.org/abs/2401.00096 + Kaplan et al. "A foundational potential energy surface dataset for + materials" (MatPES). https://arxiv.org/abs/2503.04070 +""" + +from typing import TYPE_CHECKING + +from molzoo.mace.spec import MACEMatpesSpec, MACEOMolSpec, MACESpec + +if TYPE_CHECKING: + from molrep.embedding.mace import EmbeddingBlock, EmbeddingSpec + from molrep.interaction.mace.block import InteractionBlock, InteractionSpec + from molzoo.mace.checkpoint import ( + MATPES_KEY_REMAP, + MATPES_REMAP, + OMOL_KEY_REMAP, + OMOL_REMAP, + CheckpointRemap, + ) + from molzoo.mace.potential import MACEPotential + from molzoo.mace.research import MACE, MACEResearchSpec + from molzoo.mace.variants import ( + MACEMatpes, + MACEOMol, + load_matpes_state_dict, + load_omol_state_dict, + ) + +#: Lazily exported (PEP 562): each of these pulls in cuEquivariance (or, for +#: the checkpoint names, torch). +_LAZY = { + "MACE": "molzoo.mace.research", + "MACEPotential": "molzoo.mace.potential", + "MACEResearchSpec": "molzoo.mace.research", + "MACEMatpes": "molzoo.mace.variants", + "MACEOMol": "molzoo.mace.variants", + "load_matpes_state_dict": "molzoo.mace.variants", + "load_omol_state_dict": "molzoo.mace.variants", + "CheckpointRemap": "molzoo.mace.checkpoint", + "MATPES_KEY_REMAP": "molzoo.mace.checkpoint", + "MATPES_REMAP": "molzoo.mace.checkpoint", + "OMOL_KEY_REMAP": "molzoo.mace.checkpoint", + "OMOL_REMAP": "molzoo.mace.checkpoint", + "EmbeddingBlock": "molrep.embedding.mace", + "EmbeddingSpec": "molrep.embedding.mace", + "InteractionBlock": "molrep.interaction.mace.block", + "InteractionSpec": "molrep.interaction.mace.block", +} + +__all__ = [ + "CheckpointRemap", + "EmbeddingBlock", + "EmbeddingSpec", + "InteractionBlock", + "InteractionSpec", + "MACE", + "MACEMatpes", + "MACEMatpesSpec", + "MACEOMol", + "MACEOMolSpec", + "MACEPotential", + "MACEResearchSpec", + "MACESpec", + "MATPES_KEY_REMAP", + "MATPES_REMAP", + "OMOL_KEY_REMAP", + "OMOL_REMAP", + "load_matpes_state_dict", + "load_omol_state_dict", +] + + +def __getattr__(name: str): + """Import a lazily exported name on first attribute access (PEP 562). + + Args: + name: Attribute requested on the ``molzoo.mace`` module. + + Returns: + The object named ``name``, imported from its module in :data:`_LAZY`. + + Raises: + AttributeError: If ``name`` is not a lazily exported symbol — the same + failure a missing module attribute would give. + """ + module_name = _LAZY.get(name) + if module_name is None: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + import importlib + + return getattr(importlib.import_module(module_name), name) + + +def __dir__() -> list[str]: + """Report the public surface, so completion survives the lazy table. + + Without this, ``dir(molzoo.mace)`` lists the module globals — which, under + the lazy policy, is the torch-free config names and nothing else. It + returns exactly :data:`__all__`, the same set ``from molzoo.mace import *`` + binds, so the two cannot drift apart. Same contract as ``molzoo/__init__``. + + Returns: + The names in :data:`__all__`, sorted. + """ + return sorted(__all__) diff --git a/src/molzoo/mace/checkpoint.py b/src/molzoo/mace/checkpoint.py new file mode 100644 index 0000000..eab49ba --- /dev/null +++ b/src/molzoo/mace/checkpoint.py @@ -0,0 +1,313 @@ +"""Official MACE checkpoint → molnex ``state_dict``: one key remap for both families. + +MACE is a message-passing neural network for the potential energy of a set of +atoms (arXiv:2206.07697; the energy expression molnex evaluates is written out +in :mod:`molzoo.mace.potential`). An official MACE checkpoint is a *fitted +potential energy surface*: a ``state_dict`` — PyTorch's flat ``name -> tensor`` +dump of a model's weights — trained once, on a large dataset, and then used +unchanged. Loading one into molnex is a **dialect translation**: the same +numbers, under the names the upstream module tree happened to give them. This +module owns that translation and nothing else — no unit conversion (both sides +use eV for energies, eV/Å for forces, Å for ``r_max`` and positions, and eV/atom +for the isolated-atom reference energies), no architecture inference, no saving +direction. + +The checkpoints it reads have already been converted, out of tree, into the +layout of *cuEquivariance* — NVIDIA's GPU library for the equivariant tensor +algebra MACE is built from, abbreviated ``cueq`` in key names and throughout +this module. That conversion (``mace.cli.convert_e3nn_cueq``) is an offline +step; molnex never imports ``mace-torch``. + +:class:`CheckpointRemap` replaces the two hand-copied loaders that the +pre-cutover flat modules carried (deleted in 06-wire; see git history), which +had drifted into two byte-identical implementations of three ideas: + +1. skip the entries cuEquivariance rebuilds anyway — symbolic graph constants + (``".graph.c" in key``) and *irrep* masks (``key.endswith("output_mask")``). + An irrep, short for irreducible representation, labels how a feature + transforms when the molecule is rotated (a scalar stays put, a vector turns + with it, and so on); the mask records which of those an operation emits, and + cuEquivariance derives it from the operation's own definition; +2. rename by **longest matching prefix**, with ``None`` meaning "drop" (the + value is rebuilt from a constructor argument, or is not persistent); +3. reshape a checkpoint tensor whose ``numel`` matches but whose rank does not + — MACE stores some frozen scalars as ``(1,)`` where molnex holds a 0-d + buffer. + +The **only** genuine difference between the two shipped families — MatPES, +fitted on periodic materials, and OMol, fitted on molecules — is what an +unhoused checkpoint key means, and that is the single knob ``on_unexpected``: +MatPES raises (a key with no home means the mapping is stale), OMol returns the +list (that family ships auxiliary heads this port deliberately does not model). + +Strictness is *not* on the knob +------------------------------ + +Under both policies, an ``nn.Parameter`` the checkpoint does not fill, or a +genuine shape disagreement, raises. Two incidents in this repo are why. Both +predate this module and their original line numbers no longer resolve (the same +step that created this file trimmed those two loaders to wrappers), so they are +quoted here rather than cited by anchor: + +* The doctrine, stated by ``load_matpes_state_dict`` + (``src/molzoo/mace_matpes.py``, lines 405-410 before the trim; the sentence + still stands in that function's docstring today): "a silently dropped tensor + is the failure mode that produces a model which runs, looks sane, and is + quietly wrong". Such a model does not crash; it reports plausible energies + and forces that are simply not the fitted surface. +* The ``bessel.freqs`` incident in ``load_omol_state_dict`` + (``src/molzoo/mace_omol.py``, lines 437-439 before the trim; written up in + ``src/molzoo/specs/mace_omol.md`` §7.1) is that doctrine defeated by a + shortcut. MACE expands every interatomic distance ``r`` in a *Bessel* radial + basis — the functions ``sin(ω_n r) / r``, whose frequencies ``ω_n`` start + from the analytic values ``nπ/r_max`` (units Å⁻¹) and are declared trainable + (an ``nn.Parameter``); molnex holds them as ``bessel.freqs``. The official + OMOL checkpoint stores them as the fp32 evaluation of that analytic init + (bit-for-bit; ≤1 fp32 ulp, max 2.2e-7 Å⁻¹, from molnex's fp64 re-derivation + — measured 2026-08-09, see ``src/molzoo/specs/mace_omol.md`` §7.1). The + check for "was every learnable tensor filled?" used a ``.weight`` / + ``.bias`` **name suffix** heuristic, which silently excused that parameter, + since it ends in neither, so the checkpoint's values were dropped and the + fp64 analytic ones stayed — a ~2.2e-7 Å⁻¹ offset that went straight into + the reported energy parity residual (7.0e-7 eV). Bit-exactness against the + official surface requires the checkpoint's values regardless of their + provenance. Hence: *learnable* means exactly ``nn.Parameter``, i.e. + membership in ``model.named_parameters()``, and never a name pattern. + +Both lists returned by :meth:`CheckpointRemap.load` are **sorted**. The flat +OMol loader handed back torch's ``unexpected_keys`` in checkpoint-iteration +order; sorting is a deliberate determinism normalisation, so a diff of two +loader reports cannot move on dict ordering alone. The remap tables themselves +(:data:`MATPES_KEY_REMAP`, :data:`OMOL_KEY_REMAP`) are carried over key-for-key +from those two modules — their content is a chain invariant of the +``mace-subpackage-restructure`` spec chain and is not edited here. + +Reference: + Batatia et al. "MACE: Higher Order Equivariant Message Passing Neural + Networks for Fast and Accurate Force Fields" NeurIPS 2022. + https://arxiv.org/abs/2206.07697 + Batatia et al. "A foundation model for atomistic materials chemistry" + (MACE-MP-0). https://arxiv.org/abs/2401.00096 + Kaplan et al. "A foundational potential energy surface dataset for + materials" (MatPES). https://arxiv.org/abs/2503.04070 +""" + +from __future__ import annotations + +from collections.abc import Mapping +from typing import Literal + +import torch +from torch import nn + +#: Official cueq key (or key prefix) → molnex name for the MACE-MatPES family. +#: ``None`` drops the entry: either it is rebuilt from the constructor +#: arguments (``r_max``, the cutoff scalars, the Z-indexed ``E0`` table) or it +#: is a non-persistent constant. Longest prefix wins, so an exact key overrides +#: its enclosing prefix. Anything unlisted maps through unchanged +#: (``interactions.*``, ``products.*``, ``readouts.*``, ``scale_shift.*``). +MATPES_KEY_REMAP: dict[str, str | None] = { + "node_embedding.linear.": "node_embedding.", + "radial_embedding.bessel_fn.bessel_weights": "bessel.freqs", + "radial_embedding.bessel_fn.": None, + "radial_embedding.distance_transform.": "distance_transform.", + "radial_embedding.cutoff_fn.": None, + "pair_repulsion_fn.p": None, # exponent is a plain int on ZBLRepulsion + "pair_repulsion_fn.": "pair_repulsion.", + "atomic_energies_fn.": None, + "atomic_numbers": "z_table", + "r_max": None, + "num_interactions": None, +} + +#: Official cueq key (or key prefix) → molnex name for the MACE-OMOL family, +#: mirroring :data:`MATPES_KEY_REMAP`. ``None`` drops the entry (rebuilt from +#: the constructor arguments or a non-persistent constant). Longest prefix +#: wins; anything unlisted maps through unchanged (``interactions.*``, +#: ``products.*``, ``joint_embedding.*``, ``scale_shift.*``). +OMOL_KEY_REMAP: dict[str, str | None] = { + "node_embedding.linear.": "node_embedding.", + "embedding_readout.linear.": "embedding_readout.", + "readouts.0.": "readout.", + # MACE's fitted Bessel frequencies (Å⁻¹). Without this entry the key is + # simply unexpected and `bessel.freqs` keeps its analytic n*pi/r_max init — + # the official OMOL values sit ~2.2e-7 Å⁻¹ away, which was landing straight + # in the reported energy parity residual (module docstring, §"Strictness"). + "radial_embedding.bessel_fn.bessel_weights": "bessel.freqs", + "radial_embedding.bessel_fn.": None, + "radial_embedding.cutoff_fn.": None, + "atomic_energies_fn.": None, # handled at construction (Z-indexed) + "atomic_numbers": "z_table", + "r_max": None, + "num_interactions": None, +} + + +class CheckpointRemap: + """Translate official MACE checkpoint keys into molnex ``state_dict`` keys. + + Construct one per checkpoint family (the two shipped ones are the module + constants :data:`MATPES_REMAP` and :data:`OMOL_REMAP`), then either inspect + the translation with :meth:`rename` or apply it with :meth:`load`:: + + MATPES_REMAP.load(model, torch.load(path, weights_only=True)) + + The prefix search order — longest first, so an exact key wins over the + prefix that encloses it — is computed once here rather than kept as a + second module-level constant beside the table, which is what the two flat + loaders did. + + A complete, runnable use of this surface — a synthetic checkpoint written in + the official dialect, read back, and checked against hard-coded energy / + force goldens, including both ``on_unexpected`` policies — is + ``regressions/mace-subpackage-restructure-05-checkpoint.py``. + + Args: + table: Official cueq key (or key prefix) → molnex name. ``None`` drops + the entry. Keys absent from the table pass through unchanged. + on_unexpected: What :meth:`load` does with a renamed key that has no + home in the target model. ``"raise"`` (the default — the stricter + behaviour is the default) refuses before touching the model; + ``"return"`` hands the list back to the caller, which is what the + OMol family needs because its checkpoints carry auxiliary heads + this port deliberately does not model. + + Raises: + ValueError: If ``on_unexpected`` is neither ``"raise"`` nor + ``"return"``. A typo must not degrade into the lax branch. + """ + + def __init__( + self, + table: Mapping[str, str | None], + *, + on_unexpected: Literal["raise", "return"] = "raise", + ) -> None: + if on_unexpected not in ("raise", "return"): + raise ValueError(f'on_unexpected must be "raise" or "return", got {on_unexpected!r}') + self._table: dict[str, str | None] = dict(table) + #: Prefixes longest-first: an exact key beats the prefix enclosing it. + self._prefixes: tuple[str, ...] = tuple(sorted(self._table, key=len, reverse=True)) + self._on_unexpected = on_unexpected + + def rename(self, official_state: Mapping[str, torch.Tensor]) -> dict[str, torch.Tensor]: + """Translate checkpoint keys, dropping what the model rebuilds itself. + + Pure: ``official_state`` is not modified, no model is consulted, and + the tensors are passed through by reference — nothing is copied, cast + or converted, so every value keeps the units it had upstream (the two + sides agree: energies in eV, isolated-atom references in eV/atom, + forces in eV/Å, lengths in Å, Bessel frequencies in Å⁻¹). + + Dropped entries are the cuEquivariance symbolic graph constants + (``".graph.c" in key``) and irrep masks (``output_mask``), which the + modules rebuild and which carry nothing learned, plus every key whose + table entry is ``None``. + + Args: + official_state: ``state_dict`` of the cueq-converted official model + (converted by ``mace.cli.convert_e3nn_cueq``, run out of tree — + molnex never imports ``mace-torch``). + + Returns: + A new dict of molnex key → the same tensor objects. + """ + renamed: dict[str, torch.Tensor] = {} + for key, value in official_state.items(): + if ".graph.c" in key or key.endswith("output_mask"): + continue + new_key: str | None = key + for prefix in self._prefixes: + if key.startswith(prefix): + replacement = self._table[prefix] + new_key = None if replacement is None else replacement + key[len(prefix) :] + break + if new_key is not None: + renamed[new_key] = value + return renamed + + def load( + self, model: nn.Module, official_state: Mapping[str, torch.Tensor] + ) -> tuple[list[str], list[str]]: + """Load an official checkpoint into ``model``, strictly. + + One named action: :meth:`rename`, then the unexpected-key policy, then + the numel-conserving reshape, then ``load_state_dict(strict=False)``, + then the strictness check. The policy branch and the shape check both + run *before* any tensor reaches the model, so those two rejections + never leave a half-loaded model behind. The parameter-coverage check is + the one that cannot: it can only ask what ``load_state_dict`` failed to + fill, so that failure does leave the model written to. It is fatal by + design — construct a fresh model rather than retrying on this one. + + "Learnable" is defined as membership in ``model.named_parameters()``. + Never a ``.weight`` / ``.bias`` name heuristic — that shortcut once + excused the ``nn.Parameter`` ``bessel.freqs`` from this very check; the + module docstring tells that story. + + Args: + model: Target model, already constructed with the checkpoint's + hyper-parameters. + official_state: ``state_dict`` of the cueq-converted official model. + + Returns: + ``(missing_buffers, unexpected)``: the non-learnable entries the + checkpoint did not provide (typically the graph constants cueq + rebuilds) and the renamed keys with no home in ``model``. Both + lists are **sorted** — a deterministic normalisation, so two runs + differ only when the checkpoint does. Under ``on_unexpected = + "raise"`` the second list is always empty. + + Raises: + RuntimeError: If a renamed key has no home in ``model`` and the + policy is ``"raise"``; if a checkpoint tensor disagrees in + ``numel`` with the model's (the model was built from the wrong + config); or — under **either** policy — if any + ``nn.Parameter`` of ``model`` is left unfilled, which would + silently keep whatever the constructor gave it (random for a + linear weight, the analytic ``nπ/r_max`` for ``bessel.freqs``). + """ + renamed = self.rename(official_state) + own = model.state_dict() + + unexpected = sorted(set(renamed) - set(own)) + if unexpected and self._on_unexpected == "raise": + raise RuntimeError( + f"checkpoint keys with no home in {type(model).__name__}: {unexpected}" + ) + + mismatched: list[str] = [] + for key, value in renamed.items(): + if key not in own: + continue # reported through ``unexpected`` + want = own[key].shape + if value.shape == want: + continue + # MACE stores some frozen scalars as (1,) where molnex holds a 0-d + # buffer; identical content, different rank. + if value.numel() == own[key].numel(): + renamed[key] = value.reshape(want) + else: + mismatched.append(f"{key}: checkpoint {tuple(value.shape)} vs model {tuple(want)}") + if mismatched: + raise RuntimeError( + "shape mismatch — model built with the wrong config? " + + ", ".join(sorted(mismatched)) + ) + + missing, _ = model.load_state_dict(renamed, strict=False) + parameters = {name for name, _ in model.named_parameters()} + unfilled = sorted(name for name in missing if name in parameters) + if unfilled: + raise RuntimeError(f"parameters not covered by the checkpoint: {unfilled}") + return sorted(set(missing) - parameters), unexpected + + +#: The MACE-MatPES loading policy: a checkpoint key with no home means the +#: mapping is stale, so refuse rather than load part of the surface. +MATPES_REMAP = CheckpointRemap(MATPES_KEY_REMAP) + +#: The MACE-OMOL loading policy: unhoused keys are *returned*, because the OMol +#: checkpoints carry auxiliary heads this port deliberately does not model. +#: Everything else (unfilled parameters, shape disagreements) still raises. +OMOL_REMAP = CheckpointRemap(OMOL_KEY_REMAP, on_unexpected="return") diff --git a/src/molzoo/mace/encoder.py b/src/molzoo/mace/encoder.py new file mode 100644 index 0000000..8cf4a47 --- /dev/null +++ b/src/molzoo/mace/encoder.py @@ -0,0 +1,424 @@ +"""Configuration-driven backbone shared by the MACE foundation models. + +:class:`MACEEncoder` registers the module graph of one MACE variant — which +one is decided entirely by the five ``Literal`` switches on +:class:`~molzoo.mace.spec.MACESpec`. Submodule names are byte-identical to the +flat ``MACEMatpes`` / ``MACEOMol`` models, so an official checkpoint keeps +loading without a key rewrite and the energy/force variant classes can simply +inherit this backbone (upstream's ``ScaleShiftMACE(MACE)`` shape). + +The class exposes primitives, not a pipeline: ``validate_elements`` / +``node_attrs`` / ``initial_node_features`` / ``angular_features`` / +``radial_features`` / ``conditioning`` / ``layer_features``. Composing them +into an energy — and differentiating it — is the caller's job. + +Hot-path discipline: every branch is frozen into a plain attribute (or the +presence of an optional submodule) by ``__init__``. The configuration object +survives only as ``self._spec``, for provenance; reading a pydantic attribute +inside a layer loop is a dynamo graph break (cf. ``src/molzoo/mace_matpes.py:118`` +at 0e05959, before deletion). + +Reference: + Batatia et al. "MACE: Higher Order Equivariant Message Passing Neural + Networks for Fast and Accurate Force Fields" NeurIPS 2022. + https://arxiv.org/abs/2206.07697 + Batatia et al. "A foundation model for atomistic materials chemistry" + (MACE-MP-0). https://arxiv.org/abs/2401.00096 + Kaplan et al. "A foundational potential energy surface dataset for + materials" (MatPES). https://arxiv.org/abs/2503.04070 +""" + +from __future__ import annotations + +import cuequivariance as cue +import cuequivariance_torch as cuet +import torch +import torch.nn as nn + +from molix import config +from molpot.heads.energy import AtomicReferenceEnergy +from molpot.heads.rescale import GlobalRescale +from molpot.potentials.repulsion import ZBLRepulsion +from molrep.embedding.angular import SphericalHarmonics +from molrep.embedding.cutoff import PolynomialCutoff +from molrep.embedding.node import JointFeatureEmbedding, JointFeatureSpec +from molrep.embedding.radial import AgnesiTransform, BesselRBF +from molrep.interaction.mace.density import DensityInteraction, DensityResidualInteraction +from molrep.interaction.product_basis import EquivariantProductBasis +from molrep.interaction.residual import ResidualInteraction +from molrep.readout.mace import LinearReadout, NonLinearBiasReadout, NonLinearReadout +from molzoo.mace.spec import MACEMatpesSpec, MACEOMolSpec, MACESpec + + +def _irreps(l_max: int, mul: int) -> str: + """``mul``-fold irreps up to ``l_max`` with natural parity ``(-1)^l``. + + Args: + l_max: Highest angular-momentum order (inclusive). + mul: Channel multiplicity of every order. + + Returns: + A cuEquivariance irreps string, e.g. ``"128x0e+128x1o+128x2e"``. + """ + return "+".join(f"{mul}x{l}{'e' if l % 2 == 0 else 'o'}" for l in range(l_max + 1)) + + +class MACEEncoder(nn.Module): + """MACE foundation-model backbone, built from a validated configuration. + + Args: + spec: A validated :class:`~molzoo.mace.spec.MACESpec` — in practice a + :class:`~molzoo.mace.spec.MACEMatpesSpec` or + :class:`~molzoo.mace.spec.MACEOMolSpec`, which carry the + variant-only fields the corresponding stack needs. + + Raises: + TypeError: If the ``interaction`` switch asks for a stack whose + variant-only fields the given spec does not carry. + """ + + def __init__(self, spec: MACESpec) -> None: + super().__init__() + ftype = config.ftype + #: Provenance only (05-checkpoint reads it); never touched on a hot path. + self._spec = spec + + n_el = len(spec.atomic_numbers) + l_max = spec.l_max + num_features = spec.num_features + num_interactions = spec.num_interactions + + # ---- frozen switches (plain attributes: no pydantic on the hot path) ---- + is_residual = spec.interaction == "residual" + self.num_interactions = num_interactions + # MACE-MP/MatPES multiplies the polynomial envelope into the radial + # basis; OMOL keeps it separate and hands it to the interaction, which + # applies it inside the message. Same envelope, different fold point. + self.fold_cutoff_into_radial = not is_residual + self.pass_cutoff_to_interaction = is_residual + + node_attrs_irreps = f"{n_el}x0e" + feat0 = f"{num_features}x0e" + sh = _irreps(l_max, 1) + # Transient message irreps: node scalars ⊗ Y_l, e.g. "128x0e+…+128x3o". + target = _irreps(l_max, num_features) + + # ---- per-layer irreps schedule (the one place the variants differ) ---- + if not is_residual: + if not isinstance(spec, MACEMatpesSpec): + raise TypeError( + 'interaction="density" needs a MACEMatpesSpec ' + "(max_hidden_l / radial_mlp are not on the base spec)" + ) + hidden = _irreps(spec.max_hidden_l, num_features) + edge_irr = [feat0] + [hidden] * (num_interactions - 1) + radial_mlp = list(spec.radial_mlp) + else: + if not isinstance(spec, MACEOMolSpec): + raise TypeError( + 'interaction="residual" needs a MACEOMolSpec ' + "(edge_channels is not on the base spec)" + ) + # OMOL drops the top order from the node state and squeezes the + # mid-layer edge irreps through the edge_channels bottleneck. + hidden = _irreps(l_max - 1, num_features) + edge_irr = [feat0] + [_irreps(l_max - 1, spec.edge_channels)] * (num_interactions - 1) + radial_mlp = [spec.edge_channels] * 3 + node_in = [feat0] + [hidden] * (num_interactions - 1) + # The last layer keeps scalars only; every earlier layer keeps `hidden`. + hidden_sched = [hidden] * (num_interactions - 1) + [feat0] + + # ---- embeddings ---- + self.node_embedding = cuet.Linear( + cue.Irreps("O3", node_attrs_irreps), + cue.Irreps("O3", feat0), + layout=cue.ir_mul, + dtype=ftype, + ) + self.spherical_harmonics = SphericalHarmonics(l_max=l_max) + # normalize=False + eps=0 + trainable reproduces MACE's BesselBasis exactly. + self.bessel = BesselRBF( + r_cut=spec.r_max, + num_radial=spec.num_bessel, + normalize=False, + eps=0.0, + trainable=True, + ) + self.distance_transform: AgnesiTransform | None = ( + AgnesiTransform() if spec.distance_transform == "agnesi" else None + ) + self.cutoff_fn = PolynomialCutoff(r_cut=spec.r_max, exponent=spec.num_polynomial_cutoff) + self.pair_repulsion: ZBLRepulsion | None = ( + ZBLRepulsion(exponent=spec.num_polynomial_cutoff) + if spec.pair_repulsion == "zbl" + else None + ) + self.joint_embedding: JointFeatureEmbedding | None = None + self.embedding_readout: cuet.Linear | None = None + if spec.conditioning == "charge_spin": + if not isinstance(spec, MACEOMolSpec): + raise TypeError( + 'conditioning="charge_spin" needs a MACEOMolSpec ' + "(charge/spin class counts are not on the base spec)" + ) + self.joint_embedding = JointFeatureEmbedding( + feature_specs=[ + JointFeatureSpec( + name="total_spin", + kind="categorical", + emb_dim=num_features, + num_classes=spec.spin_classes, + per="graph", + offset=spec.spin_offset, + ), + JointFeatureSpec( + name="total_charge", + kind="categorical", + emb_dim=num_features, + num_classes=spec.charge_classes, + per="graph", + offset=spec.charge_offset, + ), + ], + out_dim=num_features, + ) + self.embedding_readout = cuet.Linear( + cue.Irreps("O3", feat0), + cue.Irreps("O3", "1x0e"), + layout=cue.ir_mul, + dtype=ftype, + ) + self.atomic_energies = AtomicReferenceEnergy( + atomic_energies=spec.atomic_energies, + atomic_numbers=spec.atomic_numbers, + ) + self.register_buffer( + "z_table", torch.tensor(spec.atomic_numbers, dtype=torch.long), persistent=True + ) + self.z_table: torch.Tensor + + # ---- interaction / product stack ---- + # The density family carries per-element symmetric-contraction weights; + # the OMOL residual family shares one set across elements + # (``num_elements=1``, mace_omol.py:179). + product_elements = 1 if is_residual else n_el + self.interactions = nn.ModuleList() + self.products = nn.ModuleList() + for i in range(num_interactions): + interaction: nn.Module + if is_residual: + interaction = ResidualInteraction( + node_attrs_irreps=node_attrs_irreps, + node_feats_irreps=node_in[i], + edge_attrs_irreps=sh, + edge_feats_irreps=f"{spec.num_bessel}x0e", + edge_irreps=edge_irr[i], + target_irreps=target, + hidden_irreps=hidden_sched[i], + radial_mlp=radial_mlp, + use_fallback=spec.use_fallback, + ) + elif i == 0: + # The first layer has no residual to carry, hence no skip input. + interaction = DensityInteraction( + node_attrs_irreps=node_attrs_irreps, + node_feats_irreps=node_in[i], + edge_attrs_irreps=sh, + edge_feats_irreps=f"{spec.num_bessel}x0e", + edge_irreps=edge_irr[i], + target_irreps=target, + radial_mlp=radial_mlp, + use_fallback=spec.use_fallback, + ) + else: + interaction = DensityResidualInteraction( + node_attrs_irreps=node_attrs_irreps, + node_feats_irreps=node_in[i], + edge_attrs_irreps=sh, + edge_feats_irreps=f"{spec.num_bessel}x0e", + edge_irreps=edge_irr[i], + target_irreps=target, + hidden_irreps=hidden_sched[i], + radial_mlp=radial_mlp, + use_fallback=spec.use_fallback, + ) + self.interactions.append(interaction) + self.products.append( + EquivariantProductBasis( + node_feats_irreps=target, + target_irreps=hidden_sched[i], + correlation=spec.correlation, + num_elements=product_elements, + use_sc=is_residual or i > 0, + use_fallback=spec.use_fallback, + ) + ) + + # ---- readout ---- + self.readouts: nn.ModuleList | None = None + self.readout: NonLinearBiasReadout | None = None + if spec.readout == "per_layer": + self.readouts = nn.ModuleList( + [ + NonLinearReadout(irreps_in=hidden_sched[i], mlp_dim=spec.mlp_dim) + if i == num_interactions - 1 + else LinearReadout(irreps_in=hidden_sched[i]) + for i in range(num_interactions) + ] + ) + else: + self.readout = NonLinearBiasReadout(irreps_in=feat0, mlp_dim=spec.mlp_dim) + + self.scale_shift = GlobalRescale(scale=spec.scale, shift=spec.shift) + + def validate_elements(self, Z: torch.Tensor) -> None: + """Raise if any atomic number in ``Z`` is outside the element table. + + An element outside the table would be snapped onto a neighbouring row + by ``searchsorted`` and silently produce a wrong energy. The check + costs a host sync, so it is **not** meant for the per-step path: run it + once per (model, system) pair — ``Z`` is constant over a trajectory, + and a wrongly wired model/dataset pair fails on the first batch. + + Args: + Z: Atomic numbers ``(N,)``. + + Raises: + ValueError: Listing the atomic numbers outside the table. + """ + unknown = torch.unique(Z[~torch.isin(Z, self.z_table)]) + if unknown.numel(): + raise ValueError( + f"atomic numbers {unknown.tolist()} are outside this model's " + f"{self.z_table.numel()}-element table" + ) + + def node_attrs(self, Z: torch.Tensor, dtype: torch.dtype) -> torch.Tensor: + """One-hot node attributes over the element table. + + Args: + Z: Atomic numbers ``(N,)``; must be inside the table (see + :meth:`validate_elements`). + dtype: Floating dtype of the returned one-hot. + + Returns: + One-hot element attributes ``(N, n_elements)``. + """ + z_index = torch.searchsorted(self.z_table, Z.reshape(-1)).to(dtype=torch.long) + return torch.nn.functional.one_hot(z_index, self.z_table.numel()).to(dtype) + + def initial_node_features(self, node_attrs: torch.Tensor) -> torch.Tensor: + """Project the one-hot element attributes into the scalar node state. + + Args: + node_attrs: One-hot element attributes ``(N, n_elements)``. + + Returns: + Initial node features ``(N, num_features)``. + """ + return self.node_embedding(node_attrs) + + def angular_features(self, vectors: torch.Tensor) -> torch.Tensor: + """Real spherical harmonics ``Y_l(r̂)`` of the edge displacements. + + Args: + vectors: Edge displacements ``(E, 3)`` in Å, e.g. from + :func:`molzoo.mace.geometry.edge_vectors`. + + Returns: + Edge attributes ``(E, (l_max + 1)²)``. + """ + return self.spherical_harmonics(vectors) + + def radial_features( + self, lengths: torch.Tensor, Z: torch.Tensor, edge_index: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + """Radial basis and cutoff envelope of the edges. + + The envelope is always evaluated on the *raw* distance; whether it is + folded into the returned basis (MACE-MP/MatPES) or left for the + interaction to apply (OMOL) follows the ``interaction`` switch. + + Args: + lengths: Edge lengths ``(E,)`` in Å, e.g. from + :func:`molzoo.mace.geometry.edge_lengths`. + Z: Atomic numbers ``(N,)`` — needed by the Agnesi transform, which + is element-pair dependent. + edge_index: ``(E, 2)`` with ``[:, 0]`` = source, ``[:, 1]`` = target. + + Returns: + ``(edge_feats (E, num_bessel), cutoff (E, 1))``. + """ + cutoff = self.cutoff_fn(lengths).unsqueeze(-1) + distances = lengths + if self.distance_transform is not None: + distances = self.distance_transform(lengths, Z[edge_index[:, 0]], Z[edge_index[:, 1]]) + edge_feats = self.bessel(distances) + if self.fold_cutoff_into_radial: + edge_feats = edge_feats * cutoff + return edge_feats, cutoff + + def conditioning( + self, + batch: torch.Tensor, + *, + total_spin: torch.Tensor, + total_charge: torch.Tensor, + ) -> torch.Tensor: + """Per-atom charge / spin embedding, broadcast from per-graph scalars. + + Args: + batch: Graph index per atom ``(N,)``. + total_spin: Per-graph total spin ``(B,)``. + total_charge: Per-graph total charge ``(B,)`` in units of ``e``. + + Returns: + Conditioning features ``(N, num_features)`` to add to the initial + node features. + + Raises: + ValueError: If this variant has no charge/spin conditioning. + """ + if self.joint_embedding is None: + raise ValueError( + "this MACE variant carries no charge/spin conditioning " + '(conditioning="none"); build it from a MACEOMolSpec instead' + ) + return self.joint_embedding(batch, total_spin=total_spin, total_charge=total_charge) + + def layer_features( + self, + *, + node_feats: torch.Tensor, + node_attrs: torch.Tensor, + edge_attrs: torch.Tensor, + edge_feats: torch.Tensor, + edge_index: torch.Tensor, + cutoff: torch.Tensor, + ) -> list[torch.Tensor]: + """Run the interaction/product stack, returning the per-layer node state. + + A ``list`` rather than a stacked tensor: the layer widths differ (every + layer but the last carries ``hidden`` irreps, the last carries scalars). + + Args: + node_feats: Initial node features ``(N, num_features)``. + node_attrs: One-hot element attributes ``(N, n_elements)``. + edge_attrs: Spherical harmonics ``(E, (l_max + 1)²)``. + edge_feats: Radial basis features ``(E, num_bessel)``. + edge_index: ``(E, 2)`` with ``[:, 0]`` = source, ``[:, 1]`` = target. + cutoff: Per-edge cutoff envelope ``(E, 1)``; ignored by variants + that already folded it into ``edge_feats``. + + Returns: + One node-feature tensor per interaction layer, each ``(N, …)``. + """ + envelope = cutoff if self.pass_cutoff_to_interaction else None + per_layer: list[torch.Tensor] = [] + for i in range(self.num_interactions): + node_feats, sc = self.interactions[i]( + node_attrs, node_feats, edge_attrs, edge_feats, edge_index, envelope + ) + node_feats = self.products[i](node_feats, sc, node_attrs) + per_layer.append(node_feats) + return per_layer diff --git a/src/molzoo/mace/geometry.py b/src/molzoo/mace/geometry.py new file mode 100644 index 0000000..1eee121 --- /dev/null +++ b/src/molzoo/mace/geometry.py @@ -0,0 +1,67 @@ +"""Edge displacement / length for the MACE family. + +MACE's edge vector is ``r_ij = pos[target] - pos[source] + S_ij`` with an +*additive* periodic shift ``S_ij = n_ij · h`` (``unit_shifts @ cell``), exactly +as the upstream ``MACE`` models consume a neighbour list. ``S_ij`` is constant +with respect to ``pos``, so ``∂r_ij/∂pos`` is the open-system one and forces +stay exact. + +Why MACE owns a geometry module of its own, next to +``molzoo.pinet.geometry``: PiNet's ``edge_bond_diff`` +(``src/molzoo/pinet/geometry.py:28-49``) computes +``detach(imaged) + (raw - detach(raw))`` — the *minimum-image* value with a +straight-through gradient. The two share a gradient (``∂raw/∂pos``) but not a +value: minimum image versus an explicit integer translation are different +mathematical objects, and folding them into one helper would smuggle PiNet's +straight-through semantics into MACE's periodic path. Composition of the two +functions below is the caller's job; there is no one-step façade. + +Reference: + Batatia et al. "MACE: Higher Order Equivariant Message Passing Neural + Networks for Fast and Accurate Force Fields" NeurIPS 2022. + https://arxiv.org/abs/2206.07697 +""" + +from __future__ import annotations + +import torch + + +def edge_vectors( + pos: torch.Tensor, + edge_index: torch.Tensor, + shifts: torch.Tensor | None = None, +) -> torch.Tensor: + """Source→target edge displacements, with optional periodic shifts. + + Args: + pos: Atom positions ``(N, 3)`` in Å. + edge_index: Source/target pairs ``(E, 2)``; ``[:, 0]`` = source, + ``[:, 1]`` = target (the repo-wide edge convention). + shifts: Optional periodic shift vectors ``(E, 3)`` in Å + (``unit_shifts @ cell``), added to the raw displacement. + + Returns: + Edge displacement ``r_ij = pos[target] - pos[source] (+ S_ij)`` + ``(E, 3)`` in Å. + """ + vectors = pos[edge_index[:, 1]] - pos[edge_index[:, 0]] + if shifts is not None: + vectors = vectors + shifts + return vectors + + +def edge_lengths(vectors: torch.Tensor, *, keepdim: bool = False) -> torch.Tensor: + """Euclidean length of each edge displacement. + + Args: + vectors: Edge displacements ``(E, 3)`` in Å, e.g. from + :func:`edge_vectors`. + keepdim: Keep the contracted axis, giving ``(E, 1)`` instead of + ``(E,)``. + + Returns: + Edge lengths ``d_ij = ‖r_ij‖`` — ``(E,)``, or ``(E, 1)`` when + ``keepdim`` — in Å. + """ + return torch.linalg.norm(vectors, dim=-1, keepdim=keepdim) diff --git a/src/molzoo/mace/potential.py b/src/molzoo/mace/potential.py new file mode 100644 index 0000000..da930d6 --- /dev/null +++ b/src/molzoo/mace/potential.py @@ -0,0 +1,731 @@ +"""Unified MACE energy/force potential for the foundation-model variants. + +An *interatomic potential* is a function from an arrangement of atoms — their +chemical element numbers ``Z_i`` and Cartesian positions ``r_i``, in ångström +(Å) — to the potential energy ``E`` of that arrangement, in electronvolt (eV), +and from there to the force on each atom, which is the negative gradient of +that energy:: + + F_i = -∂E/∂r_i [eV/Å] + +That sign is the repo-wide convention (``molix.schema.FORCES_KEY``). MACE +*learns* such a function: it describes each atom's neighbourhood by the +displacement vectors to the neighbours inside a cutoff radius ``r_max`` and +turns that description into a per-atom energy through several rounds of +*message passing* — each round lets an atom mix in features of its neighbours. +A *foundation* variant is one whose weights were fitted once on a large, +chemically broad dataset and are then used unchanged; molnex ships two, MatPES +(periodic materials, arXiv:2503.04070) and OMol (molecules, additionally +conditioned on the total charge and total spin of the system). + +Per *graph* ``g`` — one graph is one molecule, or one periodic cell, inside a +batch of ``B`` of them — this class evaluates exactly:: + + E_g = Σ_{i∈g} E0(Z_i) + Σ_{i∈g} (scale · ε_i + shift) [eV] + ε_i = ZBL_i + Σ_layer readout_layer(h_i^layer) [eV] + +``E0(Z)`` is the frozen isolated-atom reference energy of element ``Z``, +``h_i^layer`` is atom ``i``'s feature vector after each message-passing layer, +``ZBL_i`` is the short-range Ziegler–Biersack–Littmark nuclear repulsion — +present only when the spec asks for it (``pair_repulsion="zbl"``, the MatPES +default; identically zero otherwise) — and ``scale`` / ``shift`` are the fitted +affine normalisation of the learned part, applied *per atom* +(:class:`molpot.heads.rescale.GlobalRescale`). MatPES reads out every layer, so +the inner sum runs over all of them; OMol reads out the last layer only, so its +inner sum has a single term, and it adds one further per-atom contribution from +its charge/spin embedding next to ``E0`` — that is, *outside* the +``scale``/``shift`` normalisation. The forces are never hand-coded: they are +the autograd derivative of this same ``E_g``. + +:class:`MACEPotential` is the one energy + force host for every MACE +foundation variant. Which variant it is comes entirely from the +:class:`~molzoo.mace.spec.MACESpec` it is built from — the MatPES density +stack (arXiv:2503.04070) or the OMOL charge/spin-conditioned residual stack — +so the two pre-cutover flat modules that each duplicated this forward +(deleted in the restructure's wiring step, 06-wire; see git history) collapse +into this single class. + +It **inherits** :class:`~molzoo.mace.encoder.MACEEncoder` rather than holding +one, mirroring upstream's ``ScaleShiftMACE(MACE)``: holding an encoder would +prefix every key of the ``state_dict`` (PyTorch's flat ``name -> tensor`` dump +of a model's weights) with ``encoder.``, and an official checkpoint would stop +loading without a key rewrite. + +Two seams, deliberately coexisting: + +``energy_core(positions, Z, edge_index, batch, num_graphs, …) -> (B,)`` + The **flat, public, compilable** one. Raw tensors in, per-graph energy in + eV out; no ``TensorDict`` (the nested batch container of the post-collate + schema, see CLAUDE.md) and no ``.item()`` *host sync* — reading a tensor + value back into Python makes the CPU wait for the GPU and forces + TorchDynamo, the tracer behind ``torch.compile``, to cut the traced region + in two (a *graph break*, which costs back the launch overhead compilation + was meant to remove). ``num_graphs`` is an argument for exactly that + reason: one traced graph, no break. This is the entry a molecular-dynamics + (MD) driver or an engine export binds to. + +``_write_energy(batch) -> batch`` + The **protocol hook**, one level up, on the batch schema. + ``molpot.derivation.protocol.call_energy`` prefers ``model._write_energy``, + so implementing it wires MACE into ``EnergyReadout`` / ``ForceReadout`` for + free. It reads the tensors off the post-collate batch, calls + :meth:`MACEPotential.energy_core`, and writes the result back through + ``protocol.write_energy`` — **without** detaching either the positions or + the energy. Detaching a tensor cuts it out of the autograd graph; an outer + readout session may have created the position *leaf* (the tensor autograd + differentiates with respect to) itself, so a detach here would sever + ``F = -∂E/∂r`` before anyone takes the derivative. (The flat models + detached inside their ``energy_forces(compute_forces=False)``; that duty + belongs to :class:`molix.md.forcefield.PotentialForceField`, not to the + potential.) + +Every branch — forces on/off, charge/spin conditioning, per-layer vs final +readout — is resolved once in ``__init__`` into a bound attribute, and +``forward`` is a single dispatch through ``self._pipeline``. There is no +``compute_forces`` flag on any public signature: "energy only" is a +*construction* (``MACEPotential(spec, compute_forces=False)``), which is also +what ``molix.md.forcefield.PotentialForceField.calc_energy`` already assumes. + +Units: ``graphs.energy`` in eV ``(B,)``, ``atoms.forces`` in eV/Å ``(N, 3)``, +positions and ``r_max`` in Å. The bridge to the MD integrator's unit system +(mass in atomic mass units, length in Å, time in fs, hence energy in +amu·Å²/fs²) lives in ``molix.md.forcefield``; nothing here converts units. + +Reference: + Batatia et al. "MACE: Higher Order Equivariant Message Passing Neural + Networks for Fast and Accurate Force Fields" NeurIPS 2022. + https://arxiv.org/abs/2206.07697 + Batatia et al. "A foundation model for atomistic materials chemistry" + (MACE-MP-0). https://arxiv.org/abs/2401.00096 + Kaplan et al. "A foundational potential energy surface dataset for + materials" (MatPES). https://arxiv.org/abs/2503.04070 +""" + +from __future__ import annotations + +import json +import re +from pathlib import Path + +import torch +from tensordict import TensorDict + +from molix.F.scatter import scatter_sum_compile_safe as _scatter_sum +from molix.schema import POS_KEY +from molpot.derivation.kernels import grad_force_pass +from molpot.derivation.protocol import write_energy +from molzoo.mace.checkpoint import MATPES_REMAP, CheckpointRemap +from molzoo.mace.encoder import MACEEncoder +from molzoo.mace.geometry import edge_lengths, edge_vectors +from molzoo.mace.spec import MACEMatpesSpec, MACEOMolSpec + +#: One term of an official irreps string, e.g. ``"128x1o"`` — multiplicity, +#: angular order ``l``, parity. Parity is fixed by ``l`` in every MACE +#: configuration, so only the first two groups are read. +_IRREPS_TERM = re.compile(r"(\d+)x(\d+)([eo])") + + +def _parse_irreps(text: str, key: str) -> tuple[int, int]: + """Scalar multiplicity and highest angular order of an official irreps string. + + ``"128x0e+128x1o"`` → ``(128, 1)``; ``"16x0e"`` → ``(16, 0)``. Upstream MACE + states channel widths this way, so a checkpoint's ``num_features`` / + ``max_hidden_l`` / ``mlp_dim`` are readable from its config instead of being + assumed (``run_nve.py``'s former ``build_model``, removed in + ``mace-subpackage-restructure-07-cleanup``, hard-coded ``128 / 1 / 16``, + which loads a differently sized checkpoint into the wrong model). + + Args: + text: Irreps string from the official config. + key: Config key it came from, for the error message. + + Returns: + ``(multiplicity of the scalar 0e term, highest l over all terms)``. + + Raises: + ValueError: If a term is not ``x``, or the string has + no scalar term — the channel width would then be a guess. + """ + terms: list[tuple[int, int]] = [] + for term in str(text).split("+"): + matched = _IRREPS_TERM.fullmatch(term.strip()) + if matched is None: + raise ValueError(f"config key {key} is not an irreps string: {text!r} (term {term!r})") + terms.append((int(matched.group(1)), int(matched.group(2)))) + scalars = [multiplicity for multiplicity, order in terms if order == 0] + if not scalars: + raise ValueError( + f"config key {key} has no scalar (0e) term: {text!r} — the channel " + "width is its multiplicity and there is no default" + ) + return scalars[0], max(order for _, order in terms) + + +class MACEPotential(MACEEncoder): + """MACE foundation-model energy (and forces) on the molnex batch schema. + + Given a batch of atomic systems it writes the per-graph potential energy + ``graphs.energy`` ``(B,)`` in eV and, unless built energy-only, the per-atom + forces ``atoms.forces`` ``(N, 3)`` in eV/Å; the energy expression it + evaluates and the ``F_i = -∂E/∂r_i`` convention are spelled out in the + module docstring. The variant is the spec's business; the only two switches + here decide *how much* is computed and *which cuEquivariance path* the + blocks take:: + + potential = MACEPotential(MACEMatpesSpec(...)) # E + F + potential = MACEPotential(spec, compute_forces=False) # E only + potential = MACEPotential(spec, use_fallback=True) # CPU / no ops wheel + + A complete, runnable use of this public surface (both variants, free and + periodic, energy-only and energy+forces, with hard-coded fp64 goldens) is + ``regressions/mace-subpackage-restructure-04-potential.py``. + + Args: + spec: Validated variant configuration — + :class:`~molzoo.mace.spec.MACEMatpesSpec` or + :class:`~molzoo.mace.spec.MACEOMolSpec`. + compute_forces: Build the energy **and** force pipeline. Defaults to + ``True``: both foundation checkpoints are used for MD and + evaluation, where the force is the point. + (:class:`~molzoo.pinet.potential.PiNetPotential` defaults to + ``False`` because its main road is energy-only training — the + disagreement is deliberate, not an oversight.) + use_fallback: Put the equivariant blocks on the pure-torch + cuEquivariance path instead of its *fused* kernels — single GPU + kernels that do a whole tensor product at once, and which need a + GPU plus the ``cuequivariance-ops-torch`` wheel. The flag only ever + escalates: ``True`` overrides a spec that asked for fused kernels, + while the ``False`` default never downgrades a spec that already + asked for the fallback (a ``bool`` cannot express "not given"). + Leave it alone on GPU — the foundation force path is always + autograd, so the functorch-compatibility argument for the + pure-torch blocks never applies here, and the fallback measured + **35.7x slower per MD step** with no correctness upside + (``.claude/notes/notes.md`` ``mol:note:topic:cueq-use-fallback``; + the guard that keeps it honest is + ``benchmarks/bench_mace_matpes.py``). CPU and test call sites, + which have no fused kernels available, pass ``True`` explicitly. + """ + + def __init__( + self, + spec: MACEMatpesSpec | MACEOMolSpec, + *, + compute_forces: bool = True, + use_fallback: bool = False, + ) -> None: + super().__init__(spec.model_copy(update={"use_fallback": True}) if use_fallback else spec) + + # ---- frozen switches: no pydantic attribute reads on the hot path ---- + self._conditioning = ( + self._condition_charge_spin + if spec.conditioning == "charge_spin" + else self._reject_conditioning + ) + self._readout_energy = ( + self._per_layer_readout_energy + if spec.readout == "per_layer" + else self._final_readout_energy + ) + #: One bound call path; ``forward`` never branches (cf. PiNetPotential). + self._pipeline = self._pipeline_ef if compute_forces else self._write_energy + #: Element-table gate, spent after the first batch (see ``_write_energy``). + self._elements_validated = False + + @classmethod + def from_checkpoint( + cls, + config_path: str | Path, + weights_path: str | Path, + *, + remap: CheckpointRemap = MATPES_REMAP, + map_location: str | torch.device = "cpu", + **ctor_kwargs, + ) -> MACEPotential: + """Build a MatPES potential from an official config json + cueq weights. + + Three steps, nothing hidden: read the config and translate it into a + :class:`~molzoo.mace.spec.MACEMatpesSpec`, construct the model, load the + weights through ``remap``. It replaced the hand-written construction in + ``run_nve.py``'s ``build_model`` (removed in + ``mace-subpackage-restructure-07-cleanup``, which re-pointed the script + here):: + + potential = MACEPotential.from_checkpoint( + weights_dir / "matpes_r2scan_config.json", + weights_dir / "matpes_r2scan_cueq_state.pt", + ) + potential.eval() # the caller's step, deliberately not done here + + A runnable end-to-end use — a synthetic official-dialect checkpoint + written to a temporary directory, read back through this classmethod, + and checked against hard-coded energy / force goldens — is + ``regressions/mace-subpackage-restructure-05-checkpoint.py``. + + ``num_features`` / ``max_hidden_l`` / ``mlp_dim`` are **derived** from + the config's ``hidden_irreps`` / ``MLP_irreps`` (see + :func:`_parse_irreps`); a missing key raises instead of falling back to + the shipped ``128 / 1 / 16``, which would load a differently sized + checkpoint into the wrong model. The config keys translate as + ``max_ell → l_max``, ``radial_MLP → radial_mlp``, ``atomic_inter_scale + → scale``, ``atomic_inter_shift → shift``; ``r_max``, ``num_bessel``, + ``num_polynomial_cutoff``, ``num_interactions``, ``correlation``, + ``atomic_numbers`` and ``atomic_energies`` keep their names. Nothing on + this path converts units: ``r_max`` is read in Å, ``atomic_energies`` in + eV/atom, ``atomic_inter_shift`` in eV, and the checkpoint's weights are + already in that same system (see :mod:`molzoo.mace.checkpoint`). + + **Validated when present: the two architecture switches the config + spells.** :class:`~molzoo.mace.spec.MACEMatpesSpec` hard-wires + ``pair_repulsion="zbl"`` and ``distance_transform="agnesi"``, so a + config that names either key with a contradicting value is refused here, + before the model is built. It has to be refused at *this* seam because + the load cannot see it: the fitted constants of both blocks are + registered as buffers, not ``nn.Parameter`` (``ZBLRepulsion`` and + ``AgnesiTransform`` are both built with ``trainable=False``), so the + unfilled-parameter guard of + :meth:`~molzoo.mace.checkpoint.CheckpointRemap.load` never fires — the + surplus term would keep its default constants and the model would run, + look sane, and compute a short-range repulsion (or transform every + distance) the checkpoint was never fitted with. ``distance_transform`` + is compared case-insensitively: the stock dump spells it ``"Agnesi"``. + + **Absence is the documented boundary.** A config naming neither key is + accepted on the spec's defaults — older dumps predate the two entries, + and ZBL + Agnesi is what a stock MatPES checkpoint was fitted with + anyway. The remaining variant flags (density interactions, per-layer + readout, conditioning) are likewise taken from the spec and compared + against nothing. For anything but a stock MatPES checkpoint, build the + spec explicitly and call ``MATPES_REMAP.load`` yourself. + + The ``(missing_buffers, unexpected)`` report of + :meth:`~molzoo.mace.checkpoint.CheckpointRemap.load` is dropped here: + under the default policy ``unexpected`` is empty by construction (an + unhoused key raises instead), and ``missing_buffers`` holds only entries + cuEquivariance rebuilds. Pass a ``remap`` with ``on_unexpected="return"`` + and that list is lost — load by hand if you need to inspect it. + + This constructor covers the **MatPES** family only. An OMol checkpoint + is loaded by constructing a :class:`~molzoo.mace.spec.MACEOMolSpec` + explicitly and calling + ``molzoo.mace.checkpoint.OMOL_REMAP.load(potential, state)`` — its + config carries charge/spin fields with no counterpart here. + + Args: + config_path: Official config json (the ``model.config`` dumped + beside the converted weights). + weights_path: ``state_dict`` of the cueq-converted official model, + saved by ``torch.save`` and read back with + ``weights_only=True``. + remap: Key-dialect translation and unexpected-key policy. Defaults + to :data:`~molzoo.mace.checkpoint.MATPES_REMAP`, which refuses + a checkpoint key with no home. + map_location: Device the weights are read onto, forwarded to + ``torch.load``. Defaults to CPU: the model is built on the + ambient device and moved by the caller. + **ctor_kwargs: Passed straight to ``cls`` — ``compute_forces`` / + ``use_fallback`` are runtime-environment switches, not + scientific content of the checkpoint, so they stay out of the + config translation. + + Returns: + A ``MACEPotential`` holding the checkpoint's weights, in training + mode (call ``.eval()`` yourself). + + Raises: + KeyError: If the config lacks a key the spec needs — including + ``hidden_irreps`` / ``MLP_irreps``, which are named explicitly + because no width default is acceptable. + ValueError: If the config contradicts one of the two hard-wired + architecture switches, if an irreps string cannot be parsed, or + if the spec rejects the translated configuration. + RuntimeError: From ``remap`` — an unhoused checkpoint key, a shape + disagreement, or an unfilled parameter. + """ + official = json.loads(Path(config_path).read_text()) + absent = [key for key in ("hidden_irreps", "MLP_irreps") if key not in official] + if absent: + raise KeyError( + f"official config {Path(config_path).name} is missing {absent}: " + "num_features / max_hidden_l / mlp_dim are derived from the irreps " + "strings and have no default" + ) + if "pair_repulsion" in official and not official["pair_repulsion"]: + raise ValueError( + f"official config {Path(config_path).name} sets pair_repulsion=" + f"{official['pair_repulsion']!r}, but MACEMatpesSpec hard-wires " + "pair_repulsion='zbl' and the ZBL constants are buffers, so the load " + "would not object to the surplus term: build the spec explicitly and " + "call MATPES_REMAP.load yourself" + ) + if ( + "distance_transform" in official + and str(official["distance_transform"]).lower() != "agnesi" + ): + raise ValueError( + f"official config {Path(config_path).name} sets distance_transform=" + f"{official['distance_transform']!r}, but MACEMatpesSpec hard-wires " + "distance_transform='agnesi' and the Agnesi constants are buffers, so " + "the load would not object to the surplus transform: build the spec " + "explicitly and call MATPES_REMAP.load yourself" + ) + + num_features, max_hidden_l = _parse_irreps(official["hidden_irreps"], "hidden_irreps") + mlp_dim, _ = _parse_irreps(official["MLP_irreps"], "MLP_irreps") + + spec = MACEMatpesSpec( + atomic_numbers=official["atomic_numbers"], + atomic_energies=official["atomic_energies"], + r_max=official["r_max"], + num_bessel=official["num_bessel"], + num_polynomial_cutoff=official["num_polynomial_cutoff"], + l_max=official["max_ell"], + num_features=num_features, + max_hidden_l=max_hidden_l, + num_interactions=official["num_interactions"], + correlation=official["correlation"], + mlp_dim=mlp_dim, + radial_mlp=official["radial_MLP"], + scale=official["atomic_inter_scale"], + shift=official["atomic_inter_shift"], + ) + + model = cls(spec, **ctor_kwargs) + remap.load(model, torch.load(weights_path, map_location=map_location, weights_only=True)) + return model + + def forward(self, batch: TensorDict) -> TensorDict: + """Run the construction-fixed pipeline, mutating ``batch`` in place. + + Reads ``atoms.{Z,pos,batch}``, ``edges.edge_index`` ``(E, 2)`` (``[:,0]`` + source, ``[:,1]`` target — never transposed on the way in), optional + ``edges.shifts`` ``(E, 3)`` and, for the OMOL variant, optional + ``graphs.{total_charge,total_spin}``. Writes ``graphs.energy`` ``(B,)`` + in eV and — unless the instance was built with ``compute_forces=False`` + — ``atoms.forces`` ``(N, 3)`` in eV/Å. + + Args: + batch: Post-collate ``TensorDict``. + + Returns: + The same ``batch`` object, with the outputs written in place. + + Raises: + ValueError: If the batch contains an atomic number outside this + model's element table (checked once per instance), or if it + carries ``graphs.total_charge`` / ``graphs.total_spin`` for a + variant that has no charge/spin conditioning. + RuntimeError: From the force kernel, if forces were requested but + the energy did not reach ``graphs.energy``. + """ + return self._pipeline(batch) + + # ------------------------------------------------------------ energy core + def energy_core( + self, + positions: torch.Tensor, + Z: torch.Tensor, + edge_index: torch.Tensor, + batch: torch.Tensor, + num_graphs: int, + shifts: torch.Tensor | None = None, + total_charge: torch.Tensor | None = None, + total_spin: torch.Tensor | None = None, + ) -> torch.Tensor: + """Per-graph total energy ``(B,)`` as a pure function of ``positions``. + + Evaluates the ``E_g`` of the module docstring: isolated-atom reference + energies, plus the scale/shift-normalised sum of the readouts (and the + ZBL pair term, where the variant has one), summed over the atoms of + each graph. + + This is the flat compile seam: raw tensors in, energy out, no + ``TensorDict`` and no host sync. Every quantity that depends on the + positions is recomputed here — the edge displacement vectors, their + lengths, the angular (spherical-harmonic) features of their directions + and the radial basis expansion of their lengths — so the result is a + differentiable function of ``positions`` (``torch.autograd.grad`` gives + the forces) and traces into a single TorchDynamo graph. + + The first six parameters keep the names and positions of the flat + MatPES model's ``_compute_energy`` + (``src/molzoo/mace_matpes.py:229-237`` at 0e05959, before deletion) so + those call sites move over by a pure rename. The flat OMOL model's core + (``src/molzoo/mace_omol.py:228-237`` at 0e05959) took ``total_charge`` / + ``total_spin`` positionally instead of ``num_graphs`` and does **not** + line up. + + Args: + positions: Atom positions ``(N, 3)`` in Å. + Z: Atomic numbers ``(N,)``; must lie inside the element table (see + :meth:`~molzoo.mace.encoder.MACEEncoder.validate_elements`). + edge_index: ``(E, 2)`` with ``[:, 0]`` = source, ``[:, 1]`` = target. + batch: Graph index per atom ``(N,)`` — ``batch[i]`` says which of + the ``B`` graphs atom ``i`` belongs to. + num_graphs: Number of graphs ``B``. An argument, not + ``int(batch.max().item()) + 1``: that host sync would break the + traced graph and stall the MD hot path. + shifts: Optional periodic shift vectors ``(E, 3)`` in Å + (``unit_shifts @ cell``), added to the edge displacements so a + neighbour across a periodic boundary enters at its imaged + position. Constant w.r.t. ``positions``, so the forces stay + exact. + total_charge: Per-graph total charge ``(B,)`` in units of the + elementary charge ``e`` (OMOL only; defaults to neutral). Used + as an integer index into an embedding table, offset by the + spec's ``charge_offset``. + total_spin: Per-graph total spin ``(B,)``, dimensionless (OMOL + only; defaults to ``1``, the closed-shell singlet — every + electron paired — so the values follow the spin-multiplicity + ``2S+1`` convention, not a count of unpaired electrons). Also + an embedding-table index, offset by ``spin_offset``. + + Returns: + Total energy per graph ``(B,)`` in eV. + + Raises: + ValueError: If a variant without charge/spin conditioning is handed + ``total_charge`` or ``total_spin`` — silently ignoring them + would return a plausible, wrong energy. + """ + vectors = edge_vectors(positions, edge_index, shifts) + lengths = edge_lengths(vectors) + + node_attrs = self.node_attrs(Z, positions.dtype) + edge_feats, cutoff = self.radial_features(lengths, Z, edge_index) + edge_attrs = self.angular_features(vectors) + + node_feats = self.initial_node_features(node_attrs) + e0 = _scatter_sum(self.atomic_energies(Z), batch, num_graphs) + node_feats, e0 = self._conditioning( + node_feats, e0, batch, num_graphs, total_charge, total_spin + ) + + per_layer = self.layer_features( + node_feats=node_feats, + node_attrs=node_attrs, + edge_attrs=edge_attrs, + edge_feats=edge_feats, + edge_index=edge_index, + cutoff=cutoff, + ) + # The short-range pair term is a per-atom energy on the same footing as + # the readouts, and it is summed first (MatPES operator order). + node_es = ( + self.pair_repulsion(lengths, Z, edge_index) + if self.pair_repulsion is not None + else node_feats.new_zeros(node_feats.shape[0]) + ) + node_es = self._readout_energy(per_layer, node_es) + return e0 + _scatter_sum(self.scale_shift(node_es), batch, num_graphs) + + # -------------------------------------------------------------- variants + def _reject_conditioning( + self, + node_feats: torch.Tensor, + e0: torch.Tensor, + batch: torch.Tensor, + num_graphs: int, + total_charge: torch.Tensor | None, + total_spin: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Pass the node state through — this variant has no charge/spin input. + + The ``None`` comparisons are constant-folded by dynamo, so the guard + costs no graph break on the compiled path. + + Args: + node_feats: Initial node features ``(N, num_features)``. + e0: Per-graph reference energy ``(B,)`` in eV. + batch: Graph index per atom ``(N,)`` (unused). + num_graphs: Number of graphs ``B`` (unused). + total_charge: Must be ``None``. + total_spin: Must be ``None``. + + Returns: + ``(node_feats, e0)`` unchanged. + + Raises: + ValueError: If either conditioning tensor was supplied. + """ + if total_charge is not None or total_spin is not None: + raise ValueError( + "this MACE variant has no charge/spin conditioning " + '(conditioning="none"): total_charge / total_spin would be ' + "ignored and the energy would be quietly wrong — build the " + "potential from a MACEOMolSpec instead" + ) + return node_feats, e0 + + def _condition_charge_spin( + self, + node_feats: torch.Tensor, + e0: torch.Tensor, + batch: torch.Tensor, + num_graphs: int, + total_charge: torch.Tensor | None, + total_spin: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Add the OMOL charge/spin embedding and its scalar readout to ``e0``. + + Args: + node_feats: Initial node features ``(N, num_features)``. + e0: Per-graph reference energy ``(B,)`` in eV. + batch: Graph index per atom ``(N,)``. + num_graphs: Number of graphs ``B``. + total_charge: Per-graph total charge ``(B,)`` in units of ``e``; + ``None`` → neutral. + total_spin: Per-graph total spin ``(B,)``, dimensionless; ``None`` + → ``1``, the closed-shell singlet (all electrons paired, i.e. + multiplicity ``2S+1 = 1``). Spin ``0`` would index an untrained + embedding row and return garbage + (``src/molzoo/mace_omol.py:330-335`` at 0e05959, before + deletion). + + Returns: + ``(conditioned node features (N, num_features), e0 (B,) eV)``. + """ + if total_charge is None: + total_charge = torch.zeros(num_graphs, dtype=torch.long, device=node_feats.device) + if total_spin is None: + total_spin = torch.ones(num_graphs, dtype=torch.long, device=node_feats.device) + node_feats = node_feats + self.conditioning( + batch, total_spin=total_spin, total_charge=total_charge + ) + embedded = self.embedding_readout(node_feats).squeeze(-1) + return node_feats, e0 + _scatter_sum(embedded, batch, num_graphs) + + def _per_layer_readout_energy( + self, per_layer: list[torch.Tensor], node_es: torch.Tensor + ) -> torch.Tensor: + """Sum one readout per interaction layer (MACE-MP / MatPES). + + Args: + per_layer: Node features of every layer, from + :meth:`~molzoo.mace.encoder.MACEEncoder.layer_features`. + node_es: Per-atom energy accumulated so far ``(N,)`` in eV. + + Returns: + Per-atom interaction energy ``(N,)`` in eV, before scale/shift. + """ + for readout, node_feats in zip(self.readouts, per_layer, strict=True): + node_es = node_es + readout(node_feats).squeeze(-1) + return node_es + + def _final_readout_energy( + self, per_layer: list[torch.Tensor], node_es: torch.Tensor + ) -> torch.Tensor: + """Read out the last layer only (OMOL). + + Args: + per_layer: Node features of every layer, from + :meth:`~molzoo.mace.encoder.MACEEncoder.layer_features`. + node_es: Per-atom energy accumulated so far ``(N,)`` in eV. + + Returns: + Per-atom interaction energy ``(N,)`` in eV, before scale/shift. + """ + return node_es + self.readout(per_layer[-1]).squeeze(-1) + + # ------------------------------------------------------------- pipelines + def _write_energy(self, batch: TensorDict) -> TensorDict: + """Energy core on the batch schema — the ``protocol.call_energy`` hook. + + Also the whole pipeline of a ``compute_forces=False`` instance. + Deliberately detaches **nothing**: an outer ``EnergyReadout`` / + ``ForceReadout`` session may own the position leaf this energy hangs + off, and detaching would cut ``F = -∂E/∂r``. + + The ``graphs`` namespace is created here — ``TensorDict({}, + batch_size=[num_graphs])`` — *before* ``protocol.write_energy`` runs, so + the post-collate schema (``graphs`` is ``batch_size=[B]``, see + CLAUDE.md) holds from this ``B``, the one this method already resolved. + It used to be load-bearing: ``protocol.ensure_graphs`` built + ``batch_size=[]`` and silently degraded the schema until + ``mace-subpackage-restructure-07-cleanup`` taught it ``num_graphs`` and + had ``write_energy`` pass ``energy.shape[0]``. The two now agree, and + this line stays as the explicit statement of the shape. + + Args: + batch: Post-collate ``TensorDict`` carrying ``atoms.{Z,pos,batch}`` + and ``edges.edge_index`` ``(E, 2)``. Optional and consumed when + present: ``edges.shifts`` ``(E, 3)`` in Å (periodic images) and, + for the OMOL variant, ``graphs.{total_charge,total_spin}`` + ``(B,)``. A ``graphs`` namespace that is already there supplies + ``B`` from its ``batch_size``; without one, ``B`` costs the one + host sync on this path (``forward`` is not the compile target, + :meth:`energy_core` is). + + Returns: + The same ``batch``, with ``graphs.energy`` ``(B,)`` in eV written. + + Raises: + ValueError: On the first call, if any atomic number lies outside + this model's element table; on every call, if a variant without + charge/spin conditioning was handed ``graphs.total_charge`` or + ``graphs.total_spin``. + """ + Z = batch["atoms", "Z"] + positions = batch[POS_KEY] + atom_batch = batch["atoms", "batch"] + edge_index = batch["edges", "edge_index"] # (E, 2) — the core's own convention + + # Once per instance: Z is constant over a trajectory and a wrongly + # wired model/dataset pair fails on the very first batch, while a + # per-call check costs a host sync and a dynamo graph break. + if not self._elements_validated: + self.validate_elements(Z) + self._elements_validated = True + + nested = batch.keys(include_nested=True) + has_graphs = "graphs" in batch.keys() + if has_graphs and batch["graphs"].batch_size: + num_graphs = int(batch["graphs"].batch_size[0]) + else: + # The only host sync on this path; ``forward`` is not the compile + # target — ``energy_core`` is, and it takes ``num_graphs`` directly. + num_graphs = int(atom_batch.max().item()) + 1 + + energy = self.energy_core( + positions, + Z, + edge_index, + atom_batch, + num_graphs, + shifts=batch["edges", "shifts"] if ("edges", "shifts") in nested else None, + total_charge=( + batch["graphs", "total_charge"] if ("graphs", "total_charge") in nested else None + ), + total_spin=( + batch["graphs", "total_spin"] if ("graphs", "total_spin") in nested else None + ), + ) + + if not has_graphs: + batch["graphs"] = TensorDict({}, batch_size=[num_graphs]) + write_energy(batch, energy) + return batch + + def _pipeline_ef(self, batch: TensorDict) -> TensorDict: + """Energy + forces: one energy forward, one ``autograd.grad`` backward. + + Delegates to the shared batch-level kernel + (:func:`molpot.derivation.kernels.grad_force_pass`) — the autograd + backend, the only one cuEquivariance's fused kernels support. This + module hand-rolls no force path of its own. + + ``detach_energy=None`` is the ``needs_leaf`` policy the flat models + implemented by hand: the kernel detaches ``graphs.energy`` exactly when + it had to create the position leaf itself, and leaves it attached when + the caller handed in a live ``requires_grad`` leaf, so a training + ``loss.backward()`` still reaches the parameters. ``create_graph`` + follows the ambient grad mode. + + Args: + batch: Post-collate ``TensorDict`` carrying ``atoms.pos``. + + Returns: + The same ``batch``, with ``graphs.energy`` ``(B,)`` in eV and + ``atoms.forces`` ``(N, 3)`` in eV/Å. + """ + return grad_force_pass(self._write_energy, batch, detach_energy=None) diff --git a/src/molzoo/mace/research.py b/src/molzoo/mace/research.py new file mode 100644 index 0000000..763721f --- /dev/null +++ b/src/molzoo/mace/research.py @@ -0,0 +1,330 @@ +"""MACE: Multi-Atomic Cluster Expansion research encoder. + +Equivariant message-passing encoder that produces per-layer node features. +Downstream readout, classical potential terms, and force derivation are +handled outside this module. + +This is the *research* encoder (freely configurable node-attribute embeddings, +element updates and layer norms), as opposed to the foundation-model backbone +in :mod:`molzoo.mace.encoder`. Its configuration model is therefore named +:class:`MACEResearchSpec`; ``molzoo.mace.MACESpec`` now refers to the shared +foundation-variant base in :mod:`molzoo.mace.spec`. + +Example: + >>> from molzoo import MACE + >>> from molrep.embedding.node import DiscreteEmbeddingSpec + >>> encoder = MACE( + ... node_attr_specs=[DiscreteEmbeddingSpec( + ... input_key="Z", num_classes=119, emb_dim=64)], + ... num_elements=118, + ... num_features=128, + ... r_max=5.0, + ... ) + >>> features = encoder( + ... Z=Z, + ... edge_dist=edge_dist, + ... edge_diff=edge_diff, + ... edge_index=edge_index, + ... ) + >>> print(features.shape) # (n_nodes, num_layers, num_features) + +Reference: + Batatia et al. "MACE: Higher Order Equivariant Message Passing Neural + Networks for Fast and Accurate Force Fields" NeurIPS 2022 + https://arxiv.org/abs/2206.07697 +""" + +from __future__ import annotations + +import cuequivariance as cue +import torch +import torch.nn as nn +from cuequivariance import O3, Irreps +from pydantic import BaseModel, ConfigDict, Field +from tensordict import TensorDict +from tensordict.nn import TensorDictModuleBase + +from molix import config +from molrep.embedding.mace import EmbeddingBlock, EmbeddingSpec +from molrep.embedding.node import ( + ContinuousEmbeddingSpec, + DiscreteEmbeddingSpec, +) +from molrep.interaction.element import ElementUpdate +from molrep.interaction.mace.block import InteractionBlock, InteractionSpec +from molrep.interaction.product import irreps_from_l_max +from molrep.readout.product import ProductHead + +#: Blocks promoted to ``molrep`` by mace-subpackage-restructure-01 and re-exported +#: here so ``from molzoo.mace import EmbeddingBlock`` keeps resolving. Kept while +#: ``tests/test_molzoo/test_imports.py::TestMolzooMaceReexports`` pins them to +#: their promoted ``molrep`` homes by object identity. +__all__ = [ + "EmbeddingBlock", + "EmbeddingSpec", + "InteractionBlock", + "InteractionSpec", + "MACE", + "MACEResearchSpec", +] + + +# =========================================================================== +# MACE Encoder (Feature Extractor) +# =========================================================================== + + +class MACE(TensorDictModuleBase): + """MACE equivariant feature encoder. + + Accepts a ``TensorDict`` TensorDict and writes ``node_features`` + into the ``atoms`` sub-dict in place, returning the same + ``TensorDict`` with the new key added. + + Architecture:: + + TensorDict(atoms, edges) + → [Embedding] → node_feats, edge_attrs, edge_feats + → [Interaction₁] → [ProductHead₁] → [ElementUpdate₁] + → ... + → [Interactionₙ] → [ProductHeadₙ] + → atoms.node_features (n_nodes, num_interactions, num_features) + + Reference: + Batatia et al. "MACE: Higher Order Equivariant Message Passing Neural + Networks for Fast and Accurate Force Fields" NeurIPS 2022 + https://arxiv.org/abs/2206.07697 + """ + + in_keys = [ + ("atoms", "Z"), + ("atoms", "pos"), + ("edges", "edge_index"), + ("edges", "edge_diff"), + ("edges", "edge_dist"), + ] + out_keys = [("atoms", "node_features")] + + def __init__( + self, + *, + node_attr_specs: list[DiscreteEmbeddingSpec | ContinuousEmbeddingSpec], + num_elements: int, + num_features: int, + r_max: float, + num_bessel: int = 8, + l_max: int = 2, + num_interactions: int = 2, + correlation: int = 2, + avg_num_neighbors: float = 1.0, + layer_norm: bool = False, + use_fallback: bool = True, + ): + """Initialize MACE feature extractor. + + Args: + node_attr_specs: Embedding specs for node attributes (e.g. Z). + num_elements: Number of atomic element types. + num_features: Scalar channel multiplicity at l=0. + r_max: Radial cutoff in Angstroms. + num_bessel: Number of Bessel radial basis functions. + l_max: Maximum angular momentum order. + num_interactions: Number of interaction-product-update layers. + correlation: Body-order correlation for symmetric contraction. + avg_num_neighbors: Average neighbor count for message normalization. + layer_norm: Whether to apply layer normalization between layers. + use_fallback: Pure-torch cuEq path (default ``True``, functorch-safe + for ``ForceDerivation(method="functorch")``). Set ``False`` for + the fused kernels when forces use the autograd backend — the + tensor product and symmetric contraction are the encoder's two + hottest blocks. + """ + super().__init__() + + # Frozen on the hot path: reading ``self.config.num_interactions`` in + # ``forward`` is a pydantic attribute lookup inside the layer loop, i.e. + # a dynamo graph break (cf. ``src/molzoo/mace_matpes.py:118`` at + # 0e05959, before deletion). ``self.config`` is kept for provenance + # only and must not be read by ``forward``. + self.num_interactions = num_interactions + + self.config = MACEResearchSpec( + node_attr_specs=node_attr_specs, + num_elements=num_elements, + num_features=num_features, + r_max=r_max, + num_bessel=num_bessel, + l_max=l_max, + num_interactions=num_interactions, + correlation=correlation, + avg_num_neighbors=avg_num_neighbors, + layer_norm=layer_norm, + use_fallback=use_fallback, + ) + + # Embedding + self.embedding = EmbeddingBlock( + node_attr_specs=node_attr_specs, + num_features=num_features, + r_max=r_max, + num_bessel=num_bessel, + l_max=l_max, + ) + # Mixed-l message dimension (transient TP output consumed by ProductHead) + irreps_str = irreps_from_l_max(l_max, num_features) + with cue.assume(O3): + irreps_dim = Irreps(irreps_str).dim + + # The node *state* carried between layers is pure scalar (num_features); + # only the per-edge messages are mixed-l. This keeps every node-state + # op equivariant. Initial projection is therefore scalar -> scalar. + self.initial_projection = nn.Linear(num_features, num_features, dtype=config.ftype) + + # Interaction blocks + self.interactions = nn.ModuleList( + [ + InteractionBlock( + num_features=num_features, + num_bessel=num_bessel, + l_max=l_max, + avg_num_neighbors=avg_num_neighbors, + use_fallback=use_fallback, + ) + for _ in range(num_interactions) + ] + ) + + # Product heads (from molrep, replaces former ProductBlock) + self.products = nn.ModuleList( + [ + ProductHead( + hidden_dim=irreps_dim, + out_dim=num_features, + num_radial=num_bessel, + l_max=l_max, + max_body_order=correlation, + num_species=num_elements, + use_fallback=use_fallback, + ) + for _ in range(num_interactions) + ] + ) + + # Projection of the (scalar) product readout back into the scalar node + # state for the residual path. Scalar -> scalar keeps it equivariant. + self.projections = nn.ModuleList( + [ + nn.Linear(num_features, num_features, dtype=config.ftype) + for _ in range(num_interactions) + ] + ) + + # Element-specific residual updates (all layers except last). Operates on + # the scalar node state, so ElementUpdate's scalar (l=0) treatment is now + # correct rather than silently mixing l>0 components. + self.element_updates = nn.ModuleList( + [ + ElementUpdate(hidden_dim=num_features, num_species=num_elements) + for _ in range(max(num_interactions - 1, 0)) + ] + ) + + # Layer normalization (all layers except last) over the scalar state. + self.layer_norms = nn.ModuleList( + [ + nn.LayerNorm(num_features, dtype=config.ftype) if layer_norm else nn.Identity() + for _ in range(max(num_interactions - 1, 0)) + ] + ) + + def forward(self, td: TensorDict) -> TensorDict: + """Extract per-layer geometric features. + + Args: + td: ``TensorDict`` with ``atoms`` and ``edges`` sub-dicts. + + Returns: + Same ``TensorDict`` with ``atoms.node_features`` + ``(n_nodes, num_interactions, num_features)`` added. + """ + Z = td["atoms", "Z"] + edge_dist = td["edges", "edge_dist"] + edge_diff = td["edges", "edge_diff"] + edge_index = td["edges", "edge_index"] + + # ---- Embedding ---- + node_feats_init, edge_attrs, edge_feats = self.embedding( + Z=Z, + edge_dist=edge_dist, + edge_diff=edge_diff, + ) + + # ---- Initial projection: scalar embeddings -> hidden irreps ---- + node_feats = self.initial_projection(node_feats_init) + + # ---- Interaction-Product-Update loop ---- + per_layer_features: list[torch.Tensor] = [] + + for i in range(self.num_interactions): + node_feats_msg, sc = self.interactions[i]( + node_feats=node_feats, + edge_attrs=edge_attrs, + edge_feats=edge_feats, + edge_index=edge_index, + ) + + h_product = self.products[i]( + node_features=node_feats_msg, + atom_types=Z, + ) + + per_layer_features.append(h_product) + + h_proj = self.projections[i](h_product) + + is_last = i == (self.num_interactions - 1) + if not is_last: + node_feats = self.element_updates[i]( + h_prev=sc, + m_curr=h_proj, + atom_types=Z, + ) + node_feats = self.layer_norms[i](node_feats) + else: + node_feats = h_proj + + td["atoms", "node_features"] = torch.stack(per_layer_features, dim=1) + return td + + +class MACEResearchSpec(BaseModel): + """Configuration for the MACE feature extractor. + + Attributes: + node_attr_specs: Embedding specs for node attributes. + num_elements: Number of atomic element types. + num_features: Scalar channel multiplicity. + r_max: Radial cutoff in Angstroms. + num_bessel: Number of Bessel basis functions. + l_max: Maximum angular momentum order. + num_interactions: Number of interaction-product layers. + correlation: Body-order correlation for symmetric contraction. + avg_num_neighbors: Average neighbor count for normalization. + layer_norm: Whether to apply layer normalization. + """ + + model_config = ConfigDict(arbitrary_types_allowed=True) + + node_attr_specs: list[DiscreteEmbeddingSpec | ContinuousEmbeddingSpec] = Field( + ..., min_length=1 + ) + num_elements: int = Field(..., gt=0) + num_features: int = Field(..., gt=0) + r_max: float = Field(..., gt=0.0) + num_bessel: int = Field(8, gt=0) + l_max: int = Field(2, ge=0) + num_interactions: int = Field(2, gt=0) + correlation: int = Field(2, ge=1, le=3) + avg_num_neighbors: float = Field(1.0, gt=0.0) + layer_norm: bool = False + use_fallback: bool = True diff --git a/src/molzoo/mace/spec.py b/src/molzoo/mace/spec.py new file mode 100644 index 0000000..fecafe2 --- /dev/null +++ b/src/molzoo/mace/spec.py @@ -0,0 +1,201 @@ +"""Validated configurations for the MACE foundation-model variants. + +Three pydantic models: :class:`MACESpec` carries everything the two shipped +variants share, :class:`MACEMatpesSpec` and :class:`MACEOMolSpec` add the +variant-only fields and pin the defaults of the constructors they replace +(``MACEMatpes.__init__`` / ``MACEOMol.__init__``). + +Units follow the rest of the repo: positions and ``r_max`` in Å, energies +(``atomic_energies``, ``shift``) in eV, ``atomic_energies`` per atom. + +This module deliberately imports **nothing** but ``typing`` and ``pydantic``. +A configuration must stay readable — from a CLI, a checkpoint manifest, or a +test — without paying for torch and the cuEquivariance stack, so tensor +construction (``atomic_energies`` → buffer) belongs to +:class:`~molzoo.mace.encoder.MACEEncoder`, not here. + +Reference: + Batatia et al. "MACE: Higher Order Equivariant Message Passing Neural + Networks for Fast and Accurate Force Fields" NeurIPS 2022. + https://arxiv.org/abs/2206.07697 + Batatia et al. "A foundation model for atomistic materials chemistry" + (MACE-MP-0). https://arxiv.org/abs/2401.00096 + Kaplan et al. "A foundational potential energy surface dataset for + materials" (MatPES). https://arxiv.org/abs/2503.04070 + Levine et al. "The Open Molecules 2025 (OMol25) Dataset, Evaluations, and + Models" https://arxiv.org/abs/2505.08762 +""" + +from __future__ import annotations + +from typing import Literal + +from pydantic import BaseModel, ConfigDict, Field, model_validator + + +class MACESpec(BaseModel): + """Shared configuration of the MACE foundation-model backbone. + + The five ``Literal`` switches at the bottom select which stack + :class:`~molzoo.mace.encoder.MACEEncoder` builds. They have no default on + this base class — a bare ``MACESpec`` describes no shipped model, so the + caller must say which one it means (the subclasses do exactly that). + + Attributes: + atomic_numbers: Element table (z-table), strictly ascending. The + encoder looks ``Z`` up with ``torch.searchsorted``, which silently + snaps to a neighbouring row on an unordered table. + atomic_energies: Per-element reference energies ``E0`` in eV/atom, in + ``atomic_numbers`` order. A plain ``list`` — see the module + docstring for why the tensor is built by the encoder. + r_max: Radial cutoff in Å. + num_bessel: Number of Bessel radial basis functions. + num_polynomial_cutoff: Polynomial cutoff exponent ``p`` (also the ZBL + envelope exponent, matching MACE). + l_max: Maximum spherical-harmonics order. + num_features: Scalar channel multiplicity (``hidden_irreps`` 0e count). + num_interactions: Number of interaction/product layers. + correlation: Body-order correlation of the symmetric contraction. + mlp_dim: Hidden width of the final non-linear readout. + scale: ``atomic_inter_scale`` — multiplies the interaction energy (eV). + shift: ``atomic_inter_shift`` — per-atom energy offset in eV. + use_fallback: Pure-torch cuEquivariance path (``True``) or fused + kernels (``False``). Fused kernels need a GPU and the + ``cuequivariance-ops-torch`` wheel. + interaction: ``"density"`` (MACE-MP/MatPES density-normalised + interactions) or ``"residual"`` (OMOL non-linear residual ones). + readout: ``"per_layer"`` (one readout per layer, summed) or + ``"final"`` (a single readout on the last layer). + distance_transform: ``"agnesi"`` applies the Agnesi radial transform + before the Bessel basis; ``"none"`` feeds raw distances. + pair_repulsion: ``"zbl"`` adds the ZBL short-range pair term. + conditioning: ``"charge_spin"`` adds the total-charge / total-spin + joint embedding to the initial node features. + """ + + model_config = ConfigDict(arbitrary_types_allowed=True) + + atomic_numbers: list[int] = Field(..., min_length=1) + atomic_energies: list[float] = Field(..., min_length=1) + + r_max: float = Field(6.0, gt=0.0) + num_bessel: int = Field(8, gt=0) + num_polynomial_cutoff: int = Field(5, gt=0) + l_max: int = Field(3, ge=0) + num_features: int = Field(128, gt=0) + num_interactions: int = Field(2, gt=0) + correlation: int = Field(2, ge=1) + mlp_dim: int = Field(16, gt=0) + scale: float = 1.0 + shift: float = 0.0 + use_fallback: bool = False + + interaction: Literal["density", "residual"] + readout: Literal["per_layer", "final"] + distance_transform: Literal["none", "agnesi"] + pair_repulsion: Literal["none", "zbl"] + conditioning: Literal["none", "charge_spin"] + + @model_validator(mode="after") + def _check_element_table(self) -> MACESpec: + """Reject element tables that would mis-index ``E0`` or the one-hot. + + Raises: + ValueError: If ``atomic_energies`` has a different length than + ``atomic_numbers``, or ``atomic_numbers`` is not strictly + ascending (unordered or duplicated). + """ + if len(self.atomic_energies) != len(self.atomic_numbers): + raise ValueError( + f"atomic_energies has {len(self.atomic_energies)} entries but " + f"atomic_numbers has {len(self.atomic_numbers)} — one E0 per element" + ) + pairs = zip(self.atomic_numbers, self.atomic_numbers[1:]) + if any(previous >= following for previous, following in pairs): + raise ValueError( + f"atomic_numbers must be strictly ascending (searchsorted lookup), " + f"got {self.atomic_numbers}" + ) + return self + + +class MACEMatpesSpec(MACESpec): + """Configuration of the MACE-MatPES foundation model. + + Defaults reproduce ``MACEMatpes.__init__`` of the pre-cutover flat module + (``src/molzoo/mace_matpes.py``, deleted in 06-wire; see git history), + including its materialised ``radial_mlp`` of ``[64, 64, 64]``. + + Attributes: + max_hidden_l: Highest ``l`` carried in the node state between layers + (1 for the shipped MatPES models, i.e. ``128x0e+128x1o``). + radial_mlp: Hidden widths of the radial weight MLP. + """ + + max_hidden_l: int = Field(1, ge=0) + radial_mlp: list[int] = Field(default_factory=lambda: [64, 64, 64], min_length=1) + + num_bessel: int = Field(10, gt=0) + correlation: int = Field(3, ge=1) + + interaction: Literal["density", "residual"] = "density" + readout: Literal["per_layer", "final"] = "per_layer" + distance_transform: Literal["none", "agnesi"] = "agnesi" + pair_repulsion: Literal["none", "zbl"] = "zbl" + conditioning: Literal["none", "charge_spin"] = "none" + + @model_validator(mode="after") + def _check_residual_layer(self) -> MACEMatpesSpec: + """MatPES's stack is one density layer plus at least one residual one. + + Raises: + ValueError: If ``num_interactions`` is below 2. + """ + if self.num_interactions < 2: + raise ValueError(f"num_interactions must be at least 2, got {self.num_interactions}") + return self + + +class MACEOMolSpec(MACESpec): + """Configuration of the MACE-OMOL foundation model. + + Defaults reproduce ``MACEOMol.__init__`` of the pre-cutover flat module + (``src/molzoo/mace_omol.py``, deleted in 06-wire; see git history). + ``use_fallback`` is not a constructor argument there — the flat model + hard-codes the fused cuEquivariance path, i.e. ``False``. + + Attributes: + edge_channels: Per-``l`` channel count of the mid-layer edge irreps and + the radial-MLP hidden width (a deliberate bottleneck below + ``num_features``). + charge_classes: Embedding rows for the total-charge conditioning. + charge_offset: Index offset applied to total charge (charge −100 → row 0). + spin_classes: Embedding rows for the total-spin conditioning. + spin_offset: Index offset applied to total spin. + """ + + edge_channels: int = Field(128, gt=0) + charge_classes: int = Field(201, gt=0) + charge_offset: int = Field(100, ge=0) + spin_classes: int = Field(101, gt=0) + spin_offset: int = Field(0, ge=0) + + num_features: int = Field(1024, gt=0) + num_interactions: int = Field(3, gt=0) + + interaction: Literal["density", "residual"] = "residual" + readout: Literal["per_layer", "final"] = "final" + distance_transform: Literal["none", "agnesi"] = "none" + pair_repulsion: Literal["none", "zbl"] = "none" + conditioning: Literal["none", "charge_spin"] = "charge_spin" + + @model_validator(mode="after") + def _check_angular_order(self) -> MACEOMolSpec: + """OMOL's mid-layer edge irreps span ``range(l_max)`` — empty at ``l_max=0``. + + Raises: + ValueError: If ``l_max`` is below 1. + """ + if self.l_max < 1: + raise ValueError(f"l_max must be at least 1, got {self.l_max}") + return self diff --git a/src/molzoo/mace/variants.py b/src/molzoo/mace/variants.py new file mode 100644 index 0000000..e58c0e6 --- /dev/null +++ b/src/molzoo/mace/variants.py @@ -0,0 +1,545 @@ +"""Named MACE foundation models: the keyword surface the scripts already bind. + +A *foundation* model is one shipped with weights already fitted on a large, +chemically broad dataset; molnex carries two, MatPES (periodic materials, +arXiv:2503.04070) and OMol (molecules, additionally conditioned on the total +charge and total spin of the system). Everything that *defines* them lives +elsewhere — the hyper-parameters and architecture switches in +:mod:`molzoo.mace.spec`, the module graph in :mod:`molzoo.mace.encoder`, the +energy and forces in :mod:`molzoo.mace.potential`, the official-checkpoint key +dialect in :mod:`molzoo.mace.checkpoint`. This module adds **no** physics, no +hyper-parameter, no *irreps* string and no ``E0`` table of its own: it is the +adapter that lets the pre-restructure call sites keep working unchanged. (An +*irrep*, short for irreducible representation, labels how a feature transforms +when the molecule is rotated — a scalar stays put, a vector turns with it — and +an irreps string such as ``128x0e+128x1o`` says how many features of each kind +a layer carries; ``E0`` is the table of isolated-atom reference energies, one +per element, in eV/atom.) + +One such call site is still real and out of this spec's reach:: + + benchmarks/bench_mace_matpes.py:41 MACEMatpes(atomic_numbers=…, …) + +``scripts/matpes_port/run_nve.py`` was the second until +``mace-subpackage-restructure-07-cleanup`` collapsed its ``build_model`` onto +:meth:`~molzoo.mace.potential.MACEPotential.from_checkpoint`, which builds the +spec itself. So :class:`MACEMatpes` and :class:`MACEOMol` stay constructible by +keyword, and :func:`load_matpes_state_dict` / :func:`load_omol_state_dict` stay +importable as free functions. Both loaders are *only* here for that +compatibility — CLAUDE.md's "no factory functions" rule would otherwise reject +them, and they forward, unchanged, to the two preset +:class:`~molzoo.mace.checkpoint.CheckpointRemap` instances. + +**New code should not use anything in this module.** Build the configuration +and the potential directly, and load a checkpoint through the remap:: + + from molzoo.mace import MACEMatpesSpec, MACEPotential + from molzoo.mace.checkpoint import MATPES_REMAP + + potential = MACEPotential(MACEMatpesSpec(atomic_numbers=…, atomic_energies=…)) + MATPES_REMAP.load(potential, state) + +That form says which variant it means in the *type* of the spec, takes the +hyper-parameters from one validated place, and is what +:meth:`~molzoo.mace.potential.MACEPotential.from_checkpoint` already does for a +stock MatPES checkpoint. + +The keyword vocabulary, once, for both classes below +---------------------------------------------------- + +MACE describes each atom by the neighbours inside a cutoff radius ``r_max`` +(Å) and refines that description over ``num_interactions`` rounds of *message +passing* — one round lets an atom mix in features of its neighbours. Each +neighbour enters through two expansions: its distance in ``num_bessel`` +*Bessel* radial functions (the shapes ``sin(ω_n r) / r``, with ``r`` the +interatomic distance in Å and the frequencies ``ω_n`` in Å⁻¹), and its +direction in *spherical harmonics* — the standard angular functions on the +sphere — up to angular order ``l_max`` (order 0 is a plain scalar, 1 turns like +a vector, and so on). ``num_features`` is how many rotation-invariant numbers +each atom carries, ``max_hidden_l`` the highest angular order kept in the node +state between layers, and ``correlation`` the *body order* of the symmetric +contraction — how many neighbours may enter one term of the many-body +expansion. ``num_polynomial_cutoff`` is the exponent ``p`` of the polynomial +envelope that takes an edge's weight smoothly to zero as its length approaches +``r_max``. ``scale`` and ``shift`` are the fitted affine normalisation of the +learned per-atom energy (eV), on top of which the ``E0`` table contributes the +isolated-atom references; the full energy expression, its ``F_i = -∂E/∂r_i`` +force convention and its units are written out in +:mod:`molzoo.mace.potential`. + +Resolved debt — the private energy core +--------------------------------------- + +``scripts/matpes_port/run_nve.py`` used to bind ``model._compute_energy`` (and +``benchmarks/bench_mace_matpes.py`` to hand it to ``Compiler``) as the compiled +energy core of the MD loop — a **private** method crossing a package boundary, +discovered by the ``mace-subpackage-restructure`` cutover. Both consumers were +re-pointed at the public +:meth:`~molzoo.mace.potential.MACEPotential.energy_core` by +``mace-subpackage-restructure-07-cleanup``; no in-tree caller binds the private +name any more. :attr:`MACEMatpes._compute_energy` is kept as a *name alias* of +``energy_core`` — the same function object, so the two can never disagree — +purely as back-compat for out-of-tree callers, and the positional signature +``(pos, Z, edge_index, batch, num_graphs, shifts)`` stays a contract for them. +:class:`MACEOMol` deliberately has **no** such alias: the flat OMol model's +``_compute_energy`` took ``(…, batch, total_charge, total_spin, shifts)``, so an +alias of ``energy_core`` there would silently reinterpret the fifth positional +argument as ``num_graphs``. + +Reference: + Batatia et al. "MACE: Higher Order Equivariant Message Passing Neural + Networks for Fast and Accurate Force Fields" NeurIPS 2022. + https://arxiv.org/abs/2206.07697 + Batatia et al. "A foundation model for atomistic materials chemistry" + (MACE-MP-0). https://arxiv.org/abs/2401.00096 + Kaplan et al. "A foundational potential energy surface dataset for + materials" (MatPES). https://arxiv.org/abs/2503.04070 + Levine et al. "The Open Molecules 2025 (OMol25) Dataset, Evaluations, and + Models" https://arxiv.org/abs/2505.08762 +""" + +from __future__ import annotations + +from collections.abc import Sequence + +import torch + +from molpot.derivation.force import autograd_forces_from_energy +from molzoo.mace.checkpoint import MATPES_REMAP, OMOL_REMAP +from molzoo.mace.potential import MACEPotential +from molzoo.mace.spec import MACEMatpesSpec, MACEOMolSpec + +__all__ = [ + "MACEMatpes", + "MACEOMol", + "load_matpes_state_dict", + "load_omol_state_dict", +] + + +def _supplied(**fields: object) -> dict[str, object]: + """Drop the ``None`` placeholders, leaving only the caller's own values. + + The keyword constructors below take ``None`` for "not given" instead of + repeating the shipped defaults. Restating ``num_bessel=10`` here would put a + second copy of every foundation hyper-parameter one module away from + :mod:`molzoo.mace.spec` — where the defaults are validated and + gate-checked against the flat constructors they replace (the + ``test_field_default_matches_the_flat_constructor`` cases in + ``tests/test_molzoo/test_mace/test_spec.py``). + Filtering instead means the spec stays the single source. + + Args: + **fields: Spec field name → the caller's value, or ``None``. + + Returns: + The subset whose value is not ``None``, ready to splat into a spec. + """ + return {name: value for name, value in fields.items() if value is not None} + + +def _energy_table(atomic_energies: Sequence[float] | torch.Tensor) -> list[float]: + """Normalise the ``E0`` table to the plain ``list[float]`` the spec holds. + + ``benchmarks/bench_mace_matpes.py:43`` passes a ``torch.Tensor`` + (``torch.zeros(len(_Z_TABLE), dtype=config.ftype)``), + while :class:`~molzoo.mace.spec.MACESpec` keeps the table torch-free so a + configuration can be read without importing torch. Both shapes are accepted + here; the tensor is rebuilt inside + :class:`~molzoo.mace.encoder.MACEEncoder`. + + Args: + atomic_energies: Per-element reference energies ``E0`` in eV/atom, in + ``atomic_numbers`` order — a sequence or a 1-D tensor. + + Returns: + The same energies as a ``list[float]`` (eV/atom). + """ + return [float(energy) for energy in atomic_energies] + + +class MACEMatpes(MACEPotential): + """MACE-MatPES foundation model, constructed by keyword. + + A thin adapter over :class:`~molzoo.mace.potential.MACEPotential`: the + keywords below are mapped one-for-one onto + :class:`~molzoo.mace.spec.MACEMatpesSpec` and nothing else happens. The + energy it evaluates, the ``F_i = -∂E/∂r_i`` convention and the units + (eV, eV/Å, Å) are documented on :mod:`molzoo.mace.potential`; the shipped + defaults on :class:`~molzoo.mace.spec.MACEMatpesSpec`. Prefer the spec form + in new code — see the module docstring, which also glosses the MACE + vocabulary the keywords below are named in. + + Only ``atomic_numbers`` and ``atomic_energies`` are required; every other + argument defaults to ``None``, meaning *not given*, in which case the + spec's own default applies. Those defaults are the ones this class' + predecessor hard-coded (the pre-cutover flat ``MACEMatpes``, deleted in + 06-wire; see git history). + + Args: + atomic_numbers: Element table (z-table) in checkpoint order, strictly + ascending. + atomic_energies: Per-element reference energies ``E0`` in eV/atom, same + order. A sequence or a 1-D ``torch.Tensor``. + r_max: Radial cutoff in Å. + num_bessel: Number of Bessel radial basis functions. + num_polynomial_cutoff: Polynomial cutoff exponent ``p`` (also the + envelope exponent of the ZBL — Ziegler–Biersack–Littmark — + short-range nuclear repulsion term, matching MACE). + l_max: Maximum spherical-harmonics order. + num_features: Scalar channel multiplicity (``hidden_irreps`` 0e count). + max_hidden_l: Highest ``l`` carried in the node state between layers. + num_interactions: Number of interaction/product layers. + correlation: Body-order correlation of the symmetric contraction. + mlp_dim: Hidden width of the final non-linear readout, the small + multi-layer perceptron (MLP) that turns features into an energy. + radial_mlp: Hidden widths of the radial weight MLP. + scale: ``atomic_inter_scale`` — multiplies the interaction energy (eV). + shift: ``atomic_inter_shift`` — per-atom energy offset in eV. + use_fallback: ``True`` puts the equivariant blocks on the pure-torch + path of cuEquivariance — NVIDIA's GPU library for the equivariant + tensor algebra MACE is built from — instead of its *fused* kernels + (``False``, the shipped default): single GPU kernels that do a + whole tensor product at once, which need a GPU plus the + ``cuequivariance-ops-torch`` wheel and are ~36x faster per step of + molecular dynamics (MD). + + Raises: + ValueError: From :class:`~molzoo.mace.spec.MACEMatpesSpec` — a table + that is not strictly ascending, an ``E0`` list of the wrong length, + or fewer than two interaction layers. + """ + + #: Name alias of :meth:`~molzoo.mace.potential.MACEPotential.energy_core`. + #: Same function object, so the positional signature + #: ``(pos, Z, edge_index, batch, num_graphs, shifts)`` cannot drift from the + #: public one. ``mace-subpackage-restructure-07-cleanup`` re-pointed the two + #: in-tree consumers (``scripts/matpes_port/run_nve.py``, + #: ``benchmarks/bench_mace_matpes.py``) onto ``energy_core``; the alias + #: remains as back-compat for out-of-tree callers and may be dropped by a + #: future spec. + _compute_energy = MACEPotential.energy_core + + def __init__( + self, + *, + atomic_numbers: Sequence[int], + atomic_energies: Sequence[float] | torch.Tensor, + r_max: float | None = None, + num_bessel: int | None = None, + num_polynomial_cutoff: int | None = None, + l_max: int | None = None, + num_features: int | None = None, + max_hidden_l: int | None = None, + num_interactions: int | None = None, + correlation: int | None = None, + mlp_dim: int | None = None, + radial_mlp: Sequence[int] | None = None, + scale: float | None = None, + shift: float | None = None, + use_fallback: bool | None = None, + ) -> None: + super().__init__( + MACEMatpesSpec( + atomic_numbers=list(atomic_numbers), + atomic_energies=_energy_table(atomic_energies), + **_supplied( + r_max=r_max, + num_bessel=num_bessel, + num_polynomial_cutoff=num_polynomial_cutoff, + l_max=l_max, + num_features=num_features, + max_hidden_l=max_hidden_l, + num_interactions=num_interactions, + correlation=correlation, + mlp_dim=mlp_dim, + radial_mlp=None if radial_mlp is None else list(radial_mlp), + scale=scale, + shift=shift, + use_fallback=use_fallback, + ), + ) + ) + + def energy_forces( + self, + positions: torch.Tensor, + Z: torch.Tensor, + edge_index: torch.Tensor, + batch: torch.Tensor, + num_graphs: int | None = None, + shifts: torch.Tensor | None = None, + ) -> dict[str, torch.Tensor]: + """Energy and forces from raw tensors, off the batch schema. + + The raw-tensor twin of :meth:`~molzoo.mace.potential.MACEPotential.forward`, + kept because ``benchmarks/bench_mace_matpes.py:115`` calls it. One energy + forward, then one ``torch.autograd.grad`` on a position *leaf* — the + tensor autograd differentiates with respect to — created and owned by + this call. That is MACE's own ``get_outputs`` shape, and the same + tensor-level kernel + (:func:`molpot.derivation.force.autograd_forces_from_energy`) that + :func:`molpot.derivation.kernels.grad_force_pass` uses on the batch + path. No third force path is derived here. + + The returned energy is always detached: this method creates the leaf, so + the caller has no graph to lose (the same ``needs_leaf`` policy + ``forward`` gets from ``detach_energy=None``). Differentiating the energy + w.r.t. the parameters needs the batch path, not this one. + + Args: + positions: Atom positions ``(N, 3)`` in Å, for the ``N`` atoms of + the system. + Z: Atomic numbers ``(N,)``. Must lie inside this model's element + table; **not** checked on this path — an off-table element is + snapped onto a neighbouring row by ``torch.searchsorted`` and + yields a plausible, wrong energy, so call + :meth:`~molzoo.mace.encoder.MACEEncoder.validate_elements` once + per (model, system) pair. + edge_index: ``(E, 2)`` for the ``E`` neighbour pairs, with + ``[:, 0]`` = source, ``[:, 1]`` = target. + batch: Graph index per atom ``(N,)`` — ``batch[i]`` says which + graph atom ``i`` belongs to. + num_graphs: Number of graphs ``B`` (one graph = one molecule or one + periodic cell). Inferred from ``batch`` when omitted, which + costs a *host sync* — the CPU waits for the GPU to hand the + value back — so pass it on hot paths. + shifts: Optional periodic shift vectors ``(E, 3)`` in Å, added to + the edge displacements so a neighbour across a periodic + boundary enters at its imaged position. + + Returns: + ``{"energy": (B,) eV, "forces": (N, 3) eV/Å}``. + """ + if num_graphs is None: + num_graphs = int(batch.max().item()) + 1 + leaf = positions.detach().requires_grad_(True) + with torch.enable_grad(): + energy = self.energy_core(leaf, Z, edge_index, batch, num_graphs, shifts) + return {"energy": energy.detach(), "forces": autograd_forces_from_energy(energy, leaf)} + + +class MACEOMol(MACEPotential): + """MACE-OMol foundation model, constructed by keyword. + + A thin adapter over :class:`~molzoo.mace.potential.MACEPotential`, mapping + its keywords onto :class:`~molzoo.mace.spec.MACEOMolSpec`; see + :class:`MACEMatpes` and the module docstring for the shape, for the keyword + vocabulary, and for why new code should build the spec directly. OMol adds + the per-graph charge/spin *conditioning*: the integer total charge and + total spin of a molecule index a learned vector in an *embedding* table — + a lookup from an integer to a trainable vector — which is added to every + atom's initial features. Its energy therefore carries one further per-atom + term outside the ``scale``/``shift`` normalisation + (:mod:`molzoo.mace.potential`). + + ``use_fallback`` is deliberately absent, as it was on the predecessor (the + pre-cutover flat ``MACEOMol``, deleted in 06-wire; see git history): that + model hard-coded the fused cuEquivariance path, which is what + :class:`~molzoo.mace.spec.MACEOMolSpec` defaults to. + + Args: + atomic_numbers: Element table (z-table) in checkpoint order, strictly + ascending. + atomic_energies: Per-element reference energies ``E0`` in eV/atom, same + order. A sequence or a 1-D ``torch.Tensor``. + r_max: Radial cutoff in Å. + num_bessel: Number of Bessel radial basis functions. + num_polynomial_cutoff: Polynomial cutoff exponent ``p``. + l_max: Maximum spherical-harmonics order (at least 1). + num_features: Scalar channel multiplicity (``hidden_irreps`` 0e count). + num_interactions: Number of interaction/product layers. + correlation: Body-order correlation of the symmetric contraction. + mlp_dim: Hidden width of the final non-linear readout, the small + multi-layer perceptron (MLP) that turns features into an energy. + edge_channels: Per-``l`` channel count of the mid-layer edge irreps and + the radial-MLP hidden width. + charge_classes: Number of rows in the total-charge embedding table. + charge_offset: Added to the total charge to get the row index (row = + ``total_charge + charge_offset``), so that negative charges still + land on a valid row — with the shipped ``100``, charge −100 → row 0. + spin_classes: Number of rows in the total-spin embedding table. + spin_offset: Added to the total spin to get the row index, on the same + convention as ``charge_offset``. + scale: ``atomic_inter_scale`` — multiplies the interaction energy (eV). + shift: ``atomic_inter_shift`` — per-atom energy offset in eV. + + Raises: + ValueError: From :class:`~molzoo.mace.spec.MACEOMolSpec` — a table that + is not strictly ascending, an ``E0`` list of the wrong length, or + ``l_max`` below 1. + """ + + def __init__( + self, + *, + atomic_numbers: Sequence[int], + atomic_energies: Sequence[float] | torch.Tensor, + r_max: float | None = None, + num_bessel: int | None = None, + num_polynomial_cutoff: int | None = None, + l_max: int | None = None, + num_features: int | None = None, + num_interactions: int | None = None, + correlation: int | None = None, + mlp_dim: int | None = None, + edge_channels: int | None = None, + charge_classes: int | None = None, + charge_offset: int | None = None, + spin_classes: int | None = None, + spin_offset: int | None = None, + scale: float | None = None, + shift: float | None = None, + ) -> None: + super().__init__( + MACEOMolSpec( + atomic_numbers=list(atomic_numbers), + atomic_energies=_energy_table(atomic_energies), + **_supplied( + r_max=r_max, + num_bessel=num_bessel, + num_polynomial_cutoff=num_polynomial_cutoff, + l_max=l_max, + num_features=num_features, + num_interactions=num_interactions, + correlation=correlation, + mlp_dim=mlp_dim, + edge_channels=edge_channels, + charge_classes=charge_classes, + charge_offset=charge_offset, + spin_classes=spin_classes, + spin_offset=spin_offset, + scale=scale, + shift=shift, + ), + ) + ) + + def energy_forces( + self, + positions: torch.Tensor, + Z: torch.Tensor, + edge_index: torch.Tensor, + batch: torch.Tensor, + total_charge: torch.Tensor, + total_spin: torch.Tensor, + shifts: torch.Tensor | None = None, + ) -> dict[str, torch.Tensor]: + """Energy and forces from raw tensors, conditioning tensors included. + + The OMol counterpart of :meth:`MACEMatpes.energy_forces`; the leaf + ownership, the detach policy and the force kernel are identical (see + there). ``num_graphs`` is read off ``total_charge``'s static shape + rather than ``batch.max().item()``, so this path costs no host sync. + + Args: + positions: Atom positions ``(N, 3)`` in Å, for the ``N`` atoms of + the system. + Z: Atomic numbers ``(N,)``. Must lie inside this model's element + table; **not** checked on this path — see + :meth:`MACEMatpes.energy_forces` for why that matters and + :meth:`~molzoo.mace.encoder.MACEEncoder.validate_elements` for + the once-per-system check. + edge_index: ``(E, 2)`` for the ``E`` neighbour pairs, with + ``[:, 0]`` = source, ``[:, 1]`` = target. + batch: Graph index per atom ``(N,)`` — ``batch[i]`` says which + graph atom ``i`` belongs to. + total_charge: Per-graph total charge ``(B,)`` in units of the + elementary charge ``e``, for the ``B`` graphs in the batch. Its + length is what fixes ``B`` here. + total_spin: Per-graph total spin ``(B,)``, dimensionless, on the + multiplicity ``2S+1`` convention with ``S`` the total electron + spin (``1`` = closed-shell singlet, every electron paired). + shifts: Optional periodic shift vectors ``(E, 3)`` in Å, added to + the edge displacements so a neighbour across a periodic + boundary enters at its imaged position. + + Returns: + ``{"energy": (B,) eV, "forces": (N, 3) eV/Å}``. + """ + leaf = positions.detach().requires_grad_(True) + with torch.enable_grad(): + energy = self.energy_core( + leaf, + Z, + edge_index, + batch, + total_charge.shape[0], + shifts, + total_charge=total_charge, + total_spin=total_spin, + ) + return {"energy": energy.detach(), "forces": autograd_forces_from_energy(energy, leaf)} + + +def load_matpes_state_dict(model: MACEPotential, cueq_state: dict[str, torch.Tensor]) -> None: + """Load a cueq-converted MACE-MatPES ``state_dict``, strictly. + + A ``state_dict`` is PyTorch's flat ``name -> tensor`` dump of a model's + weights; *cueq-converted* means it has already been rewritten, out of tree + by ``mace.cli.convert_e3nn_cueq``, into the layout of cuEquivariance — + NVIDIA's GPU library for the equivariant tensor algebra MACE is built from. + + Compatibility alias, forwarding to :data:`molzoo.mace.checkpoint.MATPES_REMAP` + — the ``on_unexpected="raise"`` preset of + :class:`~molzoo.mace.checkpoint.CheckpointRemap`, which owns the key dialect + and the strictness doctrine: every learnable tensor of ``model`` must be + filled by the checkpoint and every checkpoint weight must land somewhere. A + silently dropped tensor is the failure mode that produces a model which + runs, looks sane, and is quietly wrong. New code should call + ``MATPES_REMAP.load(model, state)`` directly, or build the model with + :meth:`~molzoo.mace.potential.MACEPotential.from_checkpoint`. + + Args: + model: Target model, constructed with the checkpoint's + hyper-parameters — a :class:`MACEMatpes` or any + :class:`~molzoo.mace.potential.MACEPotential` built from a + :class:`~molzoo.mace.spec.MACEMatpesSpec`. + cueq_state: ``state_dict`` of the cueq-converted official model + (``mace.cli.convert_e3nn_cueq``, run out of tree). + + Returns: + ``None`` — the legacy contract. The remap's + ``(missing_buffers, unexpected)`` report is dropped; under this policy + ``unexpected`` is empty by construction and ``missing_buffers`` holds + only entries cuEquivariance rebuilds. Call the remap directly to see it. + + Raises: + RuntimeError: If a checkpoint key has no home, a shape disagrees, or a + model parameter is left unfilled. + """ + MATPES_REMAP.load(model, cueq_state) + + +def load_omol_state_dict( + model: MACEPotential, cueq_state: dict[str, torch.Tensor] +) -> tuple[list[str], list[str]]: + """Load a cueq-converted MACE-OMol ``state_dict``, strictly. + + Compatibility alias, forwarding to :data:`molzoo.mace.checkpoint.OMOL_REMAP` + — the ``on_unexpected="return"`` preset of + :class:`~molzoo.mace.checkpoint.CheckpointRemap`. Unhoused checkpoint keys + are *returned* rather than raised, because the OMol checkpoint family + carries auxiliary heads this port deliberately does not model; an unfilled + parameter or a shape disagreement still raises. New code should call + ``OMOL_REMAP.load(model, state)`` directly. + + Args: + model: Target model, constructed with the checkpoint's + hyper-parameters — a :class:`MACEOMol` or any + :class:`~molzoo.mace.potential.MACEPotential` built from a + :class:`~molzoo.mace.spec.MACEOMolSpec`. + cueq_state: ``state_dict`` of the cueq-converted official model. + + Returns: + ``(missing_buffers, unexpected)`` — the non-learnable entries the + checkpoint did not provide (buffers, not ``nn.Parameter``) and the + renamed checkpoint keys with no home in ``model``. Both lists are + sorted by :meth:`~molzoo.mace.checkpoint.CheckpointRemap.load`; the + legacy contract, passed through unchanged. + + Raises: + RuntimeError: If a learnable parameter is left unfilled or a shape + disagrees. + """ + return OMOL_REMAP.load(model, cueq_state) diff --git a/src/molzoo/mace_omol.py b/src/molzoo/mace_omol.py deleted file mode 100644 index 0dc8c08..0000000 --- a/src/molzoo/mace_omol.py +++ /dev/null @@ -1,354 +0,0 @@ -"""MACE-OMOL foundation model, assembled from molrep/molpot building blocks. - -A faithful, native reimplementation of the official ``MACE-omol-0`` model -(``ScaleShiftMACE`` with charge/spin conditioning and non-linear residual -interactions) on the cuEquivariance backend. Every sub-block lives in -``molrep`` / ``molpot``; this module only wires them into the energy/force -forward of MACE's ``ScaleShiftMACE``. - -Use :func:`load_omol_state_dict` to import weights from an official MACE-OMOL -checkpoint that has been converted to cueq layout via -``mace.cli.convert_e3nn_cueq``. - -Reference: - Batatia et al. "MACE" NeurIPS 2022; OMol25 foundation model. - https://arxiv.org/abs/2206.07697 -""" - -from __future__ import annotations - -import cuequivariance as cue -import cuequivariance_torch as cuet -import torch -import torch.nn as nn -from tensordict import TensorDict - -from molix import config -from molix.F.scatter import scatter_sum_compile_safe as _scatter_sum -from molpot.derivation.force import ForceDerivation -from molpot.heads.energy import AtomicReferenceEnergy -from molpot.heads.rescale import GlobalRescale -from molrep.embedding.angular import SphericalHarmonics -from molrep.embedding.cutoff import PolynomialCutoff -from molrep.embedding.node import JointFeatureEmbedding, JointFeatureSpec -from molrep.embedding.radial import BesselRBF -from molrep.interaction.product_basis import EquivariantProductBasis -from molrep.interaction.residual import ResidualInteraction -from molrep.readout.scalar import NonLinearBiasReadout - - -class MACEOMol(nn.Module): - """Native MACE-OMOL energy/force model (cuEquivariance).""" - - def __init__( - self, - *, - atomic_numbers: list[int], - atomic_energies: torch.Tensor, - r_max: float = 6.0, - num_bessel: int = 8, - num_polynomial_cutoff: int = 5, - l_max: int = 3, - num_features: int = 1024, - num_interactions: int = 3, - correlation: int = 2, - mlp_dim: int = 16, - charge_classes: int = 201, - charge_offset: int = 100, - spin_classes: int = 101, - spin_offset: int = 0, - scale: float = 1.0, - shift: float = 0.0, - ) -> None: - super().__init__() - ftype = config.ftype - n_el = len(atomic_numbers) - self.register_buffer( - "z_table", torch.tensor(atomic_numbers, dtype=torch.long), persistent=True - ) - self.num_interactions = num_interactions - - sh = "+".join(f"1x{l}{'e' if l % 2 == 0 else 'o'}" for l in range(l_max + 1)) - feat0 = f"{num_features}x0e" - target = "+".join( - f"{num_features}x{l}{'e' if l % 2 == 0 else 'o'}" for l in range(l_max + 1) - ) - - self.node_embedding = cuet.Linear( - cue.Irreps("O3", f"{n_el}x0e"), - cue.Irreps("O3", feat0), - layout=cue.ir_mul, - dtype=ftype, - ) - self.spherical_harmonics = SphericalHarmonics(l_max=l_max) - self.bessel = BesselRBF( - r_cut=r_max, num_radial=num_bessel, normalize=False, eps=0.0, trainable=True - ) - self.cutoff_fn = PolynomialCutoff(r_cut=r_max, exponent=num_polynomial_cutoff) - - self.joint_embedding = JointFeatureEmbedding( - feature_specs=[ - JointFeatureSpec( - name="total_spin", - kind="categorical", - emb_dim=num_features, - num_classes=spin_classes, - per="graph", - offset=spin_offset, - ), - JointFeatureSpec( - name="total_charge", - kind="categorical", - emb_dim=num_features, - num_classes=charge_classes, - per="graph", - offset=charge_offset, - ), - ], - out_dim=num_features, - ) - self.embedding_readout = cuet.Linear( - cue.Irreps("O3", feat0), - cue.Irreps("O3", "1x0e"), - layout=cue.ir_mul, - dtype=ftype, - ) - self.atomic_energies = AtomicReferenceEnergy( - atomic_energies=atomic_energies, - atomic_numbers=atomic_numbers, - ) - - # per-layer irreps (mirror official ScaleShiftMACE hidden_irreps schedule) - hidden_full = target.replace(f"+{num_features}x{l_max}{'e' if l_max % 2 == 0 else 'o'}", "") - edge0 = feat0 - edge_mid = "+".join( - f"128x{l}{'e' if l % 2 == 0 else 'o'}" for l in range(l_max) - ) # 128x0e+128x1o+128x2e - node_in = [feat0, hidden_full, hidden_full] - edge_irr = [edge0, edge_mid, edge_mid] - hidden_sched = [hidden_full, hidden_full, feat0] - - self.interactions = nn.ModuleList() - self.products = nn.ModuleList() - for i in range(num_interactions): - self.interactions.append( - ResidualInteraction( - node_attrs_irreps=f"{n_el}x0e", - node_feats_irreps=node_in[i], - edge_attrs_irreps=sh, - edge_feats_irreps=f"{num_bessel}x0e", - edge_irreps=edge_irr[i], - target_irreps=target, - hidden_irreps=hidden_sched[i], - radial_mlp=[128, 128, 128], - # autograd force path → fused cuEq kernels OK - use_fallback=False, - ) - ) - self.products.append( - EquivariantProductBasis( - node_feats_irreps=target, - target_irreps=hidden_sched[i], - correlation=correlation, - num_elements=1, - use_sc=True, - use_fallback=False, - ) - ) - self.readout = NonLinearBiasReadout(irreps_in=feat0, mlp_dim=mlp_dim) - self.scale_shift = GlobalRescale(scale=scale, shift=shift) - # MACE uses cuEquivariance fused kernels → autograd (functorch rejects - # their legacy autograd.Function); matches the upstream MACE library. - self.force_derivation = ForceDerivation(method="autograd") - - def energy_forces( - self, - positions: torch.Tensor, - Z: torch.Tensor, - edge_index: torch.Tensor, - batch: torch.Tensor, - total_charge: torch.Tensor, - total_spin: torch.Tensor, - shifts: torch.Tensor | None = None, - compute_forces: bool = True, - ) -> dict: - """Return total energy and forces. - - Args: - positions: ``(N, 3)``. - Z: atomic numbers ``(N,)``. - edge_index: ``(2, E)`` sender/receiver. - batch: graph index per atom ``(N,)``. - total_charge / total_spin: per-graph ``(B,)``. - shifts: optional PBC shift vectors ``(E, 3)``. - compute_forces: whether to compute ``-dE/dx``. - """ - # Force path always goes through ForceDerivation (autograd backend — - # cuEq fused kernels). Do not hand-roll torch.autograd.grad here. - pos = positions.detach() - - def energy_fn(p: torch.Tensor) -> torch.Tensor: - return self._compute_energy( - p, Z, edge_index, batch, total_charge, total_spin, shifts - ).sum() - - total_energy = self._compute_energy( - pos, Z, edge_index, batch, total_charge, total_spin, shifts - ) - out = {"energy": total_energy} - if compute_forces: - out["forces"] = self.force_derivation(energy_fn, pos) - return out - - def _compute_energy( - self, - positions: torch.Tensor, - Z: torch.Tensor, - edge_index: torch.Tensor, - batch: torch.Tensor, - total_charge: torch.Tensor, - total_spin: torch.Tensor, - shifts: torch.Tensor | None = None, - ) -> torch.Tensor: - """Per-graph total energy ``(B,)`` as a pure function of ``positions``. - - Shared core of :meth:`energy_forces` and :meth:`forward`. Recomputes all - position-derived geometry internally so it can be differentiated with - ``torch.autograd.grad`` (ForceDerivation(method="autograd")). - - Args: - positions: ``(N, 3)``. - Z: atomic numbers ``(N,)``. - edge_index: ``(2, E)`` sender/receiver. - batch: graph index per atom ``(N,)``. - total_charge / total_spin: per-graph ``(B,)``. - shifts: optional PBC shift vectors ``(E, 3)``. - """ - num_nodes = positions.shape[0] - # Derive the graph count from the per-graph ``total_charge`` length (a - # static shape) rather than ``int(batch.max().item())``: the ``.item()`` - # forces a host sync that breaks the dynamo graph, blocking - # ``torch.compile(fullgraph=True)`` of the functorch force path. - num_graphs = total_charge.shape[0] - - sender, receiver = edge_index[0], edge_index[1] - vectors = positions[receiver] - positions[sender] - if shifts is not None: - vectors = vectors + shifts - lengths = torch.linalg.norm(vectors, dim=-1, keepdim=True) - - # one-hot node attrs over the element table - z_index = torch.searchsorted(self.z_table, Z.reshape(-1)).to(dtype=torch.long) - node_attrs = torch.zeros( - num_nodes, self.z_table.numel(), dtype=positions.dtype, device=positions.device - ) - node_attrs[torch.arange(num_nodes), z_index] = 1.0 - - node_e0 = self.atomic_energies(Z) - e0 = _scatter_sum(node_e0, batch, num_graphs) - - node_feats = self.node_embedding(node_attrs) - edge_attrs = self.spherical_harmonics(vectors) - edge_feats = self.bessel(lengths.squeeze(-1)) - cutoff = self.cutoff_fn(lengths.squeeze(-1)).unsqueeze(-1) - - node_feats = node_feats + self.joint_embedding( - batch, total_spin=total_spin, total_charge=total_charge - ) - emb_node_e = self.embedding_readout(node_feats).squeeze(-1) - e0 = e0 + _scatter_sum(emb_node_e, batch, num_graphs) - - feats_last = None - for i in range(self.num_interactions): - node_feats, sc = self.interactions[i]( - node_attrs, node_feats, edge_attrs, edge_feats, edge_index, cutoff - ) - node_feats = self.products[i](node_feats, sc, node_attrs) - feats_last = node_feats - - node_es = self.readout(feats_last).squeeze(-1) - node_inter = self.scale_shift(node_es) - inter_e = _scatter_sum(node_inter, batch, num_graphs) - return e0 + inter_e - - def forward(self, td: TensorDict) -> TensorDict: - """Run MACE-OMOL on a post-collate batch, writing energy and forces. - - Reads ``atoms.{Z,pos,batch}``, ``edges.edge_index`` (``(E, 2)`` with - ``[:,0]`` source / ``[:,1]`` target per the molnex edge convention), and - per-graph ``graphs.{total_charge,total_spin}`` (defaulting to a neutral - singlet when absent). Forces are obtained through - :class:`molpot.derivation.ForceDerivation` (``method="autograd"`` — - ``F = -∂E/∂pos`` via ``torch.autograd.grad``; cuEq-safe). Writes - ``graphs.energy`` ``(B,)`` and ``atoms.forces`` ``(N, 3)`` back into - ``td`` and returns it. - - Args: - td: post-collate ``TensorDict`` with ``atoms`` / ``edges`` (and - optionally ``graphs``) sub-dicts. - - Returns: - The same ``td`` with ``graphs.energy`` and ``atoms.forces`` added. - """ - Z = td["atoms", "Z"] - positions = td["atoms", "pos"] - batch = td["atoms", "batch"] - # molnex edge_index is (E, 2) [source, target]; the core wants (2, E). - edge_index = td["edges", "edge_index"].t().contiguous() - num_graphs = int(batch.max().item()) + 1 - - nested = td.keys(include_nested=True) - if ("graphs", "total_charge") in nested: - total_charge = td["graphs", "total_charge"] - else: - total_charge = torch.zeros(num_graphs, dtype=torch.long, device=Z.device) - if ("graphs", "total_spin") in nested: - total_spin = td["graphs", "total_spin"] - else: - # OMOL convention: 1 = closed-shell singlet (spin_offset=0 → index 1, - # a trained row). Spin 0 hits an untrained embedding row → garbage. - total_spin = torch.ones(num_graphs, dtype=torch.long, device=Z.device) - - energy = self._compute_energy(positions, Z, edge_index, batch, total_charge, total_spin) - - def energy_fn(pos: torch.Tensor) -> torch.Tensor: - return self._compute_energy(pos, Z, edge_index, batch, total_charge, total_spin).sum() - - forces = self.force_derivation(energy_fn, positions) - - if "graphs" not in td.keys(): - td["graphs"] = TensorDict({}, batch_size=[num_graphs]) - td["graphs", "energy"] = energy - td["atoms", "forces"] = forces - return td - - -def load_omol_state_dict(model: MACEOMol, cueq_state: dict) -> tuple[list, list]: - """Load a cueq-converted MACE-OMOL ``state_dict`` into :class:`MACEOMol`. - - Maps the official cueq key names onto this model's modules (direct copy); - returns ``(missing_learnable, unexpected)`` for inspection. - """ - remap = {} - for k, v in cueq_state.items(): - nk = k - if k.startswith("node_embedding.linear."): - nk = "node_embedding." + k.split("node_embedding.linear.")[1] - elif k.startswith("embedding_readout.linear."): - nk = "embedding_readout." + k.split("embedding_readout.linear.")[1] - elif k.startswith("readouts.0."): - nk = "readout." + k.split("readouts.0.")[1] - elif k.startswith("atomic_energies_fn."): - continue # handled at construction (Z-indexed) - remap[nk] = v - missing, unexpected = model.load_state_dict(remap, strict=False) - learnable_leaf = (".weight", ".bias") - learnable_scalar = ("alpha", "beta", "scale", "shift") - miss_learn = [ - m - for m in missing - if (m.endswith(learnable_leaf) or m.split(".")[-1] in learnable_scalar) - and not m.startswith("atomic_energies") - and "scale_shift" not in m - ] - return miss_learn, unexpected diff --git a/src/molzoo/pinet/__init__.py b/src/molzoo/pinet/__init__.py index 63c335b..5670205 100644 --- a/src/molzoo/pinet/__init__.py +++ b/src/molzoo/pinet/__init__.py @@ -19,11 +19,6 @@ from .properties import PiNetDipole, PiNetPolarizability, pool_layer from .spec import PiNetSpec -# Back-compat private aliases used by older call sites / docs. -_compute_d5 = compute_d5 -_edge_bond_diff = edge_bond_diff -_pool_layer = pool_layer - __all__ = [ "PiNet", "PiNetSpec", diff --git a/src/molzoo/pinet/encoder.py b/src/molzoo/pinet/encoder.py index fe497c0..33cbba4 100644 --- a/src/molzoo/pinet/encoder.py +++ b/src/molzoo/pinet/encoder.py @@ -32,6 +32,9 @@ class PiNet(TensorDictModuleBase): * ``("atoms", "node_features")``: scalar P1 states ``(N, depth, D)``. * ``("atoms", "p1_block_outputs")``: raw block outputs ``(N, depth, D)``. + + And, when ``emit_property_features`` is set (the default): + * ``("atoms", "p3_features")`` / ``("atoms", "p5_features")`` when enabled. * ``("edges", "i1_features")`` (+ ``i3`` / ``i5`` when enabled). """ @@ -60,7 +63,25 @@ def __init__( activation: str = "tanh", weighted: bool = False, rank: Literal[1, 3, 5] = 3, + emit_property_features: bool = True, ) -> None: + """Build a PiNet encoder. + + Args: + emit_property_features: Also stack and write the per-block + equivariant node states (``p3``/``p5``) and interaction states + (``i1``/``i3``/``i5``) that the PiNet property heads + (:class:`~molzoo.pinet.properties.PiNetDipole`, + :class:`~molzoo.pinet.properties.PiNetPolarizability`) consume. + The energy/force path never reads them, and the ``i*`` tracks + are per-*edge* — at a typical 30 neighbours/atom they dominate + the encoder's output volume (~150x the bytes the energy path + actually reads) and stay pinned for the force double-backward. + Leave ``True`` when attaching a property head; set ``False`` + for pure energy/force models. Every other argument is + unaffected — the tracks are still computed inside each block, + only the extra stack + write is skipped. + """ super().__init__() self.config = PiNetSpec( atom_types=atom_types or [1, 6, 7, 8], @@ -77,6 +98,7 @@ def __init__( activation=activation, weighted=weighted, rank=rank, + emit_property_features=emit_property_features, ) cfg = self.config if cfg.pp_nodes[-1] != cfg.ii_nodes[-1]: @@ -84,6 +106,7 @@ def __init__( self.rank = int(cfg.rank) self.depth = int(cfg.depth) + self.emit_property_features = bool(cfg.emit_property_features) self.feature_dim = int(cfg.ii_nodes[-1]) self.n_props = int(self.rank // 2) + 1 self.output_dim = self.feature_dim @@ -193,6 +216,12 @@ def forward(self, td: TensorDict) -> TensorDict: fc = self.cutoff(edge_dist) basis = self.basis_fn(edge_dist, fc=fc) + # The p3/p5 and i1/i3/i5 histories are only read by the property heads. + # For a pure energy/force model they are dead weight — and the ``i*`` + # tracks are per-edge, so at ~30 neighbours/atom they dominate both the + # copy cost and the memory pinned for the force double-backward. Skip + # accumulating them entirely rather than stacking and discarding. + emit_props = self.emit_property_features p1_states: list[torch.Tensor] = [] p1_block_outputs: list[torch.Tensor] = [] p3_states: list[torch.Tensor] = [] @@ -206,25 +235,29 @@ def forward(self, td: TensorDict) -> TensorDict: p1_block_outputs.append(new["p1"]) tensors["p1"] = self.res_update1[i](tensors["p1"], new["p1"]) p1_states.append(tensors["p1"]) - i1_states.append(new["i1"]) + if emit_props: + i1_states.append(new["i1"]) if self.rank >= 3: tensors["p3"] = self.res_update3[i](tensors["p3"], new["p3"]) - p3_states.append(tensors["p3"]) - i3_states.append(new["i3"]) + if emit_props: + p3_states.append(tensors["p3"]) + i3_states.append(new["i3"]) if self.rank >= 5: tensors["p5"] = self.res_update5[i](tensors["p5"], new["p5"]) - p5_states.append(tensors["p5"]) - i5_states.append(new["i5"]) + if emit_props: + p5_states.append(tensors["p5"]) + i5_states.append(new["i5"]) td["atoms", "node_features"] = torch.stack(p1_states, dim=1) td["atoms", "p1_block_outputs"] = torch.stack(p1_block_outputs, dim=1) - td["edges", "i1_features"] = torch.stack(i1_states, dim=1) - if self.rank >= 3: - td["atoms", "p3_features"] = torch.stack(p3_states, dim=1) - td["edges", "i3_features"] = torch.stack(i3_states, dim=1) - if self.rank >= 5: - td["atoms", "p5_features"] = torch.stack(p5_states, dim=1) - td["edges", "i5_features"] = torch.stack(i5_states, dim=1) + if emit_props: + td["edges", "i1_features"] = torch.stack(i1_states, dim=1) + if self.rank >= 3: + td["atoms", "p3_features"] = torch.stack(p3_states, dim=1) + td["edges", "i3_features"] = torch.stack(i3_states, dim=1) + if self.rank >= 5: + td["atoms", "p5_features"] = torch.stack(p5_states, dim=1) + td["edges", "i5_features"] = torch.stack(i5_states, dim=1) return td diff --git a/src/molzoo/pinet/potential.py b/src/molzoo/pinet/potential.py index 6e096e3..f3b7307 100644 --- a/src/molzoo/pinet/potential.py +++ b/src/molzoo/pinet/potential.py @@ -1,10 +1,7 @@ -"""PiNet energy + force potential (encoder composition + derivation). +"""PiNet energy + force potential (encoder + monomorphic init-fixed pipeline). -Physics (force derivation, energy aggregation) lives in ``molpot`` via -:class:`molpot.composition.energy_force.EnergyForceModel`. This module only -wires the PiNet encoder + per-block OutLayers — the long-term home for a fully -generic encoder→energy façade remains ``molpot``; public import stays -``from molzoo.pinet import PiNetPotential``. +All branching is in ``__init__`` (``compute_forces``, ``method``). ``forward`` +always runs one static pipeline — energy-only has **no** Derivative session. """ from __future__ import annotations @@ -15,42 +12,49 @@ import torch.nn as nn from tensordict import TensorDict -from molpot.composition.energy_force import EnergyForceModel -from molpot.derivation import EnergyAggregation +from molpot.derivation import EnergyAggregation, func_force_pass, grad_force_pass +from molpot.derivation.protocol import write_energy from molrep.interaction.pinet import OutLayer from .encoder import PiNet -class PiNetPotential(EnergyForceModel): - """PiNet energy + force prediction model — ready to use from hyperparameters. +class PiNetPotential(nn.Module): + """PiNet energy (+ optional forces) with in-place batch writes. - Pass PiNet hyperparameters directly; the encoder is built internally:: + Configuration is fixed at construction; ``forward`` has no flags:: - model = PiNetPotential(atom_types=[1, 6, 7, 8], r_max=4.5, depth=5, - hidden_dim=64, compute_forces=True) + # energy only + model = PiNetPotential(..., compute_forces=False) + batch = model(batch) # writes graphs.energy - Forces use ``ForceDerivation(method="functorch")`` (pure-PyTorch graph) - through :class:`EnergyForceModel`. All linears are fully specified at - construction — no lazy materialisation. + # energy + forces (func = compile-friendly 1-pass) + model = PiNetPotential(..., compute_forces=True, method="func") + batch = model(batch) # + atoms.forces + + # train force-matching often prefers method="grad" + model = PiNetPotential(..., compute_forces=True, method="grad") """ def __init__( self, *, hidden_dim: int = 64, - layer_reduction: Literal["mean", "sum", "last"] = "mean", compute_forces: bool = False, + method: Literal["func", "grad"] = "func", encoder: PiNet | None = None, **pinet_kwargs: object, ) -> None: - super().__init__(force_method="functorch", compute_forces=compute_forces) + super().__init__() if encoder is not None and pinet_kwargs: raise ValueError("Pass either encoder=... or PiNet kwargs, not both.") - self.encoder = encoder if encoder is not None else PiNet(**pinet_kwargs) # type: ignore[arg-type] - # Accepted for API compatibility; PiNet2 accumulates per-block OutLayers - # residually — there is no layer axis to reduce for energy. - self.layer_reduction = layer_reduction + if method not in ("func", "grad"): + raise ValueError(f"method must be 'func' or 'grad', got {method!r}") + if encoder is not None: + self.encoder = encoder + else: + pinet_kwargs.setdefault("emit_property_features", False) + self.encoder = PiNet(**pinet_kwargs) depth: int = int(getattr(self.encoder, "depth", 1)) feature_dim: int = int(getattr(self.encoder, "feature_dim", hidden_dim)) @@ -66,9 +70,24 @@ def __init__( ] ) self.energy_aggregation = EnergyAggregation(pooling="sum") - - def energy_forward(self, batch: TensorDict) -> dict[str, torch.Tensor]: - """Compilable energy: encoder → per-block OutLayer sum → aggregate.""" + self.compute_forces = compute_forces + self.method = method + + # Monomorphic pipeline — one bound call path, no runtime if/else. + if not compute_forces: + self._pipeline = self._write_energy + elif method == "func": + self._pipeline = self._pipeline_ef_func + else: + self._pipeline = self._pipeline_ef_grad + + def forward(self, batch: TensorDict) -> TensorDict: + """Run the init-fixed pipeline; mutate ``batch`` in place.""" + return self._pipeline(batch) + + # ------------------------------------------------------------------ energy + def _write_energy(self, batch: TensorDict) -> TensorDict: + """Energy core only — also the energy-only public pipeline.""" batch = self.encoder(batch) block_outputs = batch["atoms", "p1_block_outputs"] # (N, depth, D) @@ -82,12 +101,56 @@ def energy_forward(self, batch: TensorDict) -> dict[str, torch.Tensor]: atom_energy = atom_energy * batch["atoms", "mask"].to(atom_energy.dtype) num_graphs = batch["graphs"].batch_size[0] energy = self.energy_aggregation(atom_energy, atom_batch, num_graphs=num_graphs) + write_energy(batch, energy, atomic_energy=atom_energy) + return batch + + # ----------------------------------------------------------- force kernels + def _pipeline_ef_func(self, batch: TensorDict) -> TensorDict: + """Init-fixed: single ``torch.func.grad(..., has_aux=True)`` pass. + + Compile-friendly monomorphic path (no session / no set_non_tensor). + """ + return func_force_pass(self._write_energy, batch) + + def _pipeline_ef_grad(self, batch: TensorDict) -> TensorDict: + """Init-fixed: one energy pass + ``torch.autograd.grad`` on positions. + + Often faster for force-supervised training (``create_graph`` when + ``self.training``). Energy stays attached (``detach_energy=False``) so + an energy loss can share the graph with the force loss. + """ + return grad_force_pass( + self._write_energy, + batch, + create_graph=bool(self.training), + detach_energy=False, + ) - return { - "atomic_energy": atom_energy, - "energy": energy, - } - - # Back-compat alias used by older call sites / docs. - def _energy_forward(self, batch: TensorDict) -> dict[str, torch.Tensor]: - return self.energy_forward(batch) + # ---------------------------------------------------------------- compile + def compile( + self, + *, + backend: str = "inductor", + fullgraph: bool = False, + dynamic: bool | None = None, + mode: str | None = None, + ) -> nn.Module: + """Return ``torch.compile(self, ...)`` of this monomorphic module. + + Energy-only and ``method='func'`` force paths are the intended targets. + Prefer fixed shapes (``PadMolecularBatch``) for ``fullgraph=True`` / + ``mode='reduce-overhead'`` (CUDA graphs):: + + model = PiNetPotential(..., compute_forces=True, method="func") + model = model.compile(fullgraph=True) # OptimizedModule + + Also works via :class:`molix.compile.Compiler` / + ``Trainer.compile(...)``. + """ + return torch.compile( + self, + backend=backend, + fullgraph=fullgraph, + dynamic=dynamic, + mode=mode, + ) diff --git a/src/molzoo/pinet/spec.py b/src/molzoo/pinet/spec.py index 48f17bf..6c0e441 100644 --- a/src/molzoo/pinet/spec.py +++ b/src/molzoo/pinet/spec.py @@ -26,3 +26,6 @@ class PiNetSpec(BaseModel): activation: str = "tanh" weighted: bool = False rank: Literal[1, 3, 5] = 3 + #: Emit the per-block ``p3``/``p5`` and ``i1``/``i3``/``i5`` tracks that the + #: PiNet property heads consume. See :class:`~molzoo.pinet.encoder.PiNet`. + emit_property_features: bool = True diff --git a/src/molzoo/specs/mace.md b/src/molzoo/specs/mace.md index 6a3cb86..6c46c55 100644 --- a/src/molzoo/specs/mace.md +++ b/src/molzoo/specs/mace.md @@ -5,8 +5,8 @@ not a tutorial; use the MolZoo user guide for narrative and worked examples. | Field | Value | |-------|-------| -| Module | `molzoo.mace` | -| Entry point | `MACE` (config `MACESpec`) | +| Module | `molzoo.mace.research` (in the `molzoo.mace` package) | +| Entry point | `MACE` (config `MACEResearchSpec`, same module); re-exported lazily as `molzoo.MACE` | | Paper | Batatia et al., *MACE: Higher Order Equivariant Message Passing Neural Networks for Fast and Accurate Force Fields*, NeurIPS 2022 | | arXiv | https://arxiv.org/abs/2206.07697 | | DOI | not applicable (NeurIPS proceedings) | @@ -17,9 +17,12 @@ not a tutorial; use the MolZoo user guide for narrative and worked examples. | Module | Role | |--------|------| -| `molzoo.mace_omol.MACEOMol` | Full energy/force model (OMOL weights); see `mace_omol.md` | +| `molzoo.mace.potential.MACEPotential` | Full energy/force model over both foundation variants | +| `molzoo.mace.variants.MACEOMol` | Thin `MACEPotential` alias for the OMOL weights; see `mace_omol.md` | +| `molzoo.mace.variants.MACEMatpes` | Thin `MACEPotential` alias for the MatPES/MP weights; see `mace_matpes.md` | +| `molzoo.mace.checkpoint.CheckpointRemap` | Official-weight key remap (`MATPES_REMAP` / `OMOL_REMAP` presets) | | `molpot.composition` / heads / derivation | Energy readout, forces, composition | -| `molix.data.NeighborList` | Cutoff graph construction | +| `molix.data.tasks.NeighborList` | Cutoff graph construction | ## 1. Scope @@ -38,7 +41,8 @@ It does **not** own: - training loops or losses - OMOL weight import (that is `MACEOMol`) -Those are owned by `molix` and `molpot` (or `mace_omol` for the full-model path). +Those are owned by `molix` and `molpot` (or `molzoo.mace.potential.MACEPotential` +for the full-model path). ## 2. Public Contract @@ -64,7 +68,7 @@ Those are owned by `molix` and `molpot` (or `mace_omol` for the full-model path) | Symbol | Meaning | Code anchor | |--------|---------|-------------| | \(N, E\) | atoms, edges | batch sizes | -| \(L\) | `num_layers` | `MACESpec` / interaction stack | +| \(L\) | `num_interactions` | `MACEResearchSpec` / interaction stack | | \(F\) | `num_features` | scalar channel multiplicity at \(\ell=0\) | | \(\ell_{\max}\) | `l_max` | spherical harmonics / TP | @@ -83,13 +87,13 @@ kernels (see `MACEOMol`). ## 4. Configuration Contract -| `MACESpec` / ctor field | Meaning | Constraint | +| `MACEResearchSpec` / ctor field | Meaning | Constraint | |-------------------------|---------|------------| | `node_attr_specs` | Discrete/continuous embeddings (e.g. Z) | non-empty | | `num_elements` | Species table size | > 0 | | `num_features` | Channel multiplicity | > 0 | | `r_max` | Radial cutoff (Å) | > 0 | -| `num_layers` | Message-passing depth | ≥ 1 | +| `num_interactions` | Message-passing depth | ≥ 1 | | `l_max` | Angular momentum | ≥ 0 | | `num_bessel` | Radial basis size | > 0 | @@ -138,8 +142,9 @@ No dedicated `bench_mace` yet (tracked as perf work queue). | Concern | Owner | Contract | |---------|-------|----------| | Neighbor list | `molix.data.tasks.NeighborList` | cutoff graph, edge convention | -| Encoder features | `molzoo.MACE` | this spec | -| Energy / force | `molpot` or `molzoo.MACEOMol` | not this module | +| Encoder features | `molzoo.mace.research.MACE` | this spec | +| Energy / force | `molpot` or `molzoo.mace.potential.MACEPotential` | not this module | +| Foundation-block reuse | `molrep.interaction.mace` / `molrep.readout.mace` / `molrep.embedding.mace` | MACE-only blocks; owned by `molrep`, not here | | Training | `molix.Trainer` | TrainState namespaces | ## 9. Version Pinning @@ -149,6 +154,7 @@ No dedicated `bench_mace` yet (tracked as perf work queue). | Paper | Batatia et al., NeurIPS 2022 | | Reference repository | `ACEsuit/mace` (pin TBD on next audit) | | Dependencies | `cuequivariance`, `cuequivariance_torch`, `torch>=2.10` | +| Module relocation | `mace-subpackage-restructure` chain, commits `1ddd5ff..e825a51` (merged 2026-08-09): the flat `src/molzoo/mace.py` became the `src/molzoo/mace/` package (`spec` / `geometry` / `encoder` / `potential` / `checkpoint` / `variants` / `research`) and the MACE-only blocks sank into `molrep/interaction/mace/`, `molrep/readout/mace.py`, `molrep/embedding/mace.py`. This encoder is now `molzoo.mace.research.MACE`. | | Public docs mirror | `docs/molzoo/` (encoder docs partial) | ## 10. Drift Policy @@ -160,5 +166,13 @@ breaking for specs. Force backend / `use_fallback` defaults require a note in ## Appendix A. Maintenance Log -- 2026-07-29: Scaffolded from template; encoder-only contract filled from - `src/molzoo/mace.py` + industrial layout notes. +- 2026-07-29: Scaffolded from template; encoder-only contract filled from the + then-flat MACE encoder module (today `src/molzoo/mace/research.py`) + + industrial layout notes. +- 2026-08-09: Anchors re-pointed for the `mace-subpackage-restructure` chain + (`1ddd5ff..e825a51`). Header module / entry point, §1 full-model pointer, + §3.1 and §4 config-class name (`MACESpec` in this file always meant the + research encoder's config, which is now `MACEResearchSpec`; `MACESpec` in + `molzoo.mace.spec` is the *foundation* config — different class), §8 + boundary table, §9 pinning. No section added, removed or renamed; no + contract change. diff --git a/src/molzoo/specs/mace_matpes.md b/src/molzoo/specs/mace_matpes.md new file mode 100644 index 0000000..d6bc1fd --- /dev/null +++ b/src/molzoo/specs/mace_matpes.md @@ -0,0 +1,360 @@ +# MACE-MatPES Specification + +This page is the implementation contract for `molzoo.MACEMatpes` — a native +reimplementation of the official `MACE-matpes-*-0` foundation models. It is not a +tutorial; use the MolZoo user guide for narrative and worked examples. + +| Field | Value | +|-------|-------| +| Module | `molzoo.mace.variants` (in the `molzoo.mace` package) | +| Entry point | `MACEMatpes` — a thin, keyword-compatible alias over `molzoo.mace.potential.MACEPotential` configured by `molzoo.mace.spec.MACEMatpesSpec`. Weights: `MACEPotential.from_checkpoint` with the `molzoo.mace.checkpoint.MATPES_REMAP` preset; `load_matpes_state_dict` stays as a back-compat free function in `molzoo.mace.variants`. | +| Paper | Batatia et al., *MACE*, NeurIPS 2022; Batatia et al., *A foundation model for atomistic materials chemistry* (MACE-MP-0) | +| arXiv | https://arxiv.org/abs/2206.07697 · https://arxiv.org/abs/2401.00096 | +| Dataset paper | Kaplan et al., *MatPES* — https://arxiv.org/abs/2503.04070 | +| Reference implementation | `ACEsuit/mace` v0.3.16 (`ScaleShiftMACE`) | +| Reference checkpoint | `MACE-matpes-r2scan-omat-ft.model` (ACEsuit/mace-foundations, tag `mace_matpes_0`, ASL licence) | +| Spec status | partial | + +**Related modules (out of scope for this file):** + +| Module | Role | +|--------|------| +| `molzoo.mace.research.MACE` | Encoder-only trainable MACE; see `mace.md` | +| `molzoo.mace.variants.MACEOMol` | OMOL foundation model (charge/spin conditioned); see `mace_omol.md` | +| `molzoo.mace.encoder.MACEEncoder` | Shared block graph both variants inherit (`MACEPotential` subclasses it) | +| `molix.data.tasks.NeighborList` | Cutoff graph construction (PBC minimum image) | + +## 1. Scope + +`molzoo.MACEMatpes` is a **full energy/force model**. It owns: + +- the embedding → interaction → product → readout stack of `ScaleShiftMACE` + with density-normalised interactions, +- the frozen per-element reference energy `E0` and the global scale/shift, +- the ZBL pair-repulsion term, +- strict import of an official cueq-converted `state_dict`. + +It does **not** own: + +- neighbour-list construction or PBC shift vectors (caller supplies both), +- the e3nn → cueq weight conversion (out of tree, `mace.cli.convert_e3nn_cueq`), +- training loops, losses, or MD integration. + +## 2. Public Contract + +### 2.1 Required Inputs + +| Direction | Path / argument | Shape | Dtype | Contract | +|-----------|-----------------|-------|-------|----------| +| In | `("atoms", "Z")` / `Z=` | `(N,)` | int64 | Atomic numbers; must all be in the model's z-table | +| In | `("atoms", "pos")` / `positions=` | `(N, 3)` | float | Cartesian positions, Å | +| In | `("atoms", "batch")` / `batch=` | `(N,)` | int64 | Graph membership | +| In | `("edges", "edge_index")` | `(E, 2)` | int64 | Col 0 = source, col 1 = target | +| In | `edge_index=` (raw API) | `(E, 2)` | int64 | `[:, 0]` = source, `[:, 1]` = target (repo edge convention) | +| In | `("edges", "shifts")` / `shifts=` | `(E, 3)` | float | Optional PBC shift `unit_shifts @ cell` | + +### 2.2 Outputs + +| Direction | Path / return | Shape | Contract | +|-----------|---------------|-------|----------| +| Out | `("graphs", "energy")` / `["energy"]` | `(B,)` | Total energy, eV | +| Out | `("atoms", "forces")` / `["forces"]` | `(N, 3)` | `-∂E/∂pos`, eV/Å | + +## 3. Forward Contract + +### 3.1 Notation + +| Symbol | Meaning | Code anchor | +|--------|---------|-------------| +| \(N, E, B\) | atoms, edges, graphs | batch sizes | +| \(F\) | `num_features` = 128 | `hidden_irreps` 0e multiplicity (`molzoo.mace.spec.MACEMatpesSpec`) | +| \(\ell_{\max}\) | `l_max` = 3 | spherical harmonics / TP | +| \(\nu\) | `correlation` = 3 | symmetric-contraction degree | +| \(\rho_i\) | learned edge density | `molrep.interaction.mace.density.DensityInteraction._message` | +| \(E(\mathbf r)\) | per-graph energy `(B,)`, eV | `molzoo.mace.potential.MACEPotential.energy_core` | +| \(\mathbf v_e\) | edge displacement (with PBC shifts) | `molzoo.mace.geometry.edge_vectors` / `edge_lengths` | + +### 3.2 Pipeline + +``` +E = E0[Z] + scale · ( V_ZBL + Σ_l readout_l(h_l) ) + +r_ij, cutoff u(r_ij) PolynomialCutoff(p=5) on the RAW r +r̃_ij = Agnesi(r_ij; Z_i, Z_j) element-pair radial transform +edge_feats = Bessel(r̃_ij) · u(r_ij) 10 trainable Bessel channels +edge_attrs = Y_l(r_ij vector) l = 0..3 + +h_0 = Linear(one_hot(Z)) (89x0e -> 128x0e) +layer 0 DensityInteraction -> EquivariantProductBasis(use_sc=False) + readout_0 = LinearReadout(128x0e+128x1o) +layer 1 DensityResidualInteraction -> EquivariantProductBasis(use_sc=True) + readout_1 = NonLinearReadout(128x0e, MLP 16) +``` + +Density normalisation (the MatPES/MP-family divergence from stock MACE): + +\[ +\rho_i = \sum_{j \in \mathcal{N}(i)} \tanh\!\big(\mathrm{MLP}(\mathbf{e}_{ij})^2\big), +\qquad +\mathbf{m}_i \leftarrow \frac{\mathrm{Linear}(\mathbf{m}_i)}{\rho_i + 1} +\] + +replacing MACE's fixed `avg_num_neighbors` divisor. `apply_cutoff=True` in the +reference config, so the envelope is folded into `edge_feats` and the +interaction receives `cutoff=None`. + +**Ordering that matters:** the cutoff is evaluated on the **untransformed** +distance and the Bessel basis on the **transformed** one. Swapping them silently +changes every edge feature. + +## 4. Configuration Contract + +| Ctor field | MatPES-r2SCAN value | Meaning | +|------------|--------------------|---------| +| `atomic_numbers` | 89 elements | z-table, checkpoint order | +| `atomic_energies` | 89 floats | frozen `E0`, same order | +| `r_max` | 6.0 | radial cutoff, Å | +| `num_bessel` | 10 | trainable Bessel channels | +| `num_polynomial_cutoff` | 5 | envelope exponent (also ZBL's) | +| `l_max` | 3 | spherical-harmonics order | +| `num_features` | 128 | scalar multiplicity | +| `max_hidden_l` | 1 | node state is `128x0e+128x1o` | +| `num_interactions` | 2 | ≥ 2 (first + last layer differ) | +| `correlation` | 3 | body order | +| `mlp_dim` | 16 | `MLP_irreps` | +| `radial_mlp` | `[64, 64, 64]` | conv-weight MLP widths | +| `scale` / `shift` | 0.7735790334431056 / 0.0 | `atomic_inter_scale/shift` | +| `use_fallback` | `False` (default; fused kernels) — pass `True` on CPU / without the `cuequivariance-ops-torch` wheel | forces are always autograd here, so the functorch reason for the fallback never applies; fused is ~36x faster per MD step | + +## 5. Reference Crosswalk + +| Reference component | MolNex anchor | Status | +|---------------------|---------------|--------| +| `LinearNodeEmbeddingBlock` | `molzoo.mace.encoder.MACEEncoder.node_embedding` (`cuet.Linear`) | matched | +| `RadialEmbeddingBlock.bessel_fn` | `molrep.embedding.radial.BesselRBF(normalize=False, eps=0, trainable=True)` | matched | +| `AgnesiTransform` | `molrep.embedding.radial.AgnesiTransform` | matched | +| `PolynomialCutoff` | `molrep.embedding.cutoff.PolynomialCutoff` | matched | +| `ZBLBasis` | `molpot.potentials.repulsion.ZBLRepulsion` | matched | +| `AtomicEnergiesBlock` | `molpot.heads.energy.AtomicReferenceEnergy` (Z-indexed) | adapted (A1) | +| `RealAgnosticDensityInteractionBlock` | `molrep.interaction.mace.density.DensityInteraction` | matched | +| `RealAgnosticDensityResidualInteractionBlock` | `molrep.interaction.mace.density.DensityResidualInteraction` | matched | +| `EquivariantProductBasisBlock` | `molrep.interaction.product_basis.EquivariantProductBasis` (shared with non-MACE models — **not** moved into the `mace` namespace) | matched | +| `LinearReadoutBlock` | `molrep.readout.mace.LinearReadout` | matched | +| `NonLinearReadoutBlock` | `molrep.readout.mace.NonLinearReadout` | matched | +| `e3nn.nn.FullyConnectedNet` | `molrep.embedding.mlp.MomentNormalizedMLP` (generic; not moved) | matched | +| `ScaleShiftBlock` | `molpot.heads.rescale.GlobalRescale` | matched | +| `get_outputs` force path | autograd, two entries: `molpot.derivation.kernels.grad_force_pass` on the batch path (`MACEPotential.forward`) and `molpot.derivation.force.autograd_forces_from_energy` on the raw path (`MACEMatpes.energy_forces`) | matched | + +`molrep.interaction.density` and `molrep.readout.scalar` still exist as +deprecated re-export shims (`molrep.readout.product` likewise). New code must +import from the `mace` namespaces above; the shims are scheduled for removal. + +## 6. MolNex Adaptations + +| ID | Adaptation | Reason | Risk | Validation | +|----|------------|--------|------|------------| +| A1 | `E0` stored Z-indexed, not element-table-indexed | encoder takes raw `Z` | low | §7.1 parity | +| A2 | Covalent radii inlined (Cordero 2008) rather than read from `molpy.Element` | keeps a core `molrep` block free of the compiled molrs extension; molpy stores radii in fp32 | low | table asserted equal to the checkpoint buffer to 0.0 | +| A3 | cuEq `O3` group, not MACE's `O3_e3nn` | same route the OMOL port validated; CG conventions differ by ~1e-8/op | low | §7.1 parity | +| A4 | Weight import consumes a cueq-converted `state_dict`; conversion runs out of tree | molnex must not import `mace-torch` | low | strict loader, §7.2 | +| A5 | `skip_tp` is rebuilt when the module dtype changes | `cuet.FullyConnectedTensorProduct` bakes its precision in at construction, so `.double()` otherwise raises | low | `tests/test_molrep/test_interaction/test_mace/test_density.py::TestDensityInteraction::test_skip_tp_weight_survives_dtype_change` | +| A6 | `skip_tp` pinned to cuEq `method="naive"` (`molrep.interaction.mace.density.SKIP_TP_METHOD`) | cuEq's default `fused_tp` is 19x slower on a one-hot `89x0e` second operand (31.9 vs 1.8 ms); `naive` is also what MACE's own converted cueq model runs | low | output identical to `fused_tp` at 1.1e-15; §7.1 | +| A7 | Energy and forces come from one forward differentiated in place, not a re-run closure | MACE's `get_outputs` shape; the closure form cost two full forwards per step | none (same math) | §7.2 forward↔energy_forces tests | + +**Known limitation.** cuEquivariance freezes `math_dtype` at construction, and +`EquivariantProductBasis` (shared with `MACEOMol`) is not rebuilt on `.double()`. +A model built under the default fp32 and then `.double()`-d therefore still +contracts in float32 (~1e-8 eV noise). For fp64 work call +`molix.config.set_precision("fp64")` **before** construction. + +## 7. Validation Contract + +### 7.1 Research Reproduction + +Native model vs the official `MACE-matpes-r2scan-0` (e3nn, fp64), identical +neighbour lists, on CPU: + +| System | ΔE/atom (eV) | max ΔF (eV/Å) | +|---|---|---| +| Si diamond 2×1×1 (pbc) | 3.4e-08 | 6.6e-08 | +| NaCl rocksalt (pbc) | 3.6e-08 | 8.6e-10 | +| Fe bcc (pbc) | 4.3e-07 | 2.7e-08 | +| H₂O cluster (no pbc) | 3.2e-08 | 2.0e-06 | +| TiO₂ rutile (pbc) | 1.0e-07 | 8.1e-08 | + +Worst case **4.3e-07 eV/atom, 2.0e-06 eV/Å** — three orders inside the 1e-4 bar +the OMOL port set, and the same magnitude as that port's residual (7e-7 eV / +4.3e-6 eV/Å). The residual is float64 accumulation order between the e3nn and +cueq contraction paths, not a modelling difference. + +**Along a real NVE trajectory** (`mace-r2san-gh200`, 193-atom water/H3O+ box, +200 steps of 0.5 fs, fp64, fused cuEq kernels on GH200), against the official +model driven through the *same* integrator, neighbour list and initial +conditions (`mace-r2san-torch-gh200`): + +| Quantity | Worst over 200 steps | +|---|---| +| ΔE_pot/atom (both arms' own energies) | 1.8e-07 eV | +| Δ\|F\| | 6.4e-06 eV/Å | +| Δ position | 9.5e-07 Å (= float32 trajectory-storage resolution) | +| E_tot drift | -0.0119 meV/atom, **identical** in both arms | +| Mean T | 321.7 K, range 281.3-353.8 K, **identical** in both arms | + +Re-evaluating the official model on molnex's own frames gives 2.3e-07 eV/atom +and 5.0e-05 eV/Å; that force figure is dominated by the trajectory's float32 +position storage (≈ Hessian × 1e-6 Å), not by the model. + +**Against MACE's own cueq model** — the same kernels rather than the e3nn +reference — agreement is essentially bit-level, which isolates the numbers above +as e3nn-vs-cueq contraction order rather than a modelling difference: + +| Comparison (193-atom box, fp64, GH200) | ΔE | max ΔF | +|---|---|---| +| molnex vs `convert_e3nn_cueq(official)` | 4.5e-13 eV | 3.3e-14 eV/Å | +| `skip_tp` naive vs fused_tp, whole model | 2.3e-13 eV | 5.6e-15 eV/Å | + +### 7.2 Symmetry and Shape Tests + +Variant-level claims live in +`tests/test_molzoo/test_mace/test_variants.py::TestMACEMatpes` (the `::…` rows +below are relative to it); the shared pipeline they alias is covered by +`tests/test_molzoo/test_mace/test_potential.py::TestMACEPotential`. + +| Claim | Test path | Tolerance | +|-------|-----------|-----------| +| `forward(td)` == `energy_forces` | `tests/test_molzoo/test_mace/test_variants.py::TestMACEMatpes::test_forward_matches_energy_forces` | 1e-9 eV / 1e-8 eV/Å | +| energy rotation invariance | `::test_energy_is_rotation_invariant` | 1e-9 eV | +| force equivariance `F(Rx) = R F(x)` | `::test_forces_rotate_with_the_system` | 1e-8 | +| `F = -dE/dpos` vs finite differences | `::test_forces_match_finite_differences` | 1e-5 | +| net force ≈ 0 | `::test_net_force_vanishes_on_an_isolated_cluster` | 1e-8 | +| energy extensive over separated graphs | `::test_energy_is_extensive_over_separated_graphs` | 1e-9 | +| out-of-table `Z` rejected | `::test_rejects_atomic_numbers_outside_the_table`, `::test_forward_rejects_an_element_outside_the_table` | raises | +| alien checkpoint rejected | `::test_load_matpes_state_dict_refuses_an_alien_checkpoint` | raises | +| `energy_core` == `forward` energy | `tests/test_molzoo/test_mace/test_potential.py::TestMACEPotential::test_energy_core_agrees_with_the_forward_energy` | exact | +| loader is strict (unknown / missing / mis-shaped) | `tests/test_molzoo/test_mace/test_checkpoint.py::TestLoadMatpesStateDict`, `::TestCheckpointRemap` | raises | +| `from_checkpoint` builds from config + weights | `tests/test_molzoo/test_mace/test_checkpoint.py::TestFromCheckpoint` | exact | +| density normalisation, skip placement | `tests/test_molrep/test_interaction/test_mace/test_density.py` | exact | +| ZBL sign, decay, envelope, halving | `tests/test_molpot/test_potentials/test_repulsion.py` | exact | +| Agnesi monotonicity, pair symmetry | `tests/test_molrep/test_embedding/test_radial.py::TestAgnesiTransform` | exact | +| covalent table == checkpoint buffer | `tests/test_molrep/test_embedding/test_covalent.py::TestCovalentRadii` | 0.0 | + +### 7.3 Engineering Benchmark + +NVE on `wat64_h3o+` (193 atoms, 12.432 Å cubic, r_max 6.0, 17344 directed +edges) through `molix.md` on one GH200. Energy + forces, fp64, eager: + +| Configuration | ms/step | +|---|---| +| initial port (two forwards, `skip_tp` on cuEq's default `fused_tp`) | 166.0 | +| one forward differentiated in place (A7) | 117.2 | +| + `skip_tp` on `naive` (A6) | **36.9** | +| MACE's own cueq model, same graph | 36.0 | + +Block breakdown that located it (fp64, parts sum to the layer): `interaction[0]` +34.4 ms of a 53.3 ms forward, of which `skip_tp` 33.0 ms and the actual message +passing `conv_tp` 0.7 ms. + +The measurements are at 193 atoms; `fused_tp` may win at much larger node +counts, which is why `skip_tp_method` is a constructor argument rather than a +constant. + +#### `torch.compile` strategy + +Eager is launch-bound at this size, so collapsing the launches is the remaining +lever. Compiling `MACEPotential.energy_core` (the pure ``positions -> energy`` +function; measured under its former private name `_compute_energy`, which +survives as a name alias on `MACEMatpes`) and taking ``autograd.grad`` +**outside** the compiled region: + +| | fp64 ms (step/s) | fp32 ms (step/s) | +|---|---|---| +| eager, fused cuEq | 35.7 (28.0) | 35.0 (28.6) | +| `torch.compile` default inductor | 19.0 (52.5) | 19.6 (50.9) | +| **inductor + `reduce-overhead`** | **4.18 (239)** | **2.40 (417)** | +| + `fullgraph=True` | 4.19 (239) | 2.40 (417) | + +**Recommended: `molix.compile.Compiler(cuda_graphs=True)`** — molnex's existing +preset (inductor + `reduce-overhead` + `dynamic=False` + `fullgraph=True`) +transfers to MACE unchanged. Requires static shapes, which the frozen +neighbour list of an MD run already provides. + +Two facts worth keeping: + +* **cuEquivariance traces cleanly**: ``graph_breaks=0, graphs=1, ops=381``. The + fused cuEq kernels are custom autograd Functions and might have forced a graph + break; they do not, so ``fullgraph=True`` is free (identical timing — there + were no breaks to close). +* **Precision only starts mattering once the launches are gone.** Eager fp32 and + fp64 are indistinguishable (35.0 vs 35.7 ms) because the step is latency-bound; + under CUDA graphs fp32 is 1.75x fp64 (2.40 vs 4.18 ms). Compiling in fp64 is + numerically free (ΔE = 0, ΔF = 2.6e-14 vs eager). fp32 costs ~2e-3 eV and + ~1.6e-3 eV/Å against fp32 eager — that is float32 reassociation under inductor + fusion, not a compile defect, and it is the usual precision MACE foundation + models are run at for MD. + +### 7.4 Run Log + +| run_id | date | commit | dirty | dataset | config | steps | train_mae | val_mae | fwd_ms | bwd_ms | compiled | note | +|--------|------|--------|-------|---------|--------|-------|-----------|---------|--------|--------|----------|------| +| 1 | 2026-08-07 | 82c3091 | 1 | 5 crystals + H2O | official weights, fp64, CPU | 0 | n/a | n/a | n/a | n/a | no | single-point parity vs official: 4.3e-7 eV/atom, 2.0e-6 eV/Å (§7.1) | +| 2 | 2026-08-07 | 82c3091 | 1 | wat64_h3o+ | official weights, fp64, CPU, 1 step | 1 | n/a | n/a | n/a | n/a | no | `mace-r2san-cpu` smoke: E=-1015.900312 eV, \|F\|max=3.287 eV/Å, T=306 K, ~32 s/step (loaded login node) | +| 3 | 2026-08-07 | 82c3091 | 1 | wat64_h3o+ | official weights, fp64, GH200, fused cuEq | 200 | n/a | n/a | 170 | incl. | no | `mace-r2san-gh200` NVE dt=0.5 fs: same single point as CPU to all printed digits; E_tot drift -0.0119 meV/atom over 100 fs; T 281-354 K; 0.17 s/step | +| 4 | 2026-08-07 | 82c3091 | 1 | wat64_h3o+ | official mace-torch, fp64, GH200 | 200 | n/a | n/a | 35 | incl. | no | `mace-r2san-torch-gh200` control arm: identical driver/NL/ICs; E=-1015.900277 eV, same drift and T range; 0.035 s/step | +| 5 | 2026-08-07 | 82c3091 | 1 | wat64_h3o+ | official weights, fp64, GH200, after A6+A7 | 0 | n/a | n/a | 36.9 | incl. | no | energy+forces 166 -> 36.9 ms/step (4.5x), vs 36.0 ms for MACE's own cueq model; ΔE vs that model 4.5e-13 eV | +| 6 | 2026-08-07 | 82c3091 | 1 | wat64_h3o+ | fp64, GH200, inductor reduce-overhead | 0 | n/a | n/a | 4.18 | incl. | yes | 239 step/s, 8.5x eager; ΔE=0 ΔF=2.6e-14 vs eager; graph_breaks=0 | +| 7 | 2026-08-07 | 82c3091 | 1 | wat64_h3o+ | fp32, GH200, inductor reduce-overhead | 0 | n/a | n/a | 2.40 | incl. | yes | 417 step/s, 14.6x eager; fp32 reassociation ~2e-3 eV vs fp32 eager | +| 8 | 2026-08-07 | 82c3091 | 1 | wat64_h3o+ | fp64, GH200, full NVE loop `--compile` | 200 | n/a | n/a | 5.0 | incl. | yes | end-to-end MD 37.5 -> 5.0 ms/step (7.5x incl. integrator+hook); trajectory bit-identical to eager (max \|dPos\|=0, max \|dF\|=0, ΔE/atom 1.9e-12 meV); one-off compile ~60 s | +| 9 | 2026-08-08 | 82c3091 | 1 | wat64_h3o+ | launch gate: rebuild_every=5 + compile, 3 precisions | 2000 | n/a | n/a | 14.9/11.8/13.2 | incl. | yes | dead-edge dE≤1.1e-13 (atomicAdd reorder only); rebuild-vs-fresh dE=0; drift fp64 0.03 / fp32 0.04 / bf16 ~2 meV/atom/ps (bf16 heats — expected physics); bf16 needs autocast INSIDE the compiled callable (52->13.2 ms) | +| 10 | 2026-08-08 | 82c3091 | 1 | wat64_h3o+ | production 5 ns × 3 precisions (jobs 978984/978985/979020) | 10M | n/a | n/a | — | — | yes | dt=0.5 fs, stride 2000, rebuild_every=5, checkpoint 100k, auto-resume; dirs `mace-r2san-5ns-{fp64,fp32,bf16}` | +| 11 | 2026-08-09 | 43cd33d | 0 | wat64_h3o+ | official weights, fp64, GH200, full NVE `--compile`, post mace-subpackage-restructure + fix batch (dev unified) | 200 | n/a | n/a | 8.6 | incl. | yes | `mace-nve-validate-1017194` post-restructure gate: single point E=-1015.900312 eV, \|F\|max=3.2866 eV/Å — matches rows 2/3 to all printed digits; E_tot drift -0.0146 meV/atom over 0.10 ps (-0.146 meV/atom/ps, order of row 3); T 283.7-353.7 K; rebuilds 41 (~40 expected at rebuild_every=5); sbatch work/mace-nve/run_gh200_validation.sbatch, job 1017194 | +| 12 | 2026-08-09 | 7ff808d | 0 | wat64_h3o+ | precision matrix fp64/fp32 x MD/inference, GH200, `--compile`, post md-neighborlist-skin chain (list-owned policy, skin=0.2) | 2000 | n/a | n/a | 7.5/4.7/4.0 | incl. | yes | job 1042250: single point E(64/64)=-1015.900312 eV matches row 11 to all printed digits (chain is physics-neutral); E(md64/pot32)=-1015.900083 (0.23 meV off — fp64 geometry + fp32 inference), E(32/32)=-1016.004953 (105 meV off — fp32 geometry dominates, not weights); drift 0.040/0.023/0.054 meV/atom/ps over 1 ps; rebuilds 508/507/508 of 2000, ndanger=0 all arms; mixed arm = fp64-grade accuracy at 1.6x fp64 speed; sbatch benchmarks/run_gh200_mace_nve_precision.sbatch | + +## 8. System Boundary + +| Concern | Owner | Contract | +|---------|-------|----------| +| Neighbour list + PBC shifts | caller / `molix.data.tasks.NeighborList` | `(E,2)` edges; `shifts = mic_diff - (pos[t]-pos[s])` | +| Edge displacements | `molzoo.mace.geometry` | `edge_vectors` / `edge_lengths`; differentiable w.r.t. `pos` | +| Building blocks | `molrep.interaction.mace` / `molrep.readout.mace` / `molrep.embedding.mace` (+ generic `molrep` / `molpot`) | reused; not owned here | +| Model graph | `molzoo.mace.encoder.MACEEncoder` | blocks + wiring; no energy, no forces | +| Energy / forces | `molzoo.mace.potential.MACEPotential` | `energy_core` (public, compile seam) + `molpot.derivation.kernels.grad_force_pass` | +| Configuration | `molzoo.mace.spec.MACEMatpesSpec` | torch-free pydantic preset | +| Weight conversion | `mace.cli.convert_e3nn_cueq` (out of tree) | e3nn `.model` → cueq `state_dict` | +| Weight import | `molzoo.mace.checkpoint.CheckpointRemap` (`MATPES_REMAP`) via `MACEPotential.from_checkpoint`; `load_matpes_state_dict` back-compat wrapper | strict; raises on any unmapped or unfilled tensor | +| MD | `molix.md` | frozen neighbour list, `gamma=0` → NVE | +| Compile seam | `molix.compile.Compiler` | wraps `energy_core`; `autograd.grad` stays outside | +| Lazy export | `molzoo/__init__` + `molzoo/mace/__init__` | PEP 562 `__getattr__`; no eager cueq import | + +## 9. Version Pinning + +| Item | Value | +|------|-------| +| Reference repository | `ACEsuit/mace` v0.3.16 | +| Checkpoint | `MACE-matpes-r2scan-omat-ft.model`, `mace_matpes_0` release | +| Checkpoint config | `correlation=3`, `use_reduced_cg=False`, `use_agnostic_product=False`, `apply_cutoff=True`, `pair_repulsion=True`, `distance_transform=Agnesi`, `heads=['default']` | +| Dependencies | `cuequivariance` 0.10.0, `cuequivariance_torch`, `torch>=2.10` | +| Conversion oracle | e3nn 0.4.4 + mace-torch 0.3.16 in an out-of-tree venv | +| Module relocation | `mace-subpackage-restructure` chain, commits `1ddd5ff..e825a51` (merged 2026-08-09): `src/molzoo/mace_matpes.py` was retired into the `src/molzoo/mace/` package — config in `spec.py`, blocks in `encoder.py`, energy/forces in `potential.py`, key remap in `checkpoint.py`, the `MACEMatpes` alias in `variants.py`. MACE-only `molrep` blocks moved to `molrep/interaction/mace/{conv,block,density}.py`, `molrep/readout/mace.py`, `molrep/embedding/mace.py`; `molrep.interaction.density` / `molrep.readout.scalar` / `molrep.readout.product` remain as deprecated shims. Tests moved to `tests/test_molzoo/test_mace/`. Weights, hyper-parameters and numerics unchanged (§7.1 not re-run). | + +## 10. Drift Policy + +Any change to the interaction schedule, the readout placement, the density +normalisation, or the checkpoint key map **must** update this file in the same +PR and re-run §7.1. `use_reduced_cg` / `original_mace` and the `apply_cutoff` +ordering are load-bearing for weight compatibility — changing either silently +produces a model that runs and is wrong. + +## Appendix A. Maintenance Log + +- 2026-08-07: Created alongside the native port; §2/§3/§5 filled from the paper + and `ACEsuit/mace` v0.3.16; §7.1 filled from the out-of-tree parity oracle. +- 2026-08-09: Anchors re-pointed for the `mace-subpackage-restructure` chain + (`1ddd5ff..e825a51`) — header, §3.1 code anchors, §5 crosswalk, §6 A5/A6, §7.2 + test paths, §7.3 compile target, §8 boundary, §9 pinning row. Two content + corrections found while re-pointing: the §5 force row named + `ForceDerivation(method="autograd")`, but the MACE path actually runs + `molpot.derivation.kernels.grad_force_pass` (batch) and + `molpot.derivation.force.autograd_forces_from_energy` (raw) — same autograd + math, different symbol; and one force tolerance in §7.1/§7.2 was written + `eV·Å` instead of `eV/Å`. No section added, removed or renamed; §7.4 rows + untouched; no numerical claim changed. diff --git a/src/molzoo/specs/mace_omol.md b/src/molzoo/specs/mace_omol.md index 4e2f9c1..08d7adb 100644 --- a/src/molzoo/specs/mace_omol.md +++ b/src/molzoo/specs/mace_omol.md @@ -1,12 +1,13 @@ # MACEOMol Specification -This page is the implementation contract for `molzoo.mace_omol`. It is not a -tutorial; use the MolZoo user guide for theory narrative and worked examples. +This page is the implementation contract for `molzoo.mace.variants.MACEOMol`. +It is not a tutorial; use the MolZoo user guide for theory narrative and worked +examples. | Field | Value | |-------|-------| -| Module | `molzoo.mace_omol` | -| Entry point | `MACEOMol` (plain `nn.Module`; constructor kwargs, no pydantic Spec) | +| Module | `molzoo.mace.variants` (in the `molzoo.mace` package) | +| Entry point | `MACEOMol` — a thin, keyword-compatible alias over `molzoo.mace.potential.MACEPotential`. The constructor still takes the OMOL keywords directly; internally it builds a `molzoo.mace.spec.MACEOMolSpec`. Weights: `MACEPotential.from_checkpoint` with the `molzoo.mace.checkpoint.OMOL_REMAP` preset; `load_omol_state_dict` stays as a back-compat free function in `molzoo.mace.variants`. | | Paper | Batatia et al., "MACE: Higher Order Equivariant Message Passing Neural Networks for Fast and Accurate Force Fields", NeurIPS 2022. OMol25 / MACE-omol-0 foundation model. | | arXiv | https://arxiv.org/abs/2206.07697 | | DOI | not applicable | @@ -20,9 +21,14 @@ tutorial; use the MolZoo user guide for theory narrative and worked examples. `l_max=3`, 3 residual interactions, product `correlation=2`, 83 elements, single `omol` head, with `total_charge` / `total_spin` conditioning, 52.7M parameters). -It **owns**: the full energy/force forward of MACE-OMOL, assembled from generic -`molrep` / `molpot` blocks (no MACE prefix), and the `load_omol_state_dict` -converter that imports official weights (after `mace.cli.convert_e3nn_cueq`). +It **owns**: the OMOL preset of the full energy/force forward — the block graph +in `molzoo.mace.encoder`, the energy/force pipeline in +`molzoo.mace.potential`, and the `OMOL_REMAP` key dialect in +`molzoo.mace.checkpoint` that imports official weights (after +`mace.cli.convert_e3nn_cueq`), reached through `load_omol_state_dict` or +`MACEPotential.from_checkpoint`. The blocks themselves come from `molrep` / +`molpot` — some generic (`ResidualInteraction`, `EquivariantProductBasis`), +some in the MACE-only namespaces (`molrep.readout.mace`). It does **not** own: the building blocks themselves (they live in `molrep` / `molpot` and are reused by other models), neighbor-list construction, dataset / @@ -43,8 +49,8 @@ not an encoder-only feature extractor: its `forward` writes `graphs.energy` and | In | `atoms.pos` | `(N, 3)` | float (`config.ftype`) | Cartesian positions; forces are `-∂E/∂pos` | | In | `atoms.batch` | `(N,)` | long | Graph membership index `0..B-1` | | In | `edges.edge_index` | `(E, 2)` | long | `[:,0]`=source/sender, `[:,1]`=target/receiver (MolNex convention) | -| In | `graphs.total_charge` | `(B,)` | long | Optional; per-graph total charge. Absent → neutral (0) | -| In | `graphs.total_spin` | `(B,)` | long | Optional; per-graph spin. Absent → singlet (0) | +| In | `graphs.total_charge` | `(B,)` | long | Optional; per-graph total charge in units of `e`. Absent → neutral (`0`) | +| In | `graphs.total_spin` | `(B,)` | long | Optional; per-graph spin channel. Absent → `1` (singlet multiplicity `2S+1 = 1`, all electrons paired). **Not `0`** — `MACEPotential._condition_charge_spin` fills `torch.ones(...)`, and spin `0` would index an untrained embedding row (`spin_offset = 0`) and return garbage | `edges.edge_diff` / `edges.edge_dist` are **not** consumed: `forward` recomputes edge vectors from `atoms.pos` so the energy is differentiable w.r.t. positions. @@ -54,12 +60,14 @@ edge vectors from `atoms.pos` so the energy is differentiable w.r.t. positions. | Direction | TensorDict path | Shape | Written when | Contract | |-----------|------------------|-------|--------------|----------| | Out | `graphs.energy` | `(B,)` | always | Per-graph total energy = E0 + scale·shift(readout) | -| Out | `atoms.forces` | `(N, 3)` | always | `F = -∂E/∂pos` via `molpot.derivation.ForceDerivation` | +| Out | `atoms.forces` | `(N, 3)` | always | `F = -∂E/∂pos` via `molpot.derivation.kernels.grad_force_pass` | `forward` mutates `td` in place (creating the `graphs` sub-dict if absent) and returns the same object. The raw-tensor entry point `MACEOMol.energy_forces(...)` returns `{"energy", "forces"}` and is used by the `scripts/omol_port/verify_*.py` -block/E2E checks. +block/E2E checks; the differentiable energy alone is +`MACEPotential.energy_core(positions, Z, edge_index, batch, num_graphs, +shifts=None, total_charge=None, total_spin=None)`. ## 3. Forward Contract @@ -67,11 +75,11 @@ block/E2E checks. | Symbol | Meaning | Code anchor | |--------|---------|-------------| -| $N, E, B$ | atoms, edges, graphs | `MACEOMol._compute_energy` | +| $N, E, B$ | atoms, edges, graphs | `molzoo.mace.potential.MACEPotential.energy_core` | | $Z_i$ | atomic number of atom $i$ | `atoms.Z` | | $\mathbf r_i$ | position of atom $i$ | `atoms.pos` | -| $s,t$ | sender / receiver of an edge | `edge_index[0]`, `edge_index[1]` | -| $\mathbf v_e=\mathbf r_t-\mathbf r_s$ | edge vector | `_compute_energy` | +| $s,t$ | sender / receiver of an edge | `edge_index[:,0]`, `edge_index[:,1]` | +| $\mathbf v_e=\mathbf r_t-\mathbf r_s$ | edge vector | `molzoo.mace.geometry.edge_vectors` | | $h_i$ | node features (irreps, `cue.ir_mul`) | `node_feats` | | $E_0$ | per-element reference + charge/spin readout | `e0` | | $q,\sigma$ | per-graph total charge / spin | `total_charge`, `total_spin` | @@ -127,9 +135,9 @@ $$ | Quantity | Shape | Code anchor | |----------|-------|-------------| -| readout | `(N,)` | `NonLinearBiasReadout` (`molrep.readout.scalar`) | +| readout | `(N,)` | `NonLinearBiasReadout` (`molrep.readout.mace`) | | scale_shift | `(N,)` | `GlobalRescale` (`molpot.heads.rescale`) | -| $\mathbf F$ | `(N,3)` | `ForceDerivation` (`forward`) / `autograd.grad` (`energy_forces`) | +| $\mathbf F$ | `(N,3)` | `molpot.derivation.kernels.grad_force_pass` (`forward`) / `molpot.derivation.force.autograd_forces_from_energy` (`energy_forces`) — both `torch.autograd.grad` | ## 4. Configuration Contract @@ -161,12 +169,12 @@ $$ | Radial MLP | `modules.radial.RadialMLP` | `molrep.interaction.RadialMLP` | matched | | Gated nonlinearity | e3nn `Gate` | `molrep.interaction.GatedNonlinearity` | matched | | Product basis | `EquivariantProductBasisBlock` (`original_mace=True`) | `molrep.interaction.EquivariantProductBasis` | matched | -| Non-linear readout | `NonLinearBiasReadoutBlock` | `molrep.readout.NonLinearBiasReadout` | matched | +| Non-linear readout | `NonLinearBiasReadoutBlock` | `molrep.readout.mace.NonLinearBiasReadout` (re-exported as `molrep.readout.NonLinearBiasReadout`) | matched | | Per-element E0 | `AtomicEnergiesBlock` | `molpot.heads.AtomicReferenceEnergy` | matched | | Scale/shift | `ScaleShiftBlock` | `molpot.heads.GlobalRescale` | matched | | Charge/spin embed | `GenericJointEmbedding` | `molrep.embedding.JointFeatureEmbedding` | matched | -| Forces | `autograd.grad` | `molpot.derivation.ForceDerivation` (`torch.func.grad`) | adapted | -| Weight import | e3nn `state_dict` | `load_omol_state_dict` (after `convert_e3nn_cueq`) | adapted | +| Forces | `autograd.grad` | `molpot.derivation.kernels.grad_force_pass` / `molpot.derivation.force.autograd_forces_from_energy` (both `torch.autograd.grad`; cuEq's fused ops are legacy `autograd.Function`s that `torch.func.grad` rejects) | matched | +| Weight import | e3nn `state_dict` | `molzoo.mace.checkpoint.OMOL_REMAP` via `load_omol_state_dict` / `MACEPotential.from_checkpoint` (after `convert_e3nn_cueq`) | adapted | ## 6. MolNex Adaptations @@ -174,10 +182,10 @@ $$ |----|------------|--------|------|------------| | A1 | PolynomialCutoff + trainable un-normalised Bessel (`eps=0`) | OMOL variant vs standard MACE | low | `scripts/omol_port/verify_radial.py` (7e-15) | | A2 | charge/spin via `JointFeatureEmbedding` added to node feats + into E0 | OMOL conditioning | low | `scripts/omol_port/verify_joint_embed.py` (0) | -| A3 | cue `"O3"` group everywhere (no `O3_e3nn`) | weights are converted into the cue-O3 twin, so O3 is the native target; O3 vs O3_e3nn CG differ only ~1.4e-8/op and the O3 twin already matches e3nn to 1.5e-8 — O3_e3nn would not reduce the residual and would add an `e3nn` dep | low | residual 7e-7 eV / 4.3e-6 eV·Å vs official (reimplementation accumulation, not a convention diff) — inside the 1e-4 bar (`mace-omol-port-02` ac-003) | -| A4 | TensorDict `forward` forces via `ForceDerivation` (`func.grad`); `energy_forces` via `autograd.grad` | compile-friendly molnex contract | low | `tests/test_molzoo/test_mace_omol.py` (1e-8 vs autograd) | -| A5 | edge convention `v=pos[t]-pos[s]`, `edge_index (E,2)→(2,E)` | MolNex collate schema | low | `tests/test_molzoo/test_mace_omol.py` | -| A6 | `RadialMLP` honors `config.ftype` | fp64-via-config without `.double()` | low | `tests/test_molzoo/test_mace_omol.py` (fp64) | +| A3 | cue `"O3"` group everywhere (no `O3_e3nn`) | weights are converted into the cue-O3 twin, so O3 is the native target; O3 vs O3_e3nn CG differ only ~1.4e-8/op and the O3 twin already matches e3nn to 1.5e-8 — O3_e3nn would not reduce the residual and would add an `e3nn` dep | low | residual 7e-7 eV / 4.3e-6 eV/Å vs official (reimplementation accumulation, not a convention diff) — inside the 1e-4 bar (`mace-omol-port-02` ac-003) | +| A4 | TensorDict `forward` forces via `molpot.derivation.kernels.grad_force_pass`; `energy_forces` via `molpot.derivation.force.autograd_forces_from_energy`. Both are `torch.autograd.grad`; the earlier `torch.func.grad` plan was dropped because cuEquivariance's fused ops register legacy `autograd.Function`s without `setup_context` | compile-friendly molnex contract | low | `tests/test_molzoo/test_mace/test_variants.py::TestMACEOMol::test_forward_matches_energy_forces` (1e-8 vs autograd) | +| A5 | edge convention `v=pos[t]-pos[s]`, `edge_index (E,2)` end to end (upstream's `(2,E)` transposed away at the port boundary) | MolNex collate schema | low | `tests/test_molzoo/test_mace/test_geometry.py::TestEdgeVectors`, `tests/test_molzoo/test_mace/test_variants.py::TestMACEOMol` | +| A6 | `RadialMLP` honors `config.ftype` | fp64-via-config without `.double()` | low | `tests/test_molzoo/test_mace/test_variants.py::TestMACEOMol` (fp64 fixtures) | | A7 | per-layer irreps + edge-mid (128) hardcoded to OMOL dims | faithful OMOL weight load | medium | non-OMOL `l_max≥2`+small `num_features` unsupported; tracked in `mace-omol-port-02` | ## 7. Validation Contract @@ -185,25 +193,54 @@ $$ ### 7.1 Research Reproduction The accepted accuracy bar is E/F within **1e-4** of official OMOL (operator -decision, 2026-06-21). The full model with official OMOL weights reproduces the -official cueq OMOL twin on a charged molecule to **7.0e-7 eV / 4.3e-6 eV·Å** -(`scripts/omol_port/verify_e2e.py`, RESULT: PASS) — three to four orders inside -the bar; the cueq twin itself matches e3nn OMOL to 1.5e-8 eV / 3.2e-8 eV·Å -(`scripts/omol_port/verify_omol_cueq_equiv.py`). The 7e-7 residual is molnex's -own reimplementation accumulation, **not** a CG-convention difference: O3 vs -O3_e3nn Clebsch-Gordan differ only ~1.4e-8/op and the O3 twin already aligns -with e3nn to 1.5e-8, so the e3nn-convention group is neither used nor needed -(A3). +decision, 2026-06-21). + +**Historical record (2026-06-21; oracles deleted in `b85d12f`, not +reproducible in-tree — Appendix A).** Full model with official weights vs the +official cueq OMOL twin, charged molecule: **7.0e-7 eV / 4.3e-6 eV/Å** +(`verify_e2e.py`, RESULT: PASS); the cueq twin vs e3nn OMOL: 1.5e-8 eV / +3.2e-8 eV/Å (`verify_omol_cueq_equiv.py`). Three to four orders inside the +bar. The run predates the 2026-08-07 loader fix (official `bessel_weights` +silently dropped, `bessel.freqs` left at init), so the 7e-7 eV includes that +~2.2e-7 Å⁻¹ perturbation; it is otherwise molnex's own reimplementation +accumulation, **not** a CG-convention difference (O3 vs O3_e3nn CG +~1.4e-8/op, A3). Re-measuring upstream parity needs an out-of-tree oracle +(route per Appendix A, 2026-08-09). + +**Current in-tree verification (2026-08-09).** `MOLNEX_MACE_WEIGHTS_DIR`-gated +`tests/test_molzoo/test_mace/test_checkpoint.py::TestOfficialOMolWeights`: +strict 104-parameter load through `OMOL_REMAP` plus E/F stability goldens on a +five-atom cluster; the weights dump is regenerated offline by +`scripts/omol_port/convert_omol_to_cueq_state.py` (no `mace`/`e3nn`). This is +a stability lock on this machine's own output — not an upstream parity claim. + +**Bessel frequencies (measured 2026-08-09).** The official `bessel_weights` +are bit-for-bit the fp32 evaluation of the analytic init `nπ/r_max` (upcast to +fp64); the offset from the fp64 analytic values (max 2.2120e-7 Å⁻¹ at n=7, +≤ 1 fp32 ulp per entry) is fp32 rounding of an untrained parameter, **not** +fitted drift. Doctrine unchanged: `bessel.freqs` is an `nn.Parameter` and must +be filled from the checkpoint — bit-exactness against the official surface +requires the checkpoint's fp32-rounded values, not the fp64 re-derivation. ### 7.2 Symmetry and Shape Tests +Variant-level claims live in +`tests/test_molzoo/test_mace/test_variants.py::TestMACEOMol` (the `::…` rows +below are relative to it); the shared pipeline it aliases is covered by +`tests/test_molzoo/test_mace/test_potential.py::TestMACEPotential`. + | Claim | Test path | Tolerance | |-------|-----------|-----------| -| `forward(td)` energy == `energy_forces` | `tests/test_molzoo/test_mace_omol.py::test_forward_matches_energy_forces` | 1e-9 eV | -| `forward(td)` forces == `energy_forces` | same | 1e-8 eV·Å | -| net force ≈ 0 (translation invariance) | `::test_forces_translation_invariant` | 1e-7 | +| `forward(td)` energy == `energy_forces` | `tests/test_molzoo/test_mace/test_variants.py::TestMACEOMol::test_forward_matches_energy_forces` | 1e-9 eV | +| `forward(td)` forces == `energy_forces` | same | 1e-8 eV/Å | +| net force ≈ 0 (translation invariance) | `::test_net_force_vanishes_on_an_isolated_molecule` | 1e-7 | | neutral default when `graphs.*` absent | `::test_missing_charge_spin_defaults_to_neutral` | 1e-9 | +| default spin is the closed-shell singlet `1`, not `0` | `tests/test_molzoo/test_mace/test_potential.py::TestMACEPotential::test_omol_defaults_to_a_neutral_closed_shell_singlet` | exact | | per-graph batching | `::test_batched_graphs` | 1e-9 | +| force loss reaches parameters (eval mode / through `forward`) | `::test_force_loss_reaches_parameters_in_eval_mode`, `::test_force_loss_through_forward_reaches_parameters` | exact | +| alien checkpoint rejected | `::test_load_omol_state_dict_refuses_an_alien_checkpoint` | raises | +| `OMOL_REMAP` roundtrip restores every parameter | `tests/test_molzoo/test_mace/test_checkpoint.py::TestCheckpointRemap::test_roundtrip_restores_every_parameter_omol` | exact | +| edge vectors / lengths under PBC shifts | `tests/test_molzoo/test_mace/test_geometry.py` | exact | | block-level vs cueq (radial/mlp/e0/joint/interaction/product/readout) | `scripts/omol_port/verify_*.py` | 0–7e-15 | ### 7.3 Engineering Benchmark @@ -222,12 +259,15 @@ for this spec (tracked in `scripts/omol_port/SPEC.md`). | Concern | Owner | Contract | |---------|-------|----------| -| Neighbor list / edges | collate / `NeighborList` | populates `edges.edge_index` `(E,2)` | +| Neighbor list / edges | collate / `molix.data.tasks.NeighborList` | populates `edges.edge_index` `(E,2)` | +| Edge displacements | `molzoo.mace.geometry` | `edge_vectors` / `edge_lengths`; differentiable w.r.t. `pos` | | Charge/spin inputs | dataset / collate | `graphs.total_charge`, `graphs.total_spin` (optional) | -| Building blocks | `molrep` / `molpot` | reused; not owned here | -| Forces | `molpot.derivation.ForceDerivation` | `F=-∂E/∂pos` from energy closure | -| Weight import | `load_omol_state_dict` + `mace.cli.convert_e3nn_cueq` | cueq `state_dict` → `MACEOMol` | -| Lazy export | `molzoo/__init__` | PEP 562 `__getattr__`; no eager cueq import | +| Building blocks | `molrep.readout.mace` (+ generic `molrep` / `molpot`) | reused; not owned here | +| Model graph | `molzoo.mace.encoder.MACEEncoder` | blocks + wiring; no energy, no forces | +| Energy / forces | `molzoo.mace.potential.MACEPotential` | `energy_core` (public, compile seam) + `molpot.derivation.kernels.grad_force_pass` | +| Configuration | `molzoo.mace.spec.MACEOMolSpec` | torch-free pydantic preset | +| Weight import | `molzoo.mace.checkpoint.OMOL_REMAP` via `load_omol_state_dict` / `MACEPotential.from_checkpoint`, after `mace.cli.convert_e3nn_cueq` | cueq `state_dict` → `MACEOMol`; strict on missing `nn.Parameter`s since 2026-08-07 | +| Lazy export | `molzoo/__init__` + `molzoo/mace/__init__` | PEP 562 `__getattr__`; no eager cueq import | ## 9. Version Pinning @@ -237,7 +277,8 @@ for this spec (tracked in `scripts/omol_port/SPEC.md`). | Reference repository | `ACEsuit/mace` | | Reference commit | not pinned to sha; `mace==0.3.16` (PyPI) used for conversion + verify. Follow-up audit to pin exact sha. | | Dependencies | `torch==2.12.1`, `cuequivariance==0.10.0`, `cuequivariance_torch==0.10.0`, `tensordict==0.13.0` | -| Public docs mirror | `docs/molzoo/specs/mace_omol.md` | +| Module relocation | `mace-subpackage-restructure` chain, commits `1ddd5ff..e825a51` (merged 2026-08-09): `src/molzoo/mace_omol.py` was retired into the `src/molzoo/mace/` package — config in `spec.py`, blocks in `encoder.py`, energy/forces in `potential.py`, key remap in `checkpoint.py`, the `MACEOMol` alias in `variants.py`. MACE-only `molrep` blocks moved to `molrep/interaction/mace/{conv,block,density}.py`, `molrep/readout/mace.py`, `molrep/embedding/mace.py`. Tests moved to `tests/test_molzoo/test_mace/`. Weights, hyper-parameters and numerics unchanged (§7.1 not re-run). | +| Public docs mirror | `docs/molzoo/specs/mace_omol.md` — byte-identical copy of this file; re-sync both halves on every edit | ## 10. Drift Policy @@ -254,8 +295,51 @@ rows. `mace-omol-port-01/02` implementation (status draft → partial). §5 rows `matched` per `scripts/omol_port/verify_*.py`. - 2026-06-21: accuracy bar set to 1e-4 (operator); cue O3 meets it at - 7e-7 eV / 4.3e-6 eV·Å. `mace-omol-port-02` ac-003 verified; chain done. + 7e-7 eV / 4.3e-6 eV/Å. `mace-omol-port-02` ac-003 verified; chain done. - 2026-06-21: measured O3 vs O3_e3nn CG = 1.4e-8/op → the 7e-7 residual is reimplementation accumulation, not a convention diff. Dropped the O3_e3nn pursuit entirely and removed the dead `MACEOMol(group=)` hook from MACEOMol / ResidualInteraction / EquivariantProductBasis (always cue O3). +- 2026-08-09: Anchors re-pointed for the `mace-subpackage-restructure` chain + (`1ddd5ff..e825a51`) — header, §1 ownership, §2.2 / §3.1 / §3.5 code anchors, + §5 crosswalk, §6 A4/A5/A6 verification paths, §7.2 test paths, §8 boundary, + §9 pinning row. Three content corrections found while re-pointing: + (a) `graphs.total_spin` absent defaults to **`1`** (closed-shell singlet + multiplicity `2S+1 = 1`), not `0` — `MACEPotential._condition_charge_spin` + fills `torch.ones(...)`, and `0` would index an untrained embedding row; + (b) force units written `eV·Å` now read `eV/Å` (§6 A3, §7.1, §7.2 and this + log); (c) forces are `torch.autograd.grad` + (`molpot.derivation.kernels.grad_force_pass` / + `molpot.derivation.force.autograd_forces_from_energy`), never + `ForceDerivation(method="functorch")` / `torch.func.grad` — cuEquivariance's + fused ops are legacy `autograd.Function`s that `torch.func.grad` rejects. + No section added, removed or renamed; §7.4 rows untouched; no numerical + claim changed. `docs/molzoo/specs/mace_omol.md` re-synced from this file + (the two copies had drifted on §6 A5, §7.1 and §8). +- 2026-08-09: **Dangling anchors, not fixed here.** Every + `scripts/omol_port/verify_*.py` cited by §5 (`matched` source of truth), §6 + A1/A2, §7.1, §7.2, §7.3 and §10 was deleted in commit `b85d12f`; only + `README.md` and `SPEC.md` remain in that directory. The recorded numbers + (7.0e-7 eV / 4.3e-6 eV/Å, the 0–7e-15 block-level residuals) are therefore + no longer reproducible in-tree, and §10's drift trigger (b) — "`verify_e2e.py` + E/F residual regresses > 10×" — cannot fire. Restoring the oracles (or + re-homing them under `regressions/` with hard-coded goldens) is out of scope + for `mace-subpackage-restructure-07-cleanup`, which is anchor-refresh only. + Route: `/mol:fix` or a follow-up spec. Combined with the 2026-08-07 caveat + in §7.1 (the parity run predates the trainable-Bessel loader fix), §7.1 + should be treated as **stale, pending re-measurement**. +- 2026-08-09: §7.1 rewritten (molzoo-auditor, operator-directed). (a) The + mace-torch parity figures are now labelled a dated **historical record** + (oracles deleted in `b85d12f`), and the current in-tree surface is named: + `MOLNEX_MACE_WEIGHTS_DIR`-gated `TestOfficialOMolWeights` (strict 104-param + load + E/F stability goldens) with the dump regenerable via + `scripts/omol_port/convert_omol_to_cueq_state.py`. (b) The "fitted drift" + reading of `bessel.freqs` is corrected to measurement: the official + `bessel_weights` are **bit-for-bit** the fp32 evaluation of the analytic + `nπ/r_max` init upcast to fp64 (max 2.2120e-7 Å⁻¹ from the fp64 values at + n=7, ≤ 1 fp32 ulp per entry) — storage rounding of an untrained parameter, + not training drift; the strict-loading doctrine is unchanged. A ⚠️ was + printed (not applied) against `src/molzoo/mace/checkpoint.py`'s docstring + ("fitted like any other weight" / "the fitted frequencies were dropped"). + No section added, removed or renamed; §7.4 untouched; §6 A1/A2 and §7.2's + dangling `verify_*.py` anchors remain covered by the entry above. diff --git a/src/molzoo/specs/pinet2.md b/src/molzoo/specs/pinet2.md index 4f6779a..70d8d14 100644 --- a/src/molzoo/specs/pinet2.md +++ b/src/molzoo/specs/pinet2.md @@ -9,7 +9,7 @@ This spec covers the MolNex PyTorch port of Teoroo-CMC/PiNN `PiNet2` at charge-response, and polarizability models are downstream `molpot` modules. Out of scope: TensorFlow checkpoint import, PiNN YAML compatibility, BPNN, -legacy PiNet, PiNNAcLe, PiNNwall, and ASE calculator wrappers. +legacy PiNet, PiNNAcLe, and PiNNwall wrappers. ## 2. Paper↔Code Mapping @@ -40,6 +40,15 @@ MolNex adaptation: - PiNN's Keras `out_extra` readouts are not embedded in the encoder. MolNex writes raw `i1/i3/i5` representation tracks and downstream `molpot` heads perform task-specific projections. +- MolNex adds `emit_property_features` (no PiNN counterpart). The `i1/i3/i5` + tracks are per-*edge*, so at typical neighbour counts they dominate the + encoder's output volume while being read only by the property heads + (`PiNetDipole`, `PiNetPolarizability`) — never by the energy/force path. + The flag gates only the stacking and TensorDict write; the tracks are still + computed inside each block, so gated and ungated encoders produce + bit-identical `node_features`, `p1_block_outputs`, energies and forces. + Default `True` (PiNN-equivalent emission); `PiNetPotential` defaults the + encoders it constructs itself to `False`. - PiNN cutoff functions are zeroed outside `r_max` because MolNex edges may be supplied by arbitrary preprocessing tasks. - The encoder writes TensorDict keys in place rather than returning TensorFlow @@ -55,13 +64,16 @@ rank-5 basis: The encoder outputs per-block scalar states `(N, depth, D)`, optional vector states `(N, depth, 3, D)`, optional rank-5 states `(N, depth, 5, D)`, and -per-edge interaction tracks. +per-edge interaction tracks. The scalar states (`node_features`, +`p1_block_outputs`) are always written; the vector, rank-5 and per-edge tracks +are written only when `emit_property_features` is set (see §4). ## 6. Config Mapping `PiNet2Spec` mirrors PiNN constructor arguments: `atom_types`, `r_max`, `cutoff_type`, `basis_type`, `n_basis`, `gamma`, `center`, `pp_nodes`, -`pi_nodes`, `ii_nodes`, `depth`, `activation`, `weighted`, and `rank`. +`pi_nodes`, `ii_nodes`, `depth`, `activation`, `weighted`, and `rank`, plus the +MolNex-only `emit_property_features` (§4). ## 7. Benchmark Contract diff --git a/tests/conftest.py b/tests/conftest.py index e00d8a1..211807b 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -40,6 +40,7 @@ def make_graph_batch( batch: torch.Tensor, *, graphs: dict[str, torch.Tensor] | None = None, + shifts: torch.Tensor | None = None, ) -> TensorDict: """Build a ``TensorDict`` from raw tensors. @@ -51,12 +52,19 @@ def make_graph_batch( graphs: Optional per-graph fields (e.g. ``{"total_charge": tensor}``) written to the ``"graphs"`` sub-tensordict. ``num_atoms`` is auto-derived from ``batch`` and always present. + shifts: Optional PBC shift vectors ``(E, 3)`` (``unit_shifts @ cell``), + written to ``edges.shifts`` and folded additively into + ``edge_diff``. ``S_ij`` is constant w.r.t. ``pos``, so the + position gradient of the edge geometry is unchanged. Returns: - A fully-formed ``TensorDict`` with ``edge_diff = pos[dst] - pos[src]`` - and ``edge_dist = ‖edge_diff‖`` recomputed from ``pos``. + A fully-formed ``TensorDict`` with + ``edge_diff = pos[dst] - pos[src] (+ shifts)`` and + ``edge_dist = ‖edge_diff‖`` recomputed from ``pos``. """ edge_diff = pos[edge_index[:, 1]] - pos[edge_index[:, 0]] + if shifts is not None: + edge_diff = edge_diff + shifts edge_dist = edge_diff.norm(dim=-1).clamp(min=1e-6) n_atoms = pos.shape[0] n_edges = edge_index.shape[0] @@ -69,21 +77,29 @@ def make_graph_batch( for k, v in graphs.items(): graph_data[k] = v + edge_data = TensorDict( + edge_index=edge_index, + edge_diff=edge_diff, + edge_dist=edge_dist, + batch_size=[n_edges], + ) + if shifts is not None: + edge_data["shifts"] = shifts + return TensorDict( atoms=TensorDict(Z=Z, pos=pos, batch=batch, batch_size=[n_atoms]), - edges=TensorDict( - edge_index=edge_index, - edge_diff=edge_diff, - edge_dist=edge_dist, - batch_size=[n_edges], - ), + edges=edge_data, graphs=graph_data, batch_size=[], ) def translate_graph(batch: TensorDict, t: torch.Tensor) -> TensorDict: - """Shift all atomic positions by ``t``. Edge geometry recomputes from pos.""" + """Shift all atomic positions by ``t``. Edge geometry recomputes from pos. + + ``edges.shifts`` (if present) is carried unchanged: a rigid translation + does not change the box, so ``S = n · h`` is untouched. + """ pos = batch["atoms", "pos"] + t extras = _extract_graph_extras(batch) return make_graph_batch( @@ -92,24 +108,37 @@ def translate_graph(batch: TensorDict, t: torch.Tensor) -> TensorDict: edge_index=batch["edges", "edge_index"], batch=batch["atoms", "batch"], graphs=extras, + shifts=_extract_shifts(batch), ) def rotate_graph(batch: TensorDict, R: torch.Tensor) -> TensorDict: - """Rotate all atomic positions by ``R`` (3×3 rotation matrix).""" + """Rotate all atomic positions by ``R`` (3×3 rotation matrix). + + ``edges.shifts`` (if present) rotates with the positions: ``S = n · h`` + is a lattice vector, so rotating the box rotates ``S`` by the same ``R`` + while the integer image ``n`` stays put. + """ pos = batch["atoms", "pos"] @ R.T extras = _extract_graph_extras(batch) + shifts = _extract_shifts(batch) return make_graph_batch( pos=pos, Z=batch["atoms", "Z"], edge_index=batch["edges", "edge_index"], batch=batch["atoms", "batch"], graphs=extras, + shifts=None if shifts is None else shifts @ R.T, ) def permute_graph(batch: TensorDict, perm: torch.Tensor) -> TensorDict: - """Relabel atoms by ``perm``; edge_index is remapped, per-graph fields kept.""" + """Relabel atoms by ``perm``; edge_index is remapped, per-graph fields kept. + + ``edges.shifts`` is per-edge and the edge *rows* keep their order (only the + node labels inside ``edge_index`` are rewritten), so shifts stay aligned + and are carried unchanged. + """ inv_perm = torch.empty_like(perm) inv_perm[perm] = torch.arange(len(perm)) pos = batch["atoms", "pos"][perm] @@ -123,6 +152,7 @@ def permute_graph(batch: TensorDict, perm: torch.Tensor) -> TensorDict: edge_index=edge_index, batch=batch_idx, graphs=extras, + shifts=_extract_shifts(batch), ) @@ -131,10 +161,15 @@ def recompute_edge_geometry(batch: TensorDict) -> TensorDict: Call this at the start of a pipeline forward when ``pos`` carries ``requires_grad`` so autograd can trace ``∂E/∂pos`` for force tests. + ``edges.shifts`` (if present) is folded back in, exactly as + :func:`make_graph_batch` does. """ pos = batch["atoms", "pos"] edge_index = batch["edges", "edge_index"] edge_diff = pos[edge_index[:, 1]] - pos[edge_index[:, 0]] + shifts = _extract_shifts(batch) + if shifts is not None: + edge_diff = edge_diff + shifts edge_dist = edge_diff.norm(dim=-1).clamp(min=1e-6) batch["edges", "edge_diff"] = edge_diff batch["edges", "edge_dist"] = edge_dist @@ -158,3 +193,10 @@ def _extract_graph_extras(batch: TensorDict) -> dict[str, torch.Tensor] | None: graphs = batch["graphs"] extras = {k: graphs[k] for k in graphs.keys() if k not in _RESERVED_GRAPH_KEYS} return extras or None + + +def _extract_shifts(batch: TensorDict) -> torch.Tensor | None: + """Return ``edges.shifts`` if the batch carries PBC shift vectors.""" + if ("edges", "shifts") not in batch.keys(include_nested=True): + return None + return batch["edges", "shifts"] diff --git a/tests/regression/test_ff_export_goldens.py b/tests/regression/test_ff_export_goldens.py new file mode 100644 index 0000000..67e28d3 --- /dev/null +++ b/tests/regression/test_ff_export_goldens.py @@ -0,0 +1,49 @@ +"""Hard-coded force-export unit goldens (no live OpenMM). + +Spec: learnable-classical-ff-08-ff-export. + +Auto-marked ``regression`` by this directory's conftest; not part of the +default unit run. Public-API scenario script lives at +``regressions/learnable-classical-ff-08-ff-export.py``. +""" + +from __future__ import annotations + +import math + +import torch + +from molix.ff_export import ( + ForceFieldCompiler, + scale_amber_vn, + scale_bond_k, +) +from molpot.ir import BondBag, PotentialIR, ProperTorsionBag + + +def test_bond_k_golden_100_to_41840(): + assert scale_bond_k(100.0) == 41840.0 + + +def test_torsion_vn_golden_2_to_4_184(): + assert scale_amber_vn(2.0) == 4.184 + + +def test_compiler_emits_bond_and_torsion_goldens(): + ir = PotentialIR( + bonds=BondBag( + k=torch.tensor([100.0]), + r0=torch.tensor([1.0]), + ), + propers=ProperTorsionBag( + k=torch.tensor([[1.0]]), # half-barrier for AMBER Vn=2 + periodicity=torch.tensor([1], dtype=torch.long), + phase=torch.tensor([[0.0]]), + idivf=torch.tensor([1.0]), + ), + ) + spec = ForceFieldCompiler("openmm").compile(ir) + bond = next(f for f in spec.forces if f["type"] == "HarmonicBondForce") + tor = next(f for f in spec.forces if f["type"] == "PeriodicTorsionForce") + assert bond["parameters"][0]["k"] == 41840.0 + assert math.isclose(tor["parameters"][0]["k"], 4.184, abs_tol=0.0) diff --git a/tests/regression/test_timer.py b/tests/regression/test_timer.py index 65e2eb4..e60bfd7 100644 --- a/tests/regression/test_timer.py +++ b/tests/regression/test_timer.py @@ -1,4 +1,3 @@ -import ase import torch from molpot.potentials.elec import ( @@ -10,9 +9,27 @@ DTYPE = torch.float32 DEFAULT_CUTOFF = 4.4 -CHARGES_1 = torch.ones((4, 1), dtype=DTYPE) -POSITIONS_1 = 0.3 * torch.arange(12, dtype=DTYPE).reshape((4, 3)) -CELL_1 = torch.eye(3, dtype=DTYPE) + + +def _supercell( + pos: torch.Tensor, + charges: torch.Tensor, + cell: torch.Tensor, + reps: tuple[int, int, int], +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Tile a unit cell ``reps`` times along each lattice vector (numpy-free).""" + nx, ny, nz = reps + images = [] + for ix in range(nx): + for iy in range(ny): + for iz in range(nz): + shift = ix * cell[0] + iy * cell[1] + iz * cell[2] + images.append(pos + shift) + pos_sc = torch.cat(images, dim=0) + charges_sc = charges.repeat(nx * ny * nz, 1) + scale = torch.tensor([nx, ny, nz], dtype=cell.dtype, device=cell.device) + cell_sc = cell * scale.unsqueeze(1) + return pos_sc, charges_sc, cell_sc def test_timer(): @@ -20,14 +37,11 @@ def test_timer(): n_repeat_2 = 100 pos, charges, cell, _, _ = define_crystal() - # use ase to make system bigger - atoms = ase.Atoms("H" * len(pos), positions=pos.numpy(), cell=cell.numpy()) - atoms.set_initial_charges(charges.numpy().flatten()) - atoms.repeat((4, 4, 4)) - - pos = torch.tensor(atoms.positions, dtype=DTYPE) - charges = torch.tensor(atoms.get_initial_charges(), dtype=DTYPE).reshape(-1, 1) - cell = torch.tensor(atoms.cell.array, dtype=DTYPE) + # Enlarge the crystal without ASE: 4×4×4 supercell of the primitive cell. + pos = pos.to(DTYPE) + charges = charges.to(DTYPE).reshape(-1, 1) + cell = cell.to(DTYPE) + pos, charges, cell = _supercell(pos, charges, cell, (4, 4, 4)) neighbor_indices, neighbor_distances = neighbor_list( positions=pos, box=cell, cutoff=DEFAULT_CUTOFF diff --git a/tests/regression/test_values_dipole.py b/tests/regression/test_values_dipole.py index 6839b4f..cffd1e4 100644 --- a/tests/regression/test_values_dipole.py +++ b/tests/regression/test_values_dipole.py @@ -1,7 +1,7 @@ import pytest import torch -from ase.io import read +from molix.datasets._extxyz import parse_extxyz_frames from molpot.potentials.elec import CalculatorDipole, PotentialDipole from molpot.potentials.elec.prefactors import eV_A from tests.regression.conftest import ( @@ -12,11 +12,11 @@ neighbor_list, ) -frames = read(DIPOLES_TEST_FRAMES, ":3") +frames = parse_extxyz_frames(DIPOLES_TEST_FRAMES)[:3] cutoffs = [3.9986718930, 4.0000000000, 4.7363281250] alphas = [0.8819831493, 0.8956299559, 0.7215211182] -energies = [frame.get_potential_energy() for frame in frames] -forces = [frame.get_forces() for frame in frames] +energies = [frame.energy for frame in frames] +forces = [frame.forces for frame in frames] @pytest.mark.parametrize("device", DEVICES) @@ -116,9 +116,9 @@ def test_magnetostatic_ewald_crystal(self, frame, cutoff, alpha, energy, force, lr_wavelength=0.1, ) calc.to(device=device, dtype=dtype) - positions = torch.tensor(frame.get_positions(), dtype=dtype, device=device) - dipoles = torch.tensor(frame.get_array("dipoles"), dtype=dtype, device=device) - cell = torch.tensor(frame.get_cell().array, dtype=dtype, device=device) + positions = torch.tensor(frame.pos, dtype=dtype, device=device) + dipoles = torch.tensor(frame.arrays["dipoles"], dtype=dtype, device=device) + cell = torch.tensor(frame.cell, dtype=dtype, device=device) neighbor_indices, neighbor_shifts = neighbor_list( positions=positions, periodic=True, diff --git a/tests/regression/test_values_ewald.py b/tests/regression/test_values_ewald.py index 6addf33..1ff2521 100644 --- a/tests/regression/test_values_ewald.py +++ b/tests/regression/test_values_ewald.py @@ -3,9 +3,9 @@ import numpy as np import pytest import torch -from ase.io import read import molpot.potentials.elec +from molix.datasets._extxyz import parse_extxyz_frames from molpot.potentials.elec import ( CoulombPotential, EwaldCalculator, @@ -236,11 +236,11 @@ def test_random_structure( pme_order = 8 rcoulomb = 0.3 ; nm """ - frame = read(COULOMB_TEST_FRAMES, frame_index) + frame = parse_extxyz_frames(COULOMB_TEST_FRAMES)[frame_index] - positions = scaling_factor * torch.tensor(frame.positions, dtype=DTYPE) @ ortho - cell = scaling_factor * torch.tensor(frame.cell.array, dtype=DTYPE) @ ortho - charges = torch.tensor(frame.get_initial_charges(), dtype=DTYPE).reshape((-1, 1)) + positions = scaling_factor * torch.tensor(frame.pos, dtype=DTYPE) @ ortho + cell = scaling_factor * torch.tensor(frame.cell, dtype=DTYPE) @ ortho + charges = torch.tensor(frame.arrays["initial_charges"], dtype=DTYPE).reshape((-1, 1)) cutoff *= scaling_factor smearing = cutoff / 6.0 @@ -296,12 +296,13 @@ def test_random_structure( # Compute energy energy = torch.sum(potentials * charges) - energy_target = torch.tensor(frame.get_potential_energy(), dtype=DTYPE) / scaling_factor + energy_target = torch.tensor(frame.energy, dtype=DTYPE) / scaling_factor torch.testing.assert_close(energy, energy_target, atol=0.0, rtol=1e-4) # Compute forces forces = torch.autograd.grad(-energy, positions)[0] - forces_target = torch.tensor(frame.get_forces(), dtype=DTYPE) / scaling_factor**2 + assert frame.forces is not None + forces_target = torch.tensor(frame.forces, dtype=DTYPE) / scaling_factor**2 torch.testing.assert_close(forces, forces_target @ ortho, atol=0.0, rtol=5e-3) # Compute stress diff --git a/tests/test_molix/test_cli/test_check.py b/tests/test_molix/test_cli/test_check.py index 7c39852..78cafb7 100644 --- a/tests/test_molix/test_cli/test_check.py +++ b/tests/test_molix/test_cli/test_check.py @@ -44,7 +44,7 @@ def test_worst_status_ordering() -> None: def test_check_results_are_immutable() -> None: result = CheckResult("x", Status.OK, "detail") with pytest.raises((AttributeError, TypeError)): - result.detail = "mutated" # type: ignore[misc] + result.detail = "mutated" # -- the real environment (this venv must be healthy to run the suite) ----- diff --git a/tests/test_molix/test_core/test_hook_module.py b/tests/test_molix/test_core/test_hook_module.py index d9acf0d..b0289c9 100644 --- a/tests/test_molix/test_core/test_hook_module.py +++ b/tests/test_molix/test_core/test_hook_module.py @@ -105,11 +105,11 @@ def test_module_does_not_import_concrete_layers() -> None: def test_trainer_imports_hook_from_core_hook_module() -> None: - """ac-008 (partial): Trainer pulls the ``Hook`` Protocol from ``molix.core.hook``. + """Trainer pulls the ``Hook`` Protocol from ``molix.core.hook``. - The full ac-008 grep — that ``trainer.py`` has *zero* imports from - ``molix.core.hooks`` — fires after cycle 3 (concrete-hook split). - Here we only assert the Hook contract has moved to its new home. + The concrete-hook split is long done (``molix.core.hooks`` no longer + exists anywhere under ``src/``); this asserts the Hook contract lives + at its canonical home. """ import molix.core.trainer as trainer_mod diff --git a/tests/test_molix/test_core/test_losses_molecular.py b/tests/test_molix/test_core/test_losses_molecular.py index 1503254..4e97b84 100644 --- a/tests/test_molix/test_core/test_losses_molecular.py +++ b/tests/test_molix/test_core/test_losses_molecular.py @@ -3,8 +3,14 @@ from __future__ import annotations import torch +from tensordict import TensorDict from molix.core.losses import energy_force_mse, energy_mse +from molix.core.losses.molecular import ( + center_by_group, + molecule_centered_energy_mse, + parameter_bag_mse, +) from molix.data.collate import collate_molecules @@ -112,3 +118,40 @@ def test_lambda_zero_matches_energy_only(self): )(preds, batch) energy_only = energy_mse("U0")(preds, batch) assert torch.allclose(joint, energy_only) + + +class TestCenterByGroup: + def test_exact_centering(self): + values = torch.tensor([1.0, 3.0, 10.0, 14.0]) + groups = torch.tensor([0, 0, 1, 1]) + c = center_by_group(values, groups) + assert torch.allclose(c, torch.tensor([-1.0, 1.0, -2.0, 2.0])) + + +class TestMoleculeCenteredEnergyMse: + def test_zero_when_identical_after_center(self): + pred = torch.tensor([1.0, 2.0, 5.0, 7.0]) + true = torch.tensor([10.0, 11.0, 0.0, 2.0]) + groups = torch.tensor([0, 0, 1, 1]) + batch = TensorDict( + { + "graphs": TensorDict( + { + "mm_energy": true, + "molecule_id_index": groups, + }, + batch_size=[4], + ) + }, + batch_size=[], + ) + loss = molecule_centered_energy_mse()({"energy": pred}, batch) + assert float(loss) < 1e-10 + + +class TestParameterBagMse: + def test_zero_and_positive(self): + a = {"k": torch.tensor([1.0, 2.0])} + assert float(parameter_bag_mse(a, a)) == 0.0 + b = {"k": torch.tensor([1.0, 4.0])} + assert float(parameter_bag_mse(a, b)) == ((0.0 + 4.0) / 2.0) diff --git a/tests/test_molix/test_core/test_metrics.py b/tests/test_molix/test_core/test_metrics.py index 81f328d..4293ff7 100644 --- a/tests/test_molix/test_core/test_metrics.py +++ b/tests/test_molix/test_core/test_metrics.py @@ -93,3 +93,27 @@ def test_collection_gpu_inputs(self): metrics.update(preds, targets) results = metrics.compute() assert float(results["MAE"]) == pytest.approx(0.5) + + +class TestMoleculeCenteredRMSE: + def test_self_zero(self): + from molix.core.metrics import MoleculeCenteredRMSE + + m = MoleculeCenteredRMSE() + # values {0,0.25,1.0} same for pred/true → centered RMSE 0 + e = torch.tensor([0.0, 0.25, 1.0]) + g = torch.tensor([0, 0, 0]) + m.update(e, e, g) + assert float(m.compute()) == pytest.approx(0.0) + + +class TestMoleculeCenteredMAE: + def test_offsets(self): + from molix.core.metrics import MoleculeCenteredMAE + + m = MoleculeCenteredMAE() + pred = torch.tensor([1.0, 2.0, 5.0, 7.0]) + true = torch.tensor([10.0, 11.0, 0.0, 2.0]) + g = torch.tensor([0, 0, 1, 1]) + m.update(pred, true, g) + assert float(m.compute()) == pytest.approx(0.0) diff --git a/tests/test_molix/test_data/conftest.py b/tests/test_molix/test_data/conftest.py new file mode 100644 index 0000000..cb8e89d --- /dev/null +++ b/tests/test_molix/test_data/conftest.py @@ -0,0 +1,52 @@ +"""Shared fixtures for the ``molix.data`` unit tests. + +Currently holds one builder: the degenerate ``n_edges == n_atoms`` +geometry used by ``test_cache.py``, ``test_collate_packed.py`` and +``test_dataset.py`` to pin the packed-cache bucket-identity contract. +The three modules assert three consequences of the *same* contract, so +the geometry lives here rather than being copied (and drifting) three +times. +""" + +from __future__ import annotations + +import torch + + +def equal_count_samples(n: int = 3) -> list[dict]: + """Build *n* samples with ``n_edges == n_atoms == 2`` in every sample. + + A 2-atom molecule with a single bidirectional pair (``NeighborList``'s + default ``symmetry=True`` → ``E = 2 x n_pairs``) is the smallest + physically real case where the per-sample edge count collides with the + atom count. Nothing in the leading dim then distinguishes a per-atom + key from a per-edge key, so the packed schema must decide bucket + membership by key role, not by numerology. + + Every value is hand-built (no RNG), so the fixture is bit-identical + across runs and independent of global torch seed state. + + Args: + n: Number of samples. + + Returns: + Flat sample dicts with per-atom ``Z`` ``(2,)`` / ``pos`` ``(2, 3)``, + per-edge ``edge_index`` ``(2, 2)`` / ``edge_diff`` ``(2, 3)`` / + ``edge_dist`` ``(2,)``, and graph-level ``targets.U0`` ``(1,)`` + carrying the sample's identity ``float(i)``. + """ + samples: list[dict] = [] + for i in range(n): + # Atom 1 sits at +1 A along x from atom 0; the pair is stored in both + # directions, so edge_diff = pos[target] - pos[source] flips sign. + samples.append( + { + "Z": torch.tensor([1, 6], dtype=torch.long), + "pos": torch.tensor([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]]), + "edge_index": torch.tensor([[0, 1], [1, 0]], dtype=torch.long), + "edge_diff": torch.tensor([[1.0, 0.0, 0.0], [-1.0, 0.0, 0.0]]), + "edge_dist": torch.tensor([1.0, 1.0]), + "targets": {"U0": torch.tensor([float(i)])}, + } + ) + return samples diff --git a/tests/test_molix/test_data/test_cache.py b/tests/test_molix/test_data/test_cache.py index 7934252..9a84319 100644 --- a/tests/test_molix/test_data/test_cache.py +++ b/tests/test_molix/test_data/test_cache.py @@ -2,10 +2,14 @@ from __future__ import annotations +from pathlib import Path +from typing import Any + import pytest import torch from molix.data.cache import PackedCache +from tests.test_molix.test_data.conftest import equal_count_samples # --------------------------------------------------------------------------- # Helpers @@ -109,6 +113,81 @@ def test_scalar_metadata_roundtrips(self, tmp_path): assert loaded["samples"][0]["n_atoms"] == 1 +# --------------------------------------------------------------------------- +# Schema inference — bucket identity +# --------------------------------------------------------------------------- + + +class TestPackSchema: + """Edge keys keep their bucket even when ``n_edges == n_atoms``. + + ``_infer_schema_across`` classifies a key by its leading dim. For a + 2-atom molecule with one bidirectional pair (``E = 2 = N``) that dim + is ambiguous, and the tie must not be broken by numerology: absorbing + ``edge_index`` / ``edge_diff`` / ``edge_dist`` into the atoms bucket + also suppresses ``edge_ptr``, which is what every downstream edge + consumer reads (``MmapDataset.avg_num_neighbors``, + ``collate_packed``). The cache is the origin of that silent drop, so + it is pinned here at the payload level. + """ + + def _payload(self, tmp_path: Path, n: int = 3) -> dict[str, Any]: + """Save *n* equal-count samples and read the raw on-disk payload.""" + sink = tmp_path / "eq.pt" + PackedCache(sink).save(equal_count_samples(n)) + return torch.load(sink, weights_only=True) + + def test_edge_index_schema_kind_is_edge(self, tmp_path): + """Schema entry for ``edge_index`` is tagged ``"edge"``, not ``"atom"``.""" + payload = self._payload(tmp_path) + assert payload["schema"]["edge_index"][0] == "edge" + + def test_edge_geometry_schema_kind_is_edge(self, tmp_path): + """``edge_diff`` / ``edge_dist`` follow ``edge_index`` into the edge kind.""" + payload = self._payload(tmp_path) + assert payload["schema"]["edge_diff"][0] == "edge" + assert payload["schema"]["edge_dist"][0] == "edge" + + def test_edge_keys_packed_into_edges_bucket(self, tmp_path): + """The three edge keys are stored under ``payload["edges"]``.""" + payload = self._payload(tmp_path) + assert set(payload["edges"]) == {"edge_index", "edge_diff", "edge_dist"} + + def test_edge_keys_absent_from_atoms_bucket(self, tmp_path): + """No edge key leaks into ``payload["atoms"]`` (which holds Z / pos only).""" + payload = self._payload(tmp_path) + assert "edge_index" not in payload["atoms"] + assert "edge_diff" not in payload["atoms"] + assert "edge_dist" not in payload["atoms"] + assert set(payload["atoms"]) == {"Z", "pos"} + + def test_edge_ptr_is_written(self, tmp_path): + """``edge_ptr`` exists — its absence is what silences edge consumers.""" + payload = self._payload(tmp_path) + assert "edge_ptr" in payload + + def test_edge_ptr_is_per_sample_cumsum(self, tmp_path): + """``edge_ptr`` cumsums 2 edges per sample over 3 samples.""" + payload = self._payload(tmp_path) + assert payload["edge_ptr"].tolist() == [0, 2, 4, 6] + + def test_atom_ptr_unaffected(self, tmp_path): + """``atom_ptr`` still cumsums 2 atoms per sample (same numbers, other axis).""" + payload = self._payload(tmp_path) + assert payload["atom_ptr"].tolist() == [0, 2, 4, 6] + + def test_unpack_roundtrips_edge_keys(self, tmp_path): + """Per-sample unpack returns the hand-built edge tensors unchanged.""" + sink = tmp_path / "eq.pt" + samples = equal_count_samples(3) + PackedCache(sink).save(samples) + payload = PackedCache(sink).load() + s2 = PackedCache.unpack_sample(payload, 2) + assert torch.equal(s2["edge_index"], samples[2]["edge_index"]) + assert torch.equal(s2["edge_diff"], samples[2]["edge_diff"]) + assert torch.equal(s2["edge_dist"], samples[2]["edge_dist"]) + + # --------------------------------------------------------------------------- # unpack_sample (staticmethod) # --------------------------------------------------------------------------- diff --git a/tests/test_molix/test_data/test_cache_valence.py b/tests/test_molix/test_data/test_cache_valence.py new file mode 100644 index 0000000..2ac45d6 --- /dev/null +++ b/tests/test_molix/test_data/test_cache_valence.py @@ -0,0 +1,155 @@ +"""PackedCache + collate_packed round-trip for valence column buckets. + +Spec: learnable-classical-ff-02-valence-topology (ac-003). +""" + +from __future__ import annotations + +import pytest +import torch + +from molix.data.cache import PackedCache +from molix.data.collate import collate_molecules, collate_packed +from molix.data.dataset import MmapDataset + + +def _mol(n_atoms: int, **families: dict) -> dict: + sample: dict = { + "Z": torch.ones(n_atoms, dtype=torch.long), + "pos": torch.zeros(n_atoms, 3), + } + for name, cols in families.items(): + sample[name] = {k: torch.as_tensor(v, dtype=torch.long) for k, v in cols.items()} + return sample + + +def test_packed_cache_angles_roundtrip_and_collate_packed(tmp_path): + samples = [ + _mol( + 3, + angles={"atomi": [0], "atomj": [1], "atomk": [2], "type": [0]}, + ), + _mol( + 4, + angles={"atomi": [0, 1], "atomj": [1, 2], "atomk": [2, 3], "type": [1, 0]}, + ), + ] + sink = tmp_path / "valence.pt" + PackedCache(sink).save(samples) + ds = MmapDataset(sink) + + # unpack restores nested columns (local indices) + u1 = ds[1] + assert isinstance(u1["angles"], dict) + assert torch.equal(u1["angles"]["atomi"], samples[1]["angles"]["atomi"]) + assert torch.equal(u1["angles"]["atomj"], samples[1]["angles"]["atomj"]) + assert torch.equal(u1["angles"]["atomk"], samples[1]["angles"]["atomk"]) + assert torch.equal(u1["angles"]["type"], samples[1]["angles"]["type"]) + + # single-file layout still (one .pt sink) + assert sink.is_file() + assert not any(p.is_dir() for p in tmp_path.iterdir() if p.name.startswith("valence")) + + indices = [0, 1] + fast = collate_packed(ds.packed_view(), indices) + oracle = collate_molecules([ds[i] for i in indices]) + for col in ("atomi", "atomj", "atomk", "type"): + assert torch.equal(fast["angles"][col], oracle["angles"][col]) + # rebased: m2 +3 + assert fast["angles"]["atomi"].tolist() == [0, 3, 4] + assert list(fast["angles"].batch_size) == list(oracle["angles"].batch_size) + + +def test_packed_cache_propers_impropers_roundtrip(tmp_path): + samples = [ + _mol( + 4, + propers={ + "atomi": [0], + "atomj": [1], + "atomk": [2], + "atoml": [3], + "type": [0], + }, + impropers={ + "atomi": [1], + "atomj": [0], + "atomk": [2], + "atoml": [3], + }, + ), + _mol( + 5, + propers={ + "atomi": [0], + "atomj": [1], + "atomk": [2], + "atoml": [3], + "type": [1], + }, + impropers={ + "atomi": [2], + "atomj": [0], + "atomk": [1], + "atoml": [4], + }, + ), + ] + sink = tmp_path / "torsions.pt" + PackedCache(sink).save(samples) + ds = MmapDataset(sink) + + assert torch.equal(ds[0]["impropers"]["atomi"], torch.tensor([1])) + assert torch.equal(ds[1]["propers"]["atoml"], torch.tensor([3])) + + fast = collate_packed(ds.packed_view(), [0, 1]) + oracle = collate_molecules([ds[0], ds[1]]) + for fam in ("propers", "impropers"): + for col in oracle[fam].keys(): + assert torch.equal(fast[fam][col], oracle[fam][col]), f"{fam}.{col}" + # improper centers rebased: 1, 2+4 + assert fast["impropers"]["atomi"].tolist() == [1, 6] + + +def test_packed_cache_mixed_valence_all_or_none(tmp_path): + samples = [ + _mol(3, angles={"atomi": [0], "atomj": [1], "atomk": [2]}), + _mol(2), # missing angles + ] + with pytest.raises(ValueError, match="angles"): + PackedCache(tmp_path / "bad.pt").save(samples) + + +def test_bonds_and_angles_together_in_cache(tmp_path): + """Bonds COO path coexists with angle columns in one packed file.""" + samples = [ + { + "Z": torch.ones(3, dtype=torch.long), + "pos": torch.zeros(3, 3), + "bond_index": torch.tensor([[0, 1], [1, 2]], dtype=torch.long), + "bond_types": torch.tensor([0, 1], dtype=torch.long), + "angles": { + "atomi": torch.tensor([0], dtype=torch.long), + "atomj": torch.tensor([1], dtype=torch.long), + "atomk": torch.tensor([2], dtype=torch.long), + }, + }, + { + "Z": torch.ones(2, dtype=torch.long), + "pos": torch.zeros(2, 3), + "bond_index": torch.tensor([[0], [1]], dtype=torch.long), + "bond_types": torch.tensor([0], dtype=torch.long), + "angles": { + "atomi": torch.tensor([0], dtype=torch.long), + "atomj": torch.tensor([0], dtype=torch.long), + "atomk": torch.tensor([1], dtype=torch.long), + }, + }, + ] + sink = tmp_path / "both.pt" + PackedCache(sink).save(samples) + ds = MmapDataset(sink) + fast = collate_packed(ds.packed_view(), [0, 1]) + oracle = collate_molecules([ds[0], ds[1]]) + assert torch.equal(fast["bonds", "bond_index"], oracle["bonds", "bond_index"]) + assert torch.equal(fast["angles"]["atomi"], oracle["angles"]["atomi"]) diff --git a/tests/test_molix/test_data/test_collate_packed.py b/tests/test_molix/test_data/test_collate_packed.py index e866244..8e65416 100644 --- a/tests/test_molix/test_data/test_collate_packed.py +++ b/tests/test_molix/test_data/test_collate_packed.py @@ -25,6 +25,11 @@ ``(1,)`` (never 1 atom/edge), so the packed schema routes ``U0`` to graph, ``forces`` to atom, and ``n_heavy`` to scalar with no leading-dim ambiguity. + +``TestEqualAtomAndEdgeCounts`` deliberately breaks that "never equal" +property (2-atom molecules, one bidirectional pair → ``E == N``) to pin +the case where leading-dim classification is ambiguous: the fast path +must still see the edges the oracle sees. """ from __future__ import annotations @@ -39,6 +44,7 @@ from molix.data.cache import PackedCache from molix.data.collate import TargetSchema, collate_molecules from molix.data.dataset import CachedDataset, MmapDataset, SubsetDataset +from tests.test_molix.test_data.conftest import equal_count_samples # --------------------------------------------------------------------------- # Fixture sample builders @@ -189,6 +195,34 @@ def _assert_td_equal(got: TensorDict, want: TensorDict) -> None: assert torch.equal(g, w), f"{level}.{key} values differ" +def _assert_edges_equal(got: TensorDict, want: TensorDict) -> None: + """Assert the ``edges`` namespace of two batches is leaf-for-leaf identical. + + Narrower than :func:`_assert_td_equal` so a failure names the edge + namespace directly rather than whichever level the walk reaches first. + + Args: + got: Output under test (fast path). + want: Oracle output (slow path). + """ + e_got = got["edges"] + e_want = want["edges"] + assert list(e_got.batch_size) == list(e_want.batch_size), ( + f"edges batch_size: {list(e_got.batch_size)} != {list(e_want.batch_size)}" + ) + assert set(e_got.keys()) == set(e_want.keys()), ( + f"edges keys: {set(e_got.keys())} != {set(e_want.keys())}" + ) + for key in e_want.keys(): + g = e_got[key] + w = e_want[key] + assert g.dtype == w.dtype, f"edges.{key} dtype: {g.dtype} != {w.dtype}" + assert tuple(g.shape) == tuple(w.shape), ( + f"edges.{key} shape: {tuple(g.shape)} != {tuple(w.shape)}" + ) + assert torch.equal(g, w), f"edges.{key} values differ" + + def _packed_collate(dataset, indices: list[int], schema: TargetSchema) -> TensorDict: """Run the fast path: ``collate_packed(dataset.packed_view(), indices, schema)``. @@ -258,6 +292,45 @@ def test_all_edgeless_fallback_exact_fields(self, tmp_path): assert torch.equal(edges["edge_dist"], torch.zeros(0)) +class TestEqualAtomAndEdgeCounts: + """``E == N`` batches keep their edges on the fast path (ac-001). + + Two-atom molecules with one bidirectional pair make every per-edge + tensor's leading dim indistinguishable from a per-atom tensor's. The + oracle (per-sample dicts → ``collate_molecules``) is unaffected, so + any divergence here is the packed path dropping edges. + """ + + #: U0 only — these samples carry no atom-level target. + SCHEMA_EQ = TargetSchema(graph_level=frozenset({"U0"}), atom_level=frozenset()) + + def test_edges_namespace_matches_oracle(self, tmp_path): + """Every ``edges`` leaf equals the oracle's for an ``E == N`` cache.""" + ds = MmapDataset(_save(tmp_path, "eqcount", equal_count_samples(3))) + indices = [0, 2] + got = _packed_collate(ds, indices, self.SCHEMA_EQ) + want = _oracle(ds, indices, self.SCHEMA_EQ) + _assert_edges_equal(got, want) + + def test_edge_index_shape_is_not_empty(self, tmp_path): + """2 samples x 2 edges → ``edge_index`` ``(4, 2)`` on both paths.""" + ds = MmapDataset(_save(tmp_path, "eqcount_shape", equal_count_samples(3))) + indices = [0, 2] + got = _packed_collate(ds, indices, self.SCHEMA_EQ) + want = _oracle(ds, indices, self.SCHEMA_EQ) + assert tuple(want["edges", "edge_index"].shape) == (4, 2) + assert tuple(got["edges", "edge_index"].shape) == (4, 2) + + def test_full_batch_matches_oracle(self, tmp_path): + """The whole nested batch — not just edges — equals the oracle.""" + ds = MmapDataset(_save(tmp_path, "eqcount_full", equal_count_samples(3))) + indices = [0, 1] + _assert_td_equal( + _packed_collate(ds, indices, self.SCHEMA_EQ), + _oracle(ds, indices, self.SCHEMA_EQ), + ) + + class TestTargetSchemaRouting: """atom_level / graph_level / scalar targets route exactly as the oracle (ac-003).""" diff --git a/tests/test_molix/test_data/test_datamodule.py b/tests/test_molix/test_data/test_datamodule.py index 9d43d63..f0733ab 100644 --- a/tests/test_molix/test_data/test_datamodule.py +++ b/tests/test_molix/test_data/test_datamodule.py @@ -17,18 +17,54 @@ # --------------------------------------------------------------------------- +#: Atom-pair separations (A) alternated by :func:`_make_samples`. The +#: pipeline tests below build neighbor lists at ``cutoff=3.0``, so even +#: samples are in-cutoff (one bidirectional pair, 2 edges) and odd samples +#: are out-of-cutoff (0 edges). The jitter never spans the cutoff. +_IN_CUTOFF_SEP = 1.5 +_OUT_OF_CUTOFF_SEP = 5.0 +_JITTER = 0.01 + + def _make_samples(n: int = 10) -> list[dict]: - return [ - { - "Z": torch.tensor([1, 6], dtype=torch.long), - "pos": torch.randn(2, 3), - "edge_index": torch.tensor([[0, 1]], dtype=torch.long), - "edge_diff": torch.randn(1, 3), - "edge_dist": torch.tensor([1.5]), - "targets": {"U0": torch.tensor([float(i)])}, - } - for i in range(n) - ] + """Build *n* two-atom samples with deterministic geometry. + + Positions come from a local :class:`torch.Generator` (seed 0), so the + fixture is independent of global RNG state and of test execution + order. Separations alternate across the ``3.0`` cutoff used by the + pipeline tests, which keeps the per-sample edge count *varying* (2 and + 0) instead of landing on a constant ``E == n_atoms`` — the degenerate + geometry pinned in ``test_cache.py::TestPackSchema``. + + Args: + n: Number of samples. ``targets.U0`` carries the identity + ``float(i)`` used by the shuffle-order tests. + + Returns: + Flat sample dicts with per-atom ``Z`` / ``pos``, one pre-built + edge (``edge_index`` ``(1, 2)``, ``edge_diff`` ``(1, 3)``, + ``edge_dist`` ``(1,)``) and graph-level ``targets.U0`` ``(1,)``. + """ + gen = torch.Generator().manual_seed(0) + samples: list[dict] = [] + for i in range(n): + sep = _IN_CUTOFF_SEP if i % 2 == 0 else _OUT_OF_CUTOFF_SEP + pos = torch.tensor([[0.0, 0.0, 0.0], [sep, 0.0, 0.0]]) + pos = pos + _JITTER * torch.randn(2, 3, generator=gen) + # edge_index [[0, 1]] → source 0, target 1; edge_diff follows the + # repo convention pos[target] - pos[source]. + diff = (pos[1] - pos[0]).unsqueeze(0) + samples.append( + { + "Z": torch.tensor([1, 6], dtype=torch.long), + "pos": pos, + "edge_index": torch.tensor([[0, 1]], dtype=torch.long), + "edge_diff": diff, + "edge_dist": diff.norm(dim=-1), + "targets": {"U0": torch.tensor([float(i)])}, + } + ) + return samples def _write_and_load_split(tmp_path: Path, train_n: int, val_n: int, dataset_cls=CachedDataset): diff --git a/tests/test_molix/test_data/test_dataset.py b/tests/test_molix/test_data/test_dataset.py index 47d6187..8ff772d 100644 --- a/tests/test_molix/test_data/test_dataset.py +++ b/tests/test_molix/test_data/test_dataset.py @@ -20,6 +20,7 @@ MmapDataset, SubsetDataset, ) +from tests.test_molix.test_data.conftest import equal_count_samples # --------------------------------------------------------------------------- # Helpers @@ -120,6 +121,30 @@ def test_pickle_roundtrip(self, tmp_path): assert torch.equal(ds2[i]["Z"], samples[i]["Z"]) +# --------------------------------------------------------------------------- +# Connectivity statistics +# --------------------------------------------------------------------------- + + +class TestAvgNumNeighbors: + """``avg_num_neighbors`` = total_edges / total_atoms from the cache pointers.""" + + def test_equal_atom_and_edge_counts(self, tmp_path): + """3 x (2 atoms, 2 edges) → exactly 1.0 neighbour per atom. + + ``E == N`` per sample is the degenerate geometry where per-edge + keys can be misfiled as per-atom ones; the ``edge_ptr`` that this + property reads then never gets written and the property falls back + to ``0.0`` — a silently wrong normalisation constant for + Allegro/MACE rather than a loud failure. Exact equality (not a + tolerance): both totals are integers read off cumsum pointers. + """ + sink = tmp_path / "eq.pt" + PackedCache(sink).save(equal_count_samples(3)) + ds = MmapDataset(sink) + assert ds.avg_num_neighbors == 1.0 + + # --------------------------------------------------------------------------- # CachedDataset # --------------------------------------------------------------------------- diff --git a/tests/test_molix/test_data/test_group_split.py b/tests/test_molix/test_data/test_group_split.py new file mode 100644 index 0000000..d4eec9c --- /dev/null +++ b/tests/test_molix/test_data/test_group_split.py @@ -0,0 +1,17 @@ +"""Tests for group_split_indices.""" + +from __future__ import annotations + +from molix.data.group_split import group_split_indices + + +class TestGroupSplitIndices: + def test_no_leakage_801010(self): + # 10 mols × 2 confs + ids = [f"m{m}" for m in range(10) for _ in range(2)] + train, val, test = group_split_indices(ids, ratios=(0.8, 0.1, 0.1), seed=0) + assert len({ids[i] for i in train}) == 8 + assert len({ids[i] for i in val}) == 1 + assert len({ids[i] for i in test}) == 1 + assert set(ids[i] for i in train).isdisjoint(ids[i] for i in val) + assert set(ids[i] for i in train).isdisjoint(ids[i] for i in test) diff --git a/tests/test_molix/test_data/test_pipeline.py b/tests/test_molix/test_data/test_pipeline.py index 5e53f9c..64acc49 100644 --- a/tests/test_molix/test_data/test_pipeline.py +++ b/tests/test_molix/test_data/test_pipeline.py @@ -112,7 +112,7 @@ def test_add_prebuilt_node(self): def test_rejects_non_task_non_callable(self): with pytest.raises(TypeError, match="Task"): - Pipeline("p").add(42) # type: ignore[arg-type] + Pipeline("p").add(42) def test_rejects_duplicate_name_on_add(self): p = Pipeline("p").add(CountingSample(), name="shared") @@ -223,7 +223,7 @@ class TestNode: def test_is_frozen(self): n = Node(name="x", task=CountingSample()) with pytest.raises(Exception): - n.name = "y" # type: ignore[misc] + n.name = "y" def test_apply_dispatches_runnable(self): t = CountingSample() @@ -238,7 +238,7 @@ def test_apply_dispatches_bare_callable(self): assert out["tag"] is True def test_apply_rejects_non_callable(self): - node = Node(name="bad", task=42) # type: ignore[arg-type] + node = Node(name="bad", task=42) with pytest.raises(TypeError): node.apply({}) diff --git a/tests/test_molix/test_data/test_tasks/__init__.py b/tests/test_molix/test_data/test_tasks/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/test_molix/test_data/test_tasks/test_pad.py b/tests/test_molix/test_data/test_tasks/test_pad.py index c828432..b3f3508 100644 --- a/tests/test_molix/test_data/test_tasks/test_pad.py +++ b/tests/test_molix/test_data/test_tasks/test_pad.py @@ -10,9 +10,9 @@ import pytest import torch -from tests.conftest import make_graph_batch from molix.data import PadMolecularBatch +from tests.conftest import make_graph_batch def _batch() -> tuple[object, int, int]: diff --git a/tests/test_molix/test_data/test_valence_collate.py b/tests/test_molix/test_data/test_valence_collate.py new file mode 100644 index 0000000..686d1db --- /dev/null +++ b/tests/test_molix/test_data/test_valence_collate.py @@ -0,0 +1,212 @@ +"""Valence topology collate: angles / propers / impropers as column TensorDicts. + +Spec: learnable-classical-ff-02-valence-topology. + +Pre-collate samples carry nested column dicts under ``angles`` / ``propers`` / +``impropers`` (atomi/atomj/atomk[/atoml], optional type). Post-collate the same +namespaces are nested TensorDicts with 1-D long columns rebased by atom offset. +Packed COO ``angle_index [3, N]`` is intentionally NOT the collate schema. +""" + +from __future__ import annotations + +import pytest +import torch +from tensordict import TensorDict + +from molix.data.collate import collate_molecules + + +def _mol( + n_atoms: int, + *, + angles: dict | None = None, + propers: dict | None = None, + impropers: dict | None = None, + bond_index: torch.Tensor | None = None, + bond_types: torch.Tensor | None = None, +) -> dict: + sample: dict = { + "Z": torch.ones(n_atoms, dtype=torch.long), + "pos": torch.zeros(n_atoms, 3), + } + if angles is not None: + sample["angles"] = { + k: torch.as_tensor(v, dtype=torch.long) for k, v in angles.items() + } + if propers is not None: + sample["propers"] = { + k: torch.as_tensor(v, dtype=torch.long) for k, v in propers.items() + } + if impropers is not None: + sample["impropers"] = { + k: torch.as_tensor(v, dtype=torch.long) for k, v in impropers.items() + } + if bond_index is not None: + sample["bond_index"] = bond_index.long() + if bond_types is not None: + sample["bond_types"] = bond_types.long() + return sample + + +def test_angles_column_rebase_two_molecules(): + """Second molecule's angle indices equal local + n_atoms_0 (ac-001).""" + m1 = _mol( + 3, + angles={ + "atomi": [0], + "atomj": [1], # central + "atomk": [2], + "type": [0], + }, + ) + m2 = _mol( + 4, + angles={ + "atomi": [0, 1], + "atomj": [1, 2], + "atomk": [2, 3], + "type": [1, 0], + }, + ) + batch = collate_molecules([m1, m2]) + + angles = batch["angles"] + assert isinstance(angles, TensorDict) + # Nested access form preferred by the contract. + atomi = batch["angles"]["atomi"] + atomj = batch["angles"]["atomj"] + atomk = batch["angles"]["atomk"] + assert atomi.dtype == torch.long + assert atomi.shape == (3,) # 1 + 2 angles + # m1 local, m2 rebased by +3 atoms. + assert atomi.tolist() == [0, 3, 4] + assert atomj.tolist() == [1, 4, 5] + assert atomk.tolist() == [2, 5, 6] + assert batch["angles"]["type"].tolist() == [0, 1, 0] + # Prefer batch_size=[N] when all leaves share length N. + assert list(angles.batch_size) == [3] + # Primary schema is columns, not packed angle_index [3, N]. + assert "angle_index" not in angles.keys() + + +def test_propers_and_impropers_columns_with_center_first(): + """Propers/impropers expose atomi..atoml; impropers.atomi is center (ac-002).""" + m1 = _mol( + 4, + propers={ + "atomi": [0], + "atomj": [1], + "atomk": [2], + "atoml": [3], + "type": [0], + }, + impropers={ + "atomi": [1], # center (molrs) + "atomj": [0], + "atomk": [2], + "atoml": [3], + "type": [2], + }, + ) + m2 = _mol( + 5, + propers={ + "atomi": [0, 1], + "atomj": [1, 2], + "atomk": [2, 3], + "atoml": [3, 4], + "type": [1, 0], + }, + impropers={ + "atomi": [2], + "atomj": [0], + "atomk": [1], + "atoml": [3], + "type": [0], + }, + ) + batch = collate_molecules([m1, m2]) + + # m2 offset = 4 + assert batch["propers"]["atomi"].tolist() == [0, 4, 5] + assert batch["propers"]["atomj"].tolist() == [1, 5, 6] + assert batch["propers"]["atomk"].tolist() == [2, 6, 7] + assert batch["propers"]["atoml"].tolist() == [3, 7, 8] + assert batch["propers"]["type"].tolist() == [0, 1, 0] + assert list(batch["propers"].batch_size) == [3] + + assert batch["impropers"]["atomi"].tolist() == [1, 6] # centers: 1, 2+4 + assert batch["impropers"]["atomj"].tolist() == [0, 4] + assert batch["impropers"]["atomk"].tolist() == [2, 5] + assert batch["impropers"]["atoml"].tolist() == [3, 7] + assert batch["impropers"]["type"].tolist() == [2, 0] + assert list(batch["impropers"].batch_size) == [2] + + +def test_valence_all_or_none_raises(): + """Mixed presence of a valence family across samples raises ValueError.""" + m1 = _mol(3, angles={"atomi": [0], "atomj": [1], "atomk": [2]}) + m2 = _mol(2) # no angles + with pytest.raises(ValueError, match="angles"): + collate_molecules([m1, m2]) + + +def test_empty_angles_namespace(): + """Zero-count angles on every sample still emit empty columns.""" + empty = { + "atomi": torch.zeros(0, dtype=torch.long), + "atomj": torch.zeros(0, dtype=torch.long), + "atomk": torch.zeros(0, dtype=torch.long), + } + m1 = _mol(2, angles=empty) + m2 = _mol(3, angles=empty) + batch = collate_molecules([m1, m2]) + assert batch["angles"]["atomi"].shape == (0,) + assert batch["angles"]["atomj"].shape == (0,) + assert batch["angles"]["atomk"].shape == (0,) + assert list(batch["angles"].batch_size) == [0] + + +def test_no_valence_omits_namespaces(): + """Samples without valence keys do not grow angles/propers/impropers.""" + batch = collate_molecules([_mol(2), _mol(3)]) + assert "angles" not in batch.keys() + assert "propers" not in batch.keys() + assert "impropers" not in batch.keys() + + +def test_bonds_still_work_alongside_angles(): + """Existing bond_index path stays green next to angles (ac-006).""" + # bond_index is COO [2, N]: m1 has bonds 0-1 and 1-2; m2 has bond 0-1. + m1 = _mol( + 3, + angles={"atomi": [0], "atomj": [1], "atomk": [2], "type": [0]}, + bond_index=torch.tensor([[0, 1], [1, 2]], dtype=torch.long), + bond_types=torch.tensor([0, 1], dtype=torch.long), + ) + m2 = _mol( + 2, + angles={"atomi": [0], "atomj": [0], "atomk": [1], "type": [1]}, + bond_index=torch.tensor([[0], [1]], dtype=torch.long), + bond_types=torch.tensor([0], dtype=torch.long), + ) + batch = collate_molecules([m1, m2]) + assert batch["bonds", "bond_index"].tolist() == [[0, 1, 3], [1, 2, 4]] + assert batch["bonds", "bond_types"].tolist() == [0, 1, 0] + assert batch["angles"]["atomi"].tolist() == [0, 3] + + +def test_collate_does_not_import_molpy_as_batch_store(): + """Topology leaves are plain torch tensors inside TensorDict (ac-004).""" + from pathlib import Path + + import molix.data.cache as cache_mod + import molix.data.collate as collate_mod + + for mod in (collate_mod, cache_mod): + src = Path(mod.__file__).read_text() + assert "from molpy" not in src + assert "import molpy" not in src + assert "ForceField" not in src + assert "import molrs" not in src diff --git a/tests/test_molix/test_datasets/conftest.py b/tests/test_molix/test_datasets/conftest.py index 8fcba45..7cefb2a 100644 --- a/tests/test_molix/test_datasets/conftest.py +++ b/tests/test_molix/test_datasets/conftest.py @@ -1,7 +1,7 @@ """Shared fixtures for ``WaterLESSource`` / ``ChargedDimersSource`` tests. Both fixtures write extended-XYZ files to a ``tmp_path`` directory in the -exact format ASE expects (``ase.io.read(..., format="extxyz", index=":")``): +exact extended-XYZ format consumed by ``parse_extxyz_frames``: a comment line carrying ``Lattice="..."``, ``Properties=...``, ``energy=...`` (eV) and ``pbc="T T T"``, then atom rows `` x y z fx fy fz``. @@ -381,6 +381,10 @@ def molrec_qm9_record(tmp_path: Path) -> MolRecQM9Fixture: frames = [_make_molpy_frame(el, pos) for el, pos in zip(elements, positions)] rec = MolRec() rec.set_trajectory(Trajectory.from_frames(frames)) + # molpy >= 0.12 requires a record to carry a frame / system / status + # alongside its trajectory: the trajectory is the time series, the system is + # the fixed topology it is a series *of*. + rec.set_system(frames[0]) # Distinct deterministic float32 values per target so the round-trip test # can detect any cross-target mixing. @@ -444,6 +448,10 @@ def molrec_force_record(tmp_path: Path) -> MolRecForceFixture: frames = [_make_molpy_frame(elements, pos) for pos in positions] rec = MolRec() rec.set_trajectory(Trajectory.from_frames(frames)) + # molpy >= 0.12 requires a record to carry a frame / system / status + # alongside its trajectory: the trajectory is the time series, the system is + # the fixed topology it is a series *of*. + rec.set_system(frames[0]) rec.observables.add_scalar( "teacherA.energy", diff --git a/tests/test_molix/test_datasets/test_extxyz.py b/tests/test_molix/test_datasets/test_extxyz.py index e0e60ef..1b0e4bb 100644 --- a/tests/test_molix/test_datasets/test_extxyz.py +++ b/tests/test_molix/test_datasets/test_extxyz.py @@ -109,23 +109,30 @@ def test_pbc_defaults_to_true_when_missing(self, tmp_path: Path) -> None: class TestNoASE: - def test_no_ase_imports_in_datasets(self) -> None: - """ac-006: ``src/molix/datasets/`` contains no ASE imports. + def test_no_ase_imports_in_src(self) -> None: + """Package code under ``src/`` must not import ASE (or e3nn). - Replaces ASE's ``ase.io.read(..., format="extxyz")`` with the in-tree - ``_extxyz.parse_extxyz_frames`` parser. Grep is the binding rule. + Extxyz I/O uses the in-tree ``_extxyz.parse_extxyz_frames`` parser. + Grep is the binding rule — no soft optional ASE shims either. """ repo_root = Path(__file__).resolve().parents[3] - datasets_dir = repo_root / "src" / "molix" / "datasets" - result = subprocess.run( - ["grep", "-rnE", r"^\s*(import ase|from ase)", str(datasets_dir)], - capture_output=True, - text=True, - check=False, - ) - assert result.returncode == 1, ( # grep returns 1 when no matches - f"ASE import found in src/molix/datasets/:\n{result.stdout}" - ) + src_dir = repo_root / "src" + for pattern in ( + r"^\s*(import ase|from ase)\b", + r"^\s*(import e3nn|from e3nn)\b", + ): + # --include=*.py: src/ also holds C++ sources and CMake build + # trees (src/molix/op/build*) — grepping those thousands of + # files made this test take ~9 s for a rule about Python imports. + result = subprocess.run( + ["grep", "-rnE", "--include=*.py", pattern, str(src_dir)], + capture_output=True, + text=True, + check=False, + ) + assert result.returncode == 1, ( # grep returns 1 when no matches + f"Forbidden import under src/ (pattern {pattern!r}):\n{result.stdout}" + ) class TestChargedDimersRemoved: @@ -150,6 +157,7 @@ def test_no_charged_dimers_grep_under_src(self) -> None: [ "grep", "-rnE", + "--include=*.py", r"ChargedDimersSource|charged_dimers\.py", str(src_dir), ], diff --git a/tests/test_molix/test_datasets/test_valence_columns.py b/tests/test_molix/test_datasets/test_valence_columns.py new file mode 100644 index 0000000..8175e0f --- /dev/null +++ b/tests/test_molix/test_datasets/test_valence_columns.py @@ -0,0 +1,93 @@ +"""Optional kernel-local stack helpers for valence columns → COO. + +Spec: learnable-classical-ff-02-valence-topology (ac-005). + +These helpers build ``[arity, N]`` long tensors only at potential call sites; +collate output stays column form under nested TensorDict namespaces. +""" + +from __future__ import annotations + +import pytest +import torch +from tensordict import TensorDict + +from molix.data.collate import collate_molecules +from molix.datasets._valence_columns import ( + stack_angle_index, + stack_improper_index, + stack_proper_index, +) + + +def test_stack_angle_index_shape_and_values(): + angles = TensorDict( + { + "atomi": torch.tensor([0, 1], dtype=torch.long), + "atomj": torch.tensor([1, 2], dtype=torch.long), + "atomk": torch.tensor([2, 3], dtype=torch.long), + "type": torch.tensor([0, 1], dtype=torch.long), + }, + batch_size=[2], + ) + idx = stack_angle_index(angles) + assert idx.shape == (3, 2) + assert idx.dtype == torch.long + assert idx.tolist() == [[0, 1], [1, 2], [2, 3]] + + +def test_stack_proper_and_improper_index(): + propers = TensorDict( + { + "atomi": torch.tensor([0], dtype=torch.long), + "atomj": torch.tensor([1], dtype=torch.long), + "atomk": torch.tensor([2], dtype=torch.long), + "atoml": torch.tensor([3], dtype=torch.long), + }, + batch_size=[1], + ) + impropers = TensorDict( + { + "atomi": torch.tensor([1], dtype=torch.long), # center + "atomj": torch.tensor([0], dtype=torch.long), + "atomk": torch.tensor([2], dtype=torch.long), + "atoml": torch.tensor([3], dtype=torch.long), + }, + batch_size=[1], + ) + p = stack_proper_index(propers) + i = stack_improper_index(impropers) + assert p.shape == (4, 1) + assert i.shape == (4, 1) + assert p.tolist() == [[0], [1], [2], [3]] + assert i.tolist() == [[1], [0], [2], [3]] # center-first preserved + + +def test_stack_from_collated_batch_not_required_by_schema(): + """Helper is opt-in after collate; batch itself has no angle_index leaf.""" + samples = [ + { + "Z": torch.ones(3, dtype=torch.long), + "pos": torch.zeros(3, 3), + "angles": { + "atomi": torch.tensor([0], dtype=torch.long), + "atomj": torch.tensor([1], dtype=torch.long), + "atomk": torch.tensor([2], dtype=torch.long), + }, + } + ] + batch = collate_molecules(samples) + assert "angle_index" not in batch["angles"].keys() + idx = stack_angle_index(batch["angles"]) + assert idx.tolist() == [[0], [1], [2]] + + +def test_stack_length_mismatch_raises(): + # Plain mapping (not TensorDict) so construction itself does not validate shapes. + angles = { + "atomi": torch.tensor([0, 1], dtype=torch.long), + "atomj": torch.tensor([1], dtype=torch.long), + "atomk": torch.tensor([2, 3], dtype=torch.long), + } + with pytest.raises(ValueError, match="share length"): + stack_angle_index(angles) diff --git a/tests/test_molix/test_datasets/test_water_les.py b/tests/test_molix/test_datasets/test_water_les.py index 0106aee..0ecfbb0 100644 --- a/tests/test_molix/test_datasets/test_water_les.py +++ b/tests/test_molix/test_datasets/test_water_les.py @@ -162,7 +162,7 @@ def test_target_schema(self): class TestWaterLESErrors: def test_unknown_split_rejected(self, water_les_root): with pytest.raises(ValueError, match="split"): - WaterLESSource(water_les_root, split="not_a_split") # type: ignore[arg-type] + WaterLESSource(water_les_root, split="not_a_split") def test_offline_without_files_raises(self, tmp_path): empty = tmp_path / "empty" diff --git a/tests/test_molix/test_engine.py b/tests/test_molix/test_engine.py index 4d9efde..9c4abd5 100644 --- a/tests/test_molix/test_engine.py +++ b/tests/test_molix/test_engine.py @@ -28,13 +28,16 @@ class _DummyTDPotential(nn.Module): """Mimics a molnex potential: nested-TensorDict in, ``{energy, forces}`` out.""" - def forward(self, batch, *, compute_forces: bool = False): + def forward(self, batch): + # Monomorphic forward, mirroring a real potential since b85d12f: a + # potential that derives forces was built that way and always writes them. pos = batch["atoms", "pos"] bd = batch["edges", "edge_dist"] - out = {"energy": (bd**2).sum().reshape(1), "atomic_energy": torch.zeros(pos.shape[0])} - if compute_forces: - out["forces"] = torch.zeros_like(pos) - return out + return { + "energy": (bd**2).sum().reshape(1), + "atomic_energy": torch.zeros(pos.shape[0]), + "forces": torch.zeros_like(pos), + } class _FlatEFPotential(nn.Module): diff --git a/tests/test_molix/test_export.py b/tests/test_molix/test_export.py index d2ad096..c68fd4c 100644 --- a/tests/test_molix/test_export.py +++ b/tests/test_molix/test_export.py @@ -141,9 +141,9 @@ def test_direct_export_on_force_model_raises_no_fallback() -> None: def test_non_module_raises_typeerror() -> None: with pytest.raises(TypeError): - Exporter("not_a_module") # type: ignore[arg-type] + Exporter("not_a_module") def test_non_tuple_inputs_raises_typeerror(tmp_path: Path, small_mlp: nn.Sequential) -> None: with pytest.raises(TypeError): - Exporter(small_mlp).export(torch.randn(4, 10), tmp_path / "m.pt2") # type: ignore[arg-type] + Exporter(small_mlp).export(torch.randn(4, 10), tmp_path / "m.pt2") diff --git a/tests/test_molix/test_ff_export/__init__.py b/tests/test_molix/test_ff_export/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/test_molix/test_ff_export/test_compiler.py b/tests/test_molix/test_ff_export/test_compiler.py new file mode 100644 index 0000000..1b80d11 --- /dev/null +++ b/tests/test_molix/test_ff_export/test_compiler.py @@ -0,0 +1,63 @@ +"""ForceFieldCompiler + registry integration tests.""" + +from __future__ import annotations + +import pytest +import torch + +from molix.ff_export import ForceFieldCompiler, OpenMMAdapter, UnsupportedTermError +from molix.ff_export.adapter import BackendAdapter +from molix.ff_export.cases import TranslationCase +from molix.ff_export.force_spec import ForceSpec +from molpot.ir import BondBag, ImproperHarmonicBag, PotentialIR + + +class TestForceFieldCompiler: + def test_default_adapter_is_openmm_by_name(self): + compiler = ForceFieldCompiler() + assert isinstance(compiler.adapter, OpenMMAdapter) + + def test_adapter_instance_accepted(self): + compiler = ForceFieldCompiler(adapter=OpenMMAdapter()) + ir = PotentialIR(bonds=BondBag(k=torch.tensor([100.0]), r0=torch.tensor([1.0]))) + spec = compiler.compile(ir) + assert isinstance(spec, ForceSpec) + assert spec.backend == "openmm" + bond = next(f for f in spec.forces if f["type"] == "HarmonicBondForce") + assert bond["parameters"][0]["k"] == 41840.0 + + def test_unknown_adapter_name_raises(self): + with pytest.raises(ValueError, match="unknown adapter"): + ForceFieldCompiler(adapter="gromacs") + + def test_unsupported_term_raises_structured_error(self): + compiler = ForceFieldCompiler("openmm") + ir = PotentialIR( + impropers_harmonic=ImproperHarmonicBag( + k=torch.tensor([1.0]), + chi0=torch.tensor([0.0]), + ) + ) + with pytest.raises(UnsupportedTermError) as ei: + compiler.compile(ir) + err = ei.value + assert err.term == "improper_harmonic" + assert err.case is TranslationCase.UNSUPPORTED + assert "improper_harmonic" in str(err) + + def test_empty_ir_compiles_to_empty_forces(self): + spec = ForceFieldCompiler().compile(PotentialIR()) + assert spec.forces == [] + assert spec.backend == "openmm" + + def test_meta_passthrough(self): + meta = {"type_systems": {"bond": ["C-C"]}, "symbolic": None} + spec = ForceFieldCompiler().compile( + PotentialIR(), + type_systems=meta["type_systems"], + symbolic=meta["symbolic"], + ) + assert spec.metadata.get("type_systems") == {"bond": ["C-C"]} + + def test_registry_lists_openmm(self): + assert "openmm" in BackendAdapter.names() diff --git a/tests/test_molix/test_ff_export/test_conventions.py b/tests/test_molix/test_ff_export/test_conventions.py new file mode 100644 index 0000000..55a5cbd --- /dev/null +++ b/tests/test_molix/test_ff_export/test_conventions.py @@ -0,0 +1,117 @@ +"""Unit tests for TranslationCase + ConventionTable unit goldens. + +Spec: learnable-classical-ff-08-ff-export. +Hard-coded goldens only — no live OpenMM. +""" + +from __future__ import annotations + +import math + +import pytest + +from molix.ff_export.cases import TranslationCase +from molix.ff_export.conventions import ( + ANGSTROM_TO_NM, + BOND_K_IR_TO_OPENMM, + KCAL_PER_MOL_TO_KJ_PER_MOL, + ConventionTable, + scale_amber_vn, + scale_angle_k, + scale_bond_k, + scale_energy, + scale_length, + scale_torsion_k, +) + + +class TestTranslationCase: + def test_four_way_members_exist(self): + names = {c.name for c in TranslationCase} + assert names == { + "DIRECT_UNIT_SCALE", + "FORM_REPARAMETERIZE", + "DECOMPOSE", + "UNSUPPORTED", + } + + def test_values_are_stable_strings(self): + assert TranslationCase.DIRECT_UNIT_SCALE.value == "direct_unit_scale" + assert TranslationCase.FORM_REPARAMETERIZE.value == "form_reparameterize" + assert TranslationCase.DECOMPOSE.value == "decompose" + assert TranslationCase.UNSUPPORTED.value == "unsupported" + + +class TestUnitConstants: + def test_kcal_to_kj(self): + assert KCAL_PER_MOL_TO_KJ_PER_MOL == 4.184 + + def test_angstrom_to_nm(self): + assert ANGSTROM_TO_NM == 0.1 + + def test_bond_k_factor_is_418_4(self): + # k_omm = k_ir * 4.184 / (0.1 nm)^2 = k_ir * 418.4 + assert BOND_K_IR_TO_OPENMM == pytest.approx(418.4) + + +class TestScaleHelpers: + def test_bond_k_golden_100_to_41840(self): + """ac-001: k = 100 kcal mol⁻¹ Å⁻² → 41840 kJ mol⁻¹ nm⁻².""" + assert scale_bond_k(100.0) == 41840.0 + + def test_bond_k_zero(self): + assert scale_bond_k(0.0) == 0.0 + + def test_torsion_k_unit_scale(self): + """IR E=(k/s)[1+cos] prefactor in kcal/mol → OpenMM k in kJ/mol.""" + assert scale_torsion_k(1.0) == 4.184 + + def test_amber_vn_golden_2_to_4_184(self): + """ac-002: AMBER Vn=2 kcal/mol → OpenMM PeriodicTorsion k=4.184 kJ/mol. + + AMBER: E = (Vn/2)[1 + cos(...)]; OpenMM: E = k[1 + cos(...)]. + With Vn=2, k_kcal = 1 and k_omm = 4.184. + """ + assert scale_amber_vn(2.0) == 4.184 + # Same path as IR coefficient after half-barrier reparameterization. + assert scale_torsion_k(2.0 / 2.0) == 4.184 + + def test_energy_and_length_scales(self): + assert scale_energy(1.0) == 4.184 + assert scale_length(10.0) == 1.0 # 10 Å = 1 nm + + def test_angle_k_energy_only(self): + # θ in rad both sides; only energy unit changes. + assert scale_angle_k(10.0) == pytest.approx(41.84) + + +class TestConventionTable: + def test_default_table_has_core_terms(self): + table = ConventionTable.default_openmm() + for term in ( + "bond_harmonic", + "angle_harmonic", + "proper_periodic", + "lj", + "charge", + "improper_harmonic", + ): + assert term in table + + def test_bond_is_direct_unit_scale(self): + row = ConventionTable.default_openmm()["bond_harmonic"] + assert row.case is TranslationCase.DIRECT_UNIT_SCALE + assert math.isclose(row.unit_factors["k"], BOND_K_IR_TO_OPENMM) + + def test_proper_is_form_reparameterize(self): + row = ConventionTable.default_openmm()["proper_periodic"] + assert row.case is TranslationCase.FORM_REPARAMETERIZE + + def test_improper_harmonic_unsupported(self): + row = ConventionTable.default_openmm()["improper_harmonic"] + assert row.case is TranslationCase.UNSUPPORTED + + def test_lookup_unknown_raises(self): + table = ConventionTable.default_openmm() + with pytest.raises(KeyError): + table.row("not_a_term") diff --git a/tests/test_molix/test_ff_export/test_force_spec.py b/tests/test_molix/test_ff_export/test_force_spec.py new file mode 100644 index 0000000..9fd9488 --- /dev/null +++ b/tests/test_molix/test_ff_export/test_force_spec.py @@ -0,0 +1,61 @@ +"""ForceSpec schema + JSON round-trip tests.""" + +from __future__ import annotations + +import json + +from molix.ff_export.force_spec import ForceSpec + + +class TestForceSpec: + def test_construct_and_to_dict(self): + spec = ForceSpec( + backend="openmm", + forces=[ + { + "type": "HarmonicBondForce", + "case": "direct_unit_scale", + "parameters": [{"type_index": 0, "k": 41840.0, "r0": 0.15}], + } + ], + scaling={"scale_q_14": 5.0 / 6.0, "scale_lj_14": 0.5}, + metadata={"unit_system_source": "class_i_canonical"}, + ) + d = spec.to_dict() + assert d["backend"] == "openmm" + assert d["forces"][0]["type"] == "HarmonicBondForce" + assert d["scaling"]["scale_lj_14"] == 0.5 + + def test_to_dict_is_json_serializable(self): + spec = ForceSpec( + backend="openmm", + forces=[{"type": "HarmonicAngleForce", "parameters": []}], + scaling=None, + metadata={}, + ) + payload = json.dumps(spec.to_dict()) + assert "HarmonicAngleForce" in payload + + def test_from_dict_round_trip(self): + original = ForceSpec( + backend="openmm", + forces=[ + { + "type": "PeriodicTorsionForce", + "case": "form_reparameterize", + "parameters": [ + { + "type_index": 0, + "term_index": 0, + "periodicity": 2, + "phase": 0.0, + "k": 4.184, + } + ], + } + ], + scaling={"scale_q_14": 5.0 / 6.0, "scale_lj_14": 0.5}, + metadata={"note": "roundtrip"}, + ) + restored = ForceSpec.from_dict(original.to_dict()) + assert restored == original diff --git a/tests/test_molix/test_ff_export/test_openmm_adapter.py b/tests/test_molix/test_ff_export/test_openmm_adapter.py new file mode 100644 index 0000000..d6680a1 --- /dev/null +++ b/tests/test_molix/test_ff_export/test_openmm_adapter.py @@ -0,0 +1,148 @@ +"""OpenMMAdapter force-spec emission tests (no live openmm).""" + +from __future__ import annotations + +import json +import math + +import pytest +import torch + +from molix.ff_export.adapter import BackendAdapter +from molix.ff_export.cases import TranslationCase +from molix.ff_export.exceptions import UnsupportedTermError +from molix.ff_export.openmm_adapter import OpenMMAdapter +from molpot.ir import ( + AngleBag, + BondBag, + ChargeBag, + ImproperHarmonicBag, + LJBag, + NonbondedScaling, + PotentialIR, + ProperTorsionBag, +) + + +def _class_i_ir() -> PotentialIR: + return PotentialIR( + bonds=BondBag( + k=torch.tensor([100.0]), + r0=torch.tensor([1.5]), # Å + ), + angles=AngleBag( + k=torch.tensor([50.0]), + theta0=torch.tensor([1.9106332362490186]), # ~109.47° + ), + propers=ProperTorsionBag( + # AMBER Vn=2 → IR half-barrier k=1 with idivf=1 + k=torch.tensor([[1.0]]), + periodicity=torch.tensor([2], dtype=torch.long), + phase=torch.tensor([[0.0]]), + idivf=torch.tensor([1.0]), + ), + lj=LJBag( + epsilon=torch.tensor([0.1]), # kcal/mol + sigma=torch.tensor([3.5]), # Å + ), + charges=ChargeBag(q=torch.tensor([0.5, -0.5])), + scaling=NonbondedScaling(), + ) + + +class TestOpenMMAdapterRegistration: + def test_registered_as_openmm(self): + assert "openmm" in BackendAdapter.names() + adapter = BackendAdapter.from_name("openmm") + assert isinstance(adapter, OpenMMAdapter) + + def test_unknown_backend_raises(self): + with pytest.raises(ValueError, match="unknown adapter"): + BackendAdapter.from_name("not-a-backend") + + +class TestOpenMMAdapterTranslate: + def test_bond_k_and_r0_units(self): + adapter = OpenMMAdapter() + spec = adapter.translate(_class_i_ir()) + bond_force = next(f for f in spec.forces if f["type"] == "HarmonicBondForce") + p0 = bond_force["parameters"][0] + assert p0["k"] == 41840.0 + assert p0["r0"] == pytest.approx(0.15) # 1.5 Å → 0.15 nm + assert bond_force["case"] == TranslationCase.DIRECT_UNIT_SCALE.value + + def test_torsion_vn2_golden_via_ir_k(self): + """IR k=1 (half of AMBER Vn=2) → OpenMM k=4.184 kJ/mol.""" + adapter = OpenMMAdapter() + spec = adapter.translate(_class_i_ir()) + torsion = next(f for f in spec.forces if f["type"] == "PeriodicTorsionForce") + p0 = torsion["parameters"][0] + assert p0["k"] == 4.184 + assert p0["periodicity"] == 2 + assert p0["phase"] == 0.0 + + def test_idivf_absorbed_into_k(self): + ir = PotentialIR( + propers=ProperTorsionBag( + k=torch.tensor([[2.0]]), + periodicity=torch.tensor([1], dtype=torch.long), + phase=torch.tensor([[0.0]]), + idivf=torch.tensor([2.0]), + ) + ) + spec = OpenMMAdapter().translate(ir) + torsion = next(f for f in spec.forces if f["type"] == "PeriodicTorsionForce") + # (k/idivf)*4.184 = 1*4.184 + assert torsion["parameters"][0]["k"] == 4.184 + assert torsion["case"] == TranslationCase.FORM_REPARAMETERIZE.value + + def test_multi_term_proper_decomposes(self): + ir = PotentialIR( + propers=ProperTorsionBag( + k=torch.tensor([[1.0, 0.5]]), + periodicity=torch.tensor([1, 2], dtype=torch.long), + phase=torch.tensor([[0.0, math.pi]]), + idivf=torch.tensor([1.0]), + ) + ) + spec = OpenMMAdapter().translate(ir) + torsion = next(f for f in spec.forces if f["type"] == "PeriodicTorsionForce") + assert len(torsion["parameters"]) == 2 + assert torsion["case"] == TranslationCase.DECOMPOSE.value + assert torsion["parameters"][0]["k"] == 4.184 + assert torsion["parameters"][1]["k"] == pytest.approx(2.092) + assert torsion["parameters"][1]["periodicity"] == 2 + + def test_nonbonded_units_and_charges(self): + spec = OpenMMAdapter().translate(_class_i_ir()) + nb = next(f for f in spec.forces if f["type"] == "NonbondedForce") + assert nb["parameters"][0]["sigma"] == pytest.approx(0.35) # 3.5 Å + assert nb["parameters"][0]["epsilon"] == pytest.approx(0.4184) + assert nb["charges"][0] == 0.5 + assert nb["charges"][1] == -0.5 + + def test_scaling_14_from_ir(self): + """ac-006: Class-I defaults scale_q_14=5/6, scale_lj_14=0.5.""" + spec = OpenMMAdapter().translate(_class_i_ir()) + assert math.isclose(spec.scaling["scale_q_14"], 5.0 / 6.0) + assert math.isclose(spec.scaling["scale_lj_14"], 0.5) + assert math.isclose(spec.scaling["scale_q_12"], 0.0) + assert math.isclose(spec.scaling["scale_lj_12"], 0.0) + + def test_force_spec_json_serializable(self): + spec = OpenMMAdapter().translate(_class_i_ir()) + payload = json.dumps(spec.to_dict()) + assert "HarmonicBondForce" in payload + assert "PeriodicTorsionForce" in payload + + def test_unsupported_improper_harmonic_raises(self): + ir = PotentialIR( + impropers_harmonic=ImproperHarmonicBag( + k=torch.tensor([1.0]), + chi0=torch.tensor([0.0]), + ) + ) + with pytest.raises(UnsupportedTermError, match="improper_harmonic") as ei: + OpenMMAdapter().translate(ir) + assert ei.value.term == "improper_harmonic" + assert ei.value.case is TranslationCase.UNSUPPORTED diff --git a/tests/test_molix/test_md/__init__.py b/tests/test_molix/test_md/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/test_molix/test_md/conftest.py b/tests/test_molix/test_md/conftest.py new file mode 100644 index 0000000..fe23b11 --- /dev/null +++ b/tests/test_molix/test_md/conftest.py @@ -0,0 +1,76 @@ +"""Shared fixtures for the MD test package: a tiny PiNet system and a lattice. + +Imported by package path (``from tests.test_molix.test_md.conftest import …``) +per the repo test-layout rule — no free-floating helper modules. +""" + +import torch + +_DEVICE = torch.device("cpu") + + +def make_cubic_lattice(n_side: int = 3, spacing: float = 3.0) -> tuple[torch.Tensor, torch.Tensor]: + """A simple cubic lattice and its cell — a periodic system with real edges. + + The one owner of the MD suite's periodic fixture (``test_forcefield.py``, + ``test_neighbors.py`` and ``test_driver.py`` each held or wanted a + byte-identical copy). ``n_side=4, spacing=3.0`` is the 64-atom, 12 Å cube + the policy and driver suites run argon in: minimum perpendicular half-width + 6.0 Å, exact coordination shells at 3.0 / 4.2426 / 5.196 Å. + + Args: + n_side: Atoms per axis; the system holds ``n_side ** 3`` atoms. + spacing: Lattice constant in Angstrom. + + Returns: + ``(positions (n_side ** 3, 3), cell (3, 3))`` in Angstrom, ``float64``. + """ + grid = torch.arange(n_side, dtype=torch.float64) * spacing + pos = torch.stack(torch.meshgrid(grid, grid, grid, indexing="ij"), dim=-1).reshape(-1, 3) + cell = torch.eye(3, dtype=torch.float64) * (n_side * spacing) + return pos, cell + + +def make_tiny_potential(): + """A 4-species PiNet potential small enough for eager unit tests.""" + from molzoo.pinet import PiNetPotential + + torch.manual_seed(0) + return ( + PiNetPotential( + atom_types=[1, 6, 7, 8], + r_max=4.0, + n_basis=3, + pp_nodes=[8, 8], + pi_nodes=[8, 8], + ii_nodes=[8, 8], + depth=2, + rank=3, + hidden_dim=16, + # Force derivation is fixed at construction since b85d12f (the + # monomorphic pipelines torch.compile needs); an MD force field + # cannot ask for it per call. + compute_forces=True, + ) + .to(_DEVICE) + .eval() + ) + + +def make_pinet_template(): + """A 4-atom collated batch matching :func:`make_tiny_potential`.""" + from tests.conftest import make_graph_batch + + pos = torch.tensor( + [[0.0, 0.0, 0.0], [1.1, 0.1, 0.0], [0.3, 1.2, 0.2], [1.4, 1.1, -0.1]], + dtype=torch.float32, + device=_DEVICE, + ) + z = torch.tensor([1, 6, 7, 8], dtype=torch.long, device=_DEVICE) + edge_index = torch.tensor( + [[0, 1], [1, 0], [0, 2], [2, 0], [1, 3], [3, 1], [2, 3], [3, 2]], + dtype=torch.long, + device=_DEVICE, + ) + batch = torch.zeros(4, dtype=torch.long, device=_DEVICE) + return make_graph_batch(pos, z, edge_index, batch) diff --git a/tests/test_molix/test_md/oracle_bruteforce_neighbors.py b/tests/test_molix/test_md/oracle_bruteforce_neighbors.py new file mode 100644 index 0000000..412e907 --- /dev/null +++ b/tests/test_molix/test_md/oracle_bruteforce_neighbors.py @@ -0,0 +1,252 @@ +"""Independent multi-image neighbor-graph oracle for NeighborList tests. + +**Not** the MIC-only helper in ``test_neighbors._reference_pairs``. Enumerates +all lattice images with a range derived from perpendicular cell widths so +pairs that need |S| > 1 are found when the box is small relative to the cutoff. + +No imports from ``molix.md.neighbors``, ``molix.op``, matscipy, ASE, freud, or +LAMMPS. Production filter parity (``get_neighbor_pairs`` / binned path):: + + 0 < |dr| <= cutoff + +True self-edge ``i == j`` and ``S == (0,0,0)`` is excluded; periodic self-images +are kept when inside the cutoff. +""" + +from __future__ import annotations + +import math +from collections.abc import Iterable +from dataclasses import dataclass + +import torch + +# Edge key: central i, neighbor j, integer image shift of j's lattice copy. +EdgeKey = tuple[int, int, int, int, int] + + +def perpendicular_widths(cell: torch.Tensor) -> torch.Tensor: + """Perpendicular widths ``V / ||a_j × a_k||`` for each axis (Angstrom).""" + c = cell.to(dtype=torch.float64) + v = torch.det(c).abs() + cross = torch.stack( + ( + torch.linalg.cross(c[1], c[2]), + torch.linalg.cross(c[2], c[0]), + torch.linalg.cross(c[0], c[1]), + ), + dim=0, + ) + areas = cross.norm(dim=-1).clamp_min(1e-30) + return v / areas + + +def image_range(cell: torch.Tensor, cutoff: float) -> tuple[int, int, int]: + """Half-widths of the integer image stencil per axis (at least 1 if periodic).""" + widths = perpendicular_widths(cell) + r = float(cutoff) + out: list[int] = [] + for w in widths.tolist(): + if not math.isfinite(w) or w <= 0.0: + out.append(1) + continue + # Need n such that n * w can still reach distance ~ cutoff. + n = int(math.ceil(r / w + 1e-12)) + out.append(max(1, n)) + return out[0], out[1], out[2] + + +def bruteforce_edges( + pos: torch.Tensor, + *, + cell: torch.Tensor | None, + cutoff: float, + pbc: tuple[bool, bool, bool] = (True, True, True), +) -> set[EdgeKey]: + """Return the set of directed edges ``(i, j, sx, sy, sz)`` with ``0 < r <= cutoff``.""" + p = pos.detach().to(dtype=torch.float64) + n = int(p.shape[0]) + r_cut = float(cutoff) + r_cut_sq = r_cut * r_cut + keys: set[EdgeKey] = set() + + if cell is None or not any(pbc): + for i in range(n): + for j in range(n): + if i == j: + continue + d = p[j] - p[i] + d2 = float((d * d).sum()) + if 0.0 < d2 <= r_cut_sq: + keys.add((i, j, 0, 0, 0)) + return keys + + c = cell.detach().to(dtype=torch.float64) + rx, ry, rz = image_range(c, r_cut) + ranges = [] + for periodic, rmax in zip(pbc, (rx, ry, rz), strict=True): + ranges.append(range(-rmax, rmax + 1) if periodic else range(0, 1)) + + for i in range(n): + for j in range(n): + for sx in ranges[0]: + for sy in ranges[1]: + for sz in ranges[2]: + if i == j and sx == 0 and sy == 0 and sz == 0: + continue + shift_cart = float(sx) * c[0] + float(sy) * c[1] + float(sz) * c[2] + d = p[j] - p[i] + shift_cart + d2 = float((d * d).sum()) + if 0.0 < d2 <= r_cut_sq: + keys.add((i, j, int(sx), int(sy), int(sz))) + return keys + + +def shifts_to_integer_S( + shifts: torch.Tensor, + cell: torch.Tensor, + *, + atol: float = 1e-5, +) -> torch.Tensor: + """Map continuous Å remainders to integer lattice vectors ``(E, 3)``. + + Solves ``S @ cell ≈ shifts`` i.e. ``S ≈ shifts @ inv(cell)``, then rounds. + Raises if any residual exceeds *atol* (Å). + """ + s = shifts.detach().to(dtype=torch.float64) + c = cell.detach().to(dtype=torch.float64) + inv = torch.linalg.inv(c) + frac = s @ inv + S = torch.round(frac) + residual = (S @ c - s).norm(dim=-1) + bad = residual > atol + if bool(bad.any()): + idx = int(torch.nonzero(bad, as_tuple=False)[0].item()) + raise AssertionError( + f"shift→integer residual {float(residual[idx]):.3e} A exceeds atol={atol} " + f"at edge {idx}; shift={s[idx].tolist()} S={S[idx].tolist()} cell={c.tolist()}" + ) + return S.to(dtype=torch.long) + + +def neighborlist_edge_keys( + edge_index: torch.Tensor, + shifts: torch.Tensor, + num_edges: int, + cell: torch.Tensor | None, + *, + open_system: bool = False, +) -> set[EdgeKey]: + """Convert live NeighborList buffers to the oracle edge-key set.""" + n = int(num_edges) + if n == 0: + return set() + src = edge_index[:n, 0].detach().cpu().long() + tgt = edge_index[:n, 1].detach().cpu().long() + sh = shifts[:n].detach().cpu() + if open_system or cell is None: + return { + (int(src[k]), int(tgt[k]), 0, 0, 0) + for k in range(n) + if not (int(src[k]) == int(tgt[k]) and sh[k].abs().sum() < 1e-12) + } + S = shifts_to_integer_S(sh, cell) + return {(int(src[k]), int(tgt[k]), int(S[k, 0]), int(S[k, 1]), int(S[k, 2])) for k in range(n)} + + +@dataclass(frozen=True) +class GraphCompare: + """Diagnostics for one oracle vs SUT comparison.""" + + n_ref: int + n_sut: int + missing: frozenset[EdgeKey] + extra: frozenset[EdgeKey] + max_dr_mismatch: float + + @property + def ok(self) -> bool: + return not self.missing and not self.extra + + def summary(self) -> str: + return ( + f"n_ref={self.n_ref} n_sut={self.n_sut} " + f"missing={len(self.missing)} extra={len(self.extra)} " + f"max_dr_mismatch={self.max_dr_mismatch:.3e}" + ) + + +def compare_graphs( + ref: set[EdgeKey], + sut: set[EdgeKey], + *, + pos: torch.Tensor | None = None, + cell: torch.Tensor | None = None, +) -> GraphCompare: + """Set difference + optional max displacement mismatch on the intersection.""" + missing = frozenset(ref - sut) + extra = frozenset(sut - ref) + max_dr = 0.0 + if pos is not None and cell is not None and (ref & sut): + p = pos.detach().to(dtype=torch.float64) + c = cell.detach().to(dtype=torch.float64) + for i, j, sx, sy, sz in ref & sut: + shift = float(sx) * c[0] + float(sy) * c[1] + float(sz) * c[2] + dr = p[j] - p[i] + shift + # Both keys mean the same S; mismatch is numerical only. + max_dr = max(max_dr, float(dr.norm())) # not a mismatch — keep 0 + max_dr = 0.0 # keys match ⇒ same S by construction; residual checked elsewhere + return GraphCompare( + n_ref=len(ref), + n_sut=len(sut), + missing=missing, + extra=extra, + max_dr_mismatch=max_dr, + ) + + +def assert_graphs_equal( + ref: set[EdgeKey], + sut: set[EdgeKey], + *, + label: str = "", +) -> GraphCompare: + """Hard-fail with diagnostics if missing/extra edges exist.""" + cmp = compare_graphs(ref, sut) + if not cmp.ok: + head_m = list(sorted(cmp.missing))[:12] + head_e = list(sorted(cmp.extra))[:12] + raise AssertionError( + f"{label}neighbor graph mismatch ({cmp.summary()}); " + f"cutoff convention is 0 < r <= r_c. " + f"missing sample={head_m} extra sample={head_e}" + ) + return cmp + + +def physical_dr_multiset( + keys: Iterable[EdgeKey], + pos: torch.Tensor, + cell: torch.Tensor | None, + *, + decimals: int = 6, +) -> set[tuple[float, float, float]]: + """Multiset of rounded physical displacements for wrap-invariance checks.""" + p = pos.detach().to(dtype=torch.float64) + c = ( + cell.detach().to(dtype=torch.float64) + if cell is not None + else torch.eye(3, dtype=torch.float64) + ) + out: set[tuple[float, float, float]] = set() + for i, j, sx, sy, sz in keys: + shift = float(sx) * c[0] + float(sy) * c[1] + float(sz) * c[2] + dr = p[j] - p[i] + shift + out.add( + ( + round(float(dr[0]), decimals), + round(float(dr[1]), decimals), + round(float(dr[2]), decimals), + ) + ) + return out diff --git a/tests/test_molix/test_md_compile.py b/tests/test_molix/test_md/test_compile.py similarity index 51% rename from tests/test_molix/test_md_compile.py rename to tests/test_molix/test_md/test_compile.py index 452abf6..8814f83 100644 --- a/tests/test_molix/test_md_compile.py +++ b/tests/test_molix/test_md/test_compile.py @@ -1,42 +1,32 @@ -"""Compile + typed-contract tests for the MD component engine (spec ac-003/004/005). - -- ``MDState`` / ``ForceOutput`` are pytrees (flatten → unflatten round-trip). -- ``LennardJonesForceField`` closed-form force == ``-∂E/∂pos`` (autograd ref). -- ``Integrator.step`` / ``rollout`` ``torch.compile(fullgraph=True)`` with 0 graph - breaks — including over a real PiNet force field — and == eager where the - result is deterministic (step with fixed noise; NVE rollout where noise has - zero weight). +"""Compile tests for the MD component engine. + +``Integrator.step`` / ``rollout`` must ``torch.compile(fullgraph=True)`` with 0 +graph breaks — including over a real PiNet force field — and match eager where +the result is deterministic (step with fixed noise; NVE rollout where noise has +zero weight). The pytree/force-consistency contracts live in +``test_types.py`` / ``test_forcefield.py``. """ import pytest import torch -from torch.utils._pytree import tree_flatten, tree_unflatten from molix.md import ( - ForceOutput, HarmonicForceField, LangevinVerletIntegrator, + LennardJonesCutForceField, LennardJonesForceField, - MDState, + NeighborList, PotentialForceField, ) -from tests.test_molix.test_md_dynamics import _template, _tiny_potential +from tests.test_molix.test_md.conftest import ( + make_cubic_lattice, + make_pinet_template, + make_tiny_potential, +) _DTYPE = torch.float64 -def test_mdstate_forceoutput_are_pytrees(): - s = MDState(torch.randn(4, 3), torch.randn(4, 3), torch.randn(4, 3), torch.tensor(1.0)) - leaves, spec = tree_flatten(s) - assert len(leaves) == 4 - rebuilt = tree_unflatten(leaves, spec) - assert isinstance(rebuilt, MDState) - assert torch.equal(rebuilt.pos, s.pos) and torch.equal(rebuilt.energy, s.energy) - fo = ForceOutput(torch.tensor(1.0), torch.randn(4, 3)) - leaves2, spec2 = tree_flatten(fo) - assert isinstance(tree_unflatten(leaves2, spec2), ForceOutput) - - def test_lj_force_matches_autograd(): torch.manual_seed(0) pos = torch.randn(8, 3, dtype=_DTYPE, requires_grad=True) @@ -46,6 +36,66 @@ def test_lj_force_matches_autograd(): assert torch.allclose(out.forces, -ref, atol=1e-8), "LJ closed-form force != -dE/dx" +def _ljcut_ig( + *, rebuild: bool, skin: float = 0.0 +) -> tuple[LangevinVerletIntegrator, NeighborList, torch.Tensor]: + """Argon lj/cut over the 27-atom lattice, with the rebuild switch explicit. + + ``rebuild`` is passed on purpose in both arms: the switch defaults to the + force field's ``rebuilds_neighbors`` (``True`` here), and a live policy is + an eager, host-syncing decision that cannot live inside a traced step — + a compiled-path test must therefore freeze it (``rebuild=False``). + """ + pos, cell = make_cubic_lattice(n_side=3, spacing=3.0) + torch.manual_seed(2) + pos = pos + 0.2 * torch.randn_like(pos) + nl = NeighborList(cell=cell, cutoff=3.5, positions=pos, skin=skin) + ff = LennardJonesCutForceField(epsilon=0.7, sigma=2.5, neighbors=nl).to(_DTYPE) + ig = LangevinVerletIntegrator( + ff, dt=0.5, gamma=0.0, kbt=0.0, mass=39.95, seed=2, rebuild=rebuild + ) + return ig, nl, pos + + +def test_ljcut_step_fullgraph_compiles_and_matches_eager(): + """lj/cut over the fixed-capacity list (index_add + cutoff mask) traces fullgraph.""" + ig, _, pos = _ljcut_ig(rebuild=False) + st = ig.initial(pos, torch.zeros_like(pos)) + noise = torch.zeros_like(pos) + eager = ig.step(st, noise) + comp = torch.compile(ig.step, fullgraph=True, backend=_BACKEND)(st, noise) + assert torch.allclose(eager.pos, comp.pos, atol=1e-12) + assert torch.allclose(eager.forces, comp.forces, atol=1e-12) + assert torch.allclose(eager.energy, comp.energy, atol=1e-12) + + +def test_frozen_ljcut_rollout_compiles_fullgraph_and_leaves_the_list_alone(): + """``rebuild=False`` is the compiled-path invariant: dynamo specialises the + Python bool, the policy body stays dead, and ``rollout`` — not just ``step`` + — still traces to one graph and matches eager. A single rebuild inside the + loop would show up here as a graph break *and* as a nonzero count.""" + ig, nl, pos = _ljcut_ig(rebuild=False) + vel = torch.zeros_like(pos) + eager = ig.rollout(ig.initial(pos.clone(), vel.clone()), 4) + comp = torch.compile(ig.rollout, fullgraph=True, backend=_BACKEND)( + ig.initial(pos.clone(), vel.clone()), 4 + ) + assert torch.allclose(eager.pos, comp.pos, atol=1e-12) + assert torch.allclose(eager.forces, comp.forces, atol=1e-12) + assert torch.allclose(eager.energy, comp.energy, atol=1e-12) + assert nl.rebuild_count == 0, "a frozen integrator rebuilt the list" + + +def test_live_ljcut_advances_eagerly_and_drives_the_policy(): + """The production counterpart: ``rebuild=True`` keeps the loop eager + (``advance_n``) and the list's policy actually fires — the compiled arm + above must not be passing because nothing ever rebuilds.""" + ig, nl, pos = _ljcut_ig(rebuild=True) + state = ig.initial(pos, torch.zeros_like(pos)) + ig.advance_n(state, 4) + assert nl.rebuild_count > 0 + + def _harm_ig(gamma: float): return LangevinVerletIntegrator( HarmonicForceField(1.0).to(_DTYPE), dt=0.01, gamma=gamma, kbt=1.0, mass=1.0, seed=3 @@ -75,9 +125,9 @@ def test_rollout_compile_nve_matches_eager(): def _pinet_ig(dtype: torch.dtype): - template = _template() - model = _tiny_potential() - model(template.clone(), compute_forces=False) # warmup lazy params + template = make_pinet_template() + model = make_tiny_potential() + model(template.clone()) # warmup lazy params if dtype == torch.float64: model = model.to(torch.float64) template["atoms", "pos"] = template["atoms", "pos"].to(torch.float64) @@ -106,7 +156,7 @@ def test_pinet_step_fullgraph_compiles_and_matches_eager(dtype): comp = torch.compile(ig.step, fullgraph=True, backend=_BACKEND)(st, noise) atol = 1e-10 if dtype == torch.float64 else 1e-4 assert torch.allclose(eager.pos, comp.pos, atol=atol) - assert torch.allclose(eager.force, comp.force, atol=atol) + assert torch.allclose(eager.forces, comp.forces, atol=atol) assert torch.allclose(eager.energy, comp.energy, atol=atol) diff --git a/tests/test_molix/test_md/test_driver.py b/tests/test_molix/test_md/test_driver.py new file mode 100644 index 0000000..06f37ed --- /dev/null +++ b/tests/test_molix/test_md/test_driver.py @@ -0,0 +1,299 @@ +"""Tests for molix.md.driver — constructor wiring and per-call contracts only.""" + +import pytest +import torch + +from molix.md import ( + EV_PER_AMU_A2_FS2, + MD, + HarmonicForceField, + LangevinVerletIntegrator, + LennardJonesCutForceField, + MaxwellBoltzmann, + MDHook, + MDObservables, + MDRunner, + MDState, + NeighborList, +) +from tests.test_molix.test_md.conftest import make_cubic_lattice + + +@pytest.fixture +def system(): + torch.manual_seed(0) + return torch.randn(6, 3, dtype=torch.float64), torch.zeros(6, 3, dtype=torch.float64) + + +def _lj_cut_argon(skin: float) -> tuple[LennardJonesCutForceField, torch.Tensor]: + """Argon lj/cut over the 64-atom lattice with a policy-configured list. + + Argon in the integrator's (amu, Å, fs) system: ε = 0.0103 eV, σ = 2.5 Å, + r_cut = 3.5 Å over a 12 Å cube (minimum perpendicular half-width 6.0 Å, so + ``r_build = 3.5 + skin`` stays admissible up to ``skin = 2.5``). + + ``capacity_factor=2.5`` is measured, not defensive: a lattice start + under-counts the edges a warm run reaches (600 live against the 519 rows + the default would allocate), and the overflow guard is not what these + tests pin. + """ + pos, cell = make_cubic_lattice(n_side=4, spacing=3.0) + neighbors = NeighborList( + cell=cell, + cutoff=3.5, + positions=pos, + skin=skin, + every=1, + delay=0, + check=True, + capacity_factor=2.5, + ) + force = LennardJonesCutForceField( + epsilon=0.0103 / EV_PER_AMU_A2_FS2, # argon well depth, eV -> amu A^2/fs^2 + sigma=2.5, + neighbors=neighbors, + cutoff=3.5, + ) + return force, pos + + +class _TotalEnergyHook(MDHook): + """Sample the conserved quantity once per step — the drift observable.""" + + def __init__(self) -> None: + self.totals: list[torch.Tensor] = [] + + def on_step_end(self, runner: MDRunner, step: int, obs: MDObservables) -> None: + self.totals.append(obs.total.detach().clone()) + + +def _argon_nve(skin: float, *, n_steps: int = 100) -> tuple[NeighborList, MDState, list[float]]: + """100 steps of NVE argon through the **public** ``MD`` path at one skin. + + Deterministic CPU float64: seeded Maxwell-Boltzmann velocities, γ = 0, no + wall clock, no filesystem, no network. No cadence kwarg — the driver + derives the integrator's switch from the force field, and the list owns + when to rebuild. + + Returns: + ``(neighbors, final_state, total_energies)`` — the list (for + ``rebuild_count`` / ``ndanger``), the final :class:`MDState`, and the + per-step total energy in amu·Å²/fs². + """ + force, pos = _lj_cut_argon(skin) + sampler = _TotalEnergyHook() + velocities = MaxwellBoltzmann(39.95, n_atoms=64).sample(300.0, seed=0) + md = MD(force, mass=39.95, dt=4.0, gamma=0.0, dtype=torch.float64, hooks=[sampler]) + md.set_potential_dtype(torch.float64) + final = md.run(pos, velocities, n_steps, chunk=1) + return force.neighbors, final, [float(total) for total in sampler.totals] + + +def _drift(totals: list[float]) -> float: + """``max_t |E(t) − E(0)| / |E(0)|`` — the dimensionless conservation metric.""" + reference = totals[0] + return max(abs(total - reference) for total in totals) / abs(reference) + + +class TestMD: + """Test the MD component.""" + + def test_runs_and_returns_final_state(self, system): + """The run yields the advanced typed state, not the initial one.""" + pos, vel = system + md = MD(HarmonicForceField(k=1.0), mass=1.0, dt=0.01, dtype=torch.float64) + out = md.run(pos, vel, 3) + assert out.pos.shape == pos.shape + assert not torch.allclose(out.pos, pos) + + def test_dtype_governs_the_md_side_only(self, system): + """``MD(dtype=)`` casts state + integrator constants, never the model: + the MD process and the inference process are studied separately.""" + pos, vel = system + force = HarmonicForceField(k=1.0) # constructed fp32 + md = MD(force, mass=1.0, dt=0.01, dtype=torch.float64) + out = md.run(pos.float(), vel.float(), 1) + assert out.pos.dtype == torch.float64 # trajectory in the MD dtype + assert out.forces.dtype == torch.float64 # boundary cast into the state + assert force.k.dtype == torch.float32 # the potential was left alone + assert md.integrator.dt.dtype == torch.float64 # step constants follow + + def test_set_potential_dtype_casts_the_model_explicitly(self, system): + """The inference-side precision is its own explicit axis.""" + force = HarmonicForceField(k=1.0) + md = MD(force, mass=1.0, dt=0.01, dtype=torch.float64) + md.set_potential_dtype(torch.float64) + assert force.k.dtype == torch.float64 + + def test_dtype_none_leaves_precision_alone(self, system): + """Omitting dtype must not silently downcast a float64 caller.""" + pos, vel = system + md = MD(HarmonicForceField(k=1.0).double(), mass=1.0, dt=0.01) + assert md.run(pos, vel, 1).pos.dtype == torch.float64 + + def test_autocast_keeps_state_in_the_run_dtype(self, system): + """bf16-mixed: the model may run reduced, the trajectory may not.""" + pos, vel = system + md = MD( + HarmonicForceField(k=1.0), + mass=1.0, + dt=0.01, + dtype=torch.float32, + autocast_dtype=torch.bfloat16, + ) + out = md.run(pos, vel, 1) + assert out.pos.dtype == torch.float32 + assert out.forces.dtype == torch.float32 + + def test_accepts_a_constructed_integrator(self): + """The integrator seam: the caller-built integrator is wired through.""" + force = HarmonicForceField(k=1.0).double() + integrator = LangevinVerletIntegrator(force, dt=0.01, gamma=0.0, kbt=0.0, mass=1.0) + md = MD(force, mass=1.0, integrator=integrator) + assert md.integrator is integrator + assert md.runner.integrator is integrator + + def test_integrator_excludes_langevin_parameters(self): + """dt/kbt/temperature parameterise the default integrator only.""" + force = HarmonicForceField(k=1.0) + integrator = LangevinVerletIntegrator(force, dt=0.01, gamma=0.0, kbt=0.0, mass=1.0) + with pytest.raises(ValueError, match="mutually exclusive"): + MD(force, mass=1.0, dt=0.5, integrator=integrator) + + def test_integrator_must_wrap_the_same_force(self): + other = HarmonicForceField(k=2.0) + integrator = LangevinVerletIntegrator(other, dt=0.01, gamma=0.0, kbt=0.0, mass=1.0) + with pytest.raises(ValueError, match="must be the force field"): + MD(HarmonicForceField(k=1.0), mass=1.0, integrator=integrator) + + def test_dt_required_without_integrator(self): + with pytest.raises(ValueError, match="dt is required"): + MD(HarmonicForceField(k=1.0), mass=1.0) + + def test_temperature_sets_kbt(self): + """`temperature=` is the ergonomic form of `kbt=`; both cannot be given.""" + force = HarmonicForceField(k=1.0) + MD(force, mass=1.0, dt=0.01, gamma=0.1, temperature=300.0) + with pytest.raises(ValueError, match="not both"): + MD(force, mass=1.0, dt=0.01, gamma=0.1, temperature=300.0, kbt=1.0) + + def test_thermostat_without_temperature_is_rejected(self): + """gamma>0 with kbt=0 is a thermostat at 0 K — almost certainly a mistake.""" + with pytest.raises(ValueError, match="needs kbt or temperature"): + MD(HarmonicForceField(k=1.0), mass=1.0, dt=0.01, gamma=0.1) + + +class TestMDNeighborPolicy: + """The driver derives the rebuild switch; it never owns a cadence. + + ``MD(rebuild_every=)`` was a second owner of a decision the neighbour list + already makes (``skin`` / ``every`` / ``delay`` / ``check``), and the + step-start ``NeighborListHook`` was a third at the wrong seam. Both are + gone: ``MD`` builds the integrator and lets it read + ``ForceField.rebuilds_neighbors``. + """ + + def test_the_cadence_kwarg_is_gone(self): + """Hard removal (``stage: experimental``), no back-compat shim: the + migration is ``NeighborList(skin=, every=, delay=, check=)``.""" + with pytest.raises(TypeError): + MD(HarmonicForceField(k=1.0), mass=1.0, dt=0.01, rebuild_every=1) + + def test_no_step_start_hook_is_installed(self): + """The wrong-seam hook is not silently replaced by another default.""" + md = MD(HarmonicForceField(k=1.0), mass=1.0, dt=0.01) + assert md.runner.hooks == [] + + def test_the_switch_follows_the_force_field(self): + """Listless ⇒ off, list-backed ⇒ on, with no kwarg either way.""" + listless = MD(HarmonicForceField(k=1.0), mass=1.0, dt=0.01) + force, _ = _lj_cut_argon(skin=1.0) + live = MD(force, mass=39.95, dt=4.0, dtype=torch.float64) + assert listless.integrator.rebuild is False + assert live.integrator.rebuild is True + + def test_the_autocast_wrapper_delegates_the_capability(self): + """``autocast_dtype`` wraps the force field *before* the integrator is + built, so the wrapper must forward the flag — otherwise a bf16 run + silently freezes its neighbour list.""" + force, _ = _lj_cut_argon(skin=1.0) + md = MD(force, mass=39.95, dt=4.0, autocast_dtype=torch.bfloat16) + assert md.force is not force # the wrapper is in the way + assert md.integrator.rebuild is True + + def test_a_compiled_force_field_delegates_the_capability(self): + """``torch.compile(ff)`` returns an ``OptimizedModule`` that forwards + attribute reads to the wrapped module — the assumption the GPU + benchmark's compiled-force-field path depends on.""" + force, _ = _lj_cut_argon(skin=1.0) + compiled = torch.compile(force, backend="eager") + md = MD(compiled, mass=39.95, dt=4.0, dtype=torch.float64) + assert md.integrator.rebuild is True + + def test_a_skin_costs_no_energy_conservation(self): + """Invariant (c): a skin-gated run must conserve energy as well as one + that runs the policy at every force evaluation. + + A pair inside ``r_cut`` but missing from the list contributes an O(1) + force error (lj/cut shifts the *energy* continuous at the cutoff, not + the force), which integrates into a one-signed leak — so a broken skin + shows up as drift, not as noise. The ``skin=0`` baseline is itself + nonzero (finite-``dt`` velocity-Verlet), which is asserted so the ratio + cannot pass vacuously. + """ + _, _, gated = _argon_nve(skin=1.0) + _, _, every_eval = _argon_nve(skin=0.0) + baseline = _drift(every_eval) + assert baseline > 0.0, "the no-skin baseline must have real discretisation drift" + assert _drift(gated) <= 3.0 * baseline + + def test_the_skin_buys_rebuilds_without_moving_the_physics(self): + """Invariant (f): the observables the public path exposes. + + ``rebuild_count(skin=0) == 100`` is the anti-vacuity pin — a wiring + that never calls the policy satisfies every equality and monotonicity + assertion here but not that literal. One policy call per force + evaluation over 100 steps, and the entry evaluation sits at the build + positions (``max_d2 == 0``, strict ``>``), so it does not rebuild. + ``ndanger(skin=0) == 99`` (not 100): the declined entry evaluation + consumes one ``ago`` tick, so step 1's rebuild lands at ``ago == 2`` + (not dangerous) and only the remaining 99 rebuilds — each at + ``ago == 1 == max(every, delay)`` after a reset — count. Link 04's + degenerate-limit alarm, off by exactly the entry evaluation. + """ + arms = {skin: _argon_nve(skin=skin) for skin in (0.0, 0.5, 1.0)} + counts = [arms[skin][0].rebuild_count for skin in (0.0, 0.5, 1.0)] + assert counts == sorted(counts, reverse=True) + assert counts[0] == 100 + assert counts[-1] < counts[0] + assert arms[0.0][0].ndanger == 99 + assert arms[0.5][0].ndanger == 0 + assert arms[1.0][0].ndanger == 0 + reference = arms[0.0][1].energy + for skin in (0.5, 1.0): + torch.testing.assert_close( + arms[skin][1].energy, reference, atol=1e-10, rtol=0, msg=f"skin={skin}" + ) + + +class TestMaxwellBoltzmann: + """Initial-velocity sampling — its own component, off the run driver.""" + + def test_removes_com_momentum(self): + """Net momentum must vanish, which is what makes the NVE dof 3N-3.""" + mass = torch.full((8,), 2.0) + vel = MaxwellBoltzmann(mass).sample(300.0, seed=1) + assert vel.shape == (8, 3) + assert float((mass.reshape(-1, 1).double() * vel).sum(0).abs().max()) < 1e-12 + + def test_is_reproducible(self): + """Same seed, same velocities — a run must be reproducible from its args.""" + sampler = MaxwellBoltzmann(torch.ones(5)) + assert torch.equal(sampler.sample(300.0, seed=3), sampler.sample(300.0, seed=3)) + + def test_scalar_mass_needs_n_atoms(self): + """A scalar mass carries no system size — silent (1, 3) output was a bug.""" + with pytest.raises(ValueError, match="n_atoms"): + MaxwellBoltzmann(1.0) + vel = MaxwellBoltzmann(1.0, n_atoms=7).sample(300.0, seed=0) + assert vel.shape == (7, 3) diff --git a/tests/test_molix/test_md/test_forcefield.py b/tests/test_molix/test_md/test_forcefield.py new file mode 100644 index 0000000..3593ef3 --- /dev/null +++ b/tests/test_molix/test_md/test_forcefield.py @@ -0,0 +1,461 @@ +"""Tests for molix.md.forcefield — binding models to systems.""" + +import pytest +import torch +from tensordict import TensorDict + +from molix.md import ( + CallableForceField, + ForceField, + ForceOutput, + HarmonicForceField, + LennardJonesCutForceField, + LennardJonesForceField, + NeighborList, + PeriodicPotentialForceField, + PotentialForceField, +) +from tests.test_molix.test_md.conftest import ( + make_cubic_lattice, + make_pinet_template, + make_tiny_potential, +) + + +class _RecordingNeighbors: + """Neighbour-strategy stub recording *which* entry point a force field used. + + The discriminator for the seam's semantic shift: ``rebuild_neighbors`` + means "run the list's policy" (:meth:`update`), never "build now" + (:meth:`rebuild`, which stays available as the explicit escape hatch). + """ + + edge_index = torch.zeros(1, 2, dtype=torch.long) + shifts = torch.zeros(1, 3) + num_edges = 0 + capacity = 1 + cutoff = 5.0 + skin = 0.0 + + def __init__(self) -> None: + self.rebuilds: list[torch.Tensor] = [] + self.updates: list[torch.Tensor] = [] + + def rebuild(self, positions: torch.Tensor) -> None: + self.rebuilds.append(positions.detach().clone()) + + def update(self, positions: torch.Tensor) -> bool: + self.updates.append(positions.detach().clone()) + return False + + def to( + self, + device: torch.device | str | torch.dtype | None = None, + dtype: torch.dtype | None = None, + ) -> "_RecordingNeighbors": + return self + + +class TestPotentialForceField: + """The molpot-potential adapter over a fixed open-system template.""" + + def test_force_seam_tracks_live_geometry(self): + """PotentialForceField must recompute edge geometry from the live + positions, not a frozen template ``edge_diff`` — regression for the + constant-PES bug where swapping only ``pos`` left energy/force pinned + to the initial geometry.""" + template = make_pinet_template() + ref = make_tiny_potential() + ref(template.clone()) # warmup lazy params + ff = PotentialForceField(ref, template) + pos0 = template["atoms", "pos"] + out0 = ff(pos0) + torch.manual_seed(1) + pos1 = pos0 + 0.3 * torch.randn_like(pos0) # non-rigid displacement + out1 = ff(pos1) + assert (out1.energy - out0.energy).abs().item() > 1e-6, "energy frozen at initial geometry" + assert (out1.forces - out0.forces).abs().max().item() > 1e-6, ( + "forces frozen at initial geometry" + ) + + def test_to_dtype_casts_the_bound_system(self): + """``.to(float64)`` must reach the working batch, not just parameters — + the silent-no-op ``_apply`` gap was a review finding.""" + template = make_pinet_template() + potential = make_tiny_potential() + potential(template.clone()) # warmup lazy params + ff = PotentialForceField(potential, template).to(torch.float64) + assert ff._dtype == torch.float64 + assert ff._work["atoms", "pos"].dtype == torch.float64 + out = ff(template["atoms", "pos"].to(torch.float64)) + assert out.energy.dtype == torch.float64 + assert out.forces.dtype == torch.float64 + + def test_accepts_positions_in_any_dtype(self): + """The force field owns its precision: an fp64 trajectory may drive an + fp32 model, and the output stays in the model's dtype (the integrator + casts at the boundary).""" + template = make_pinet_template() + potential = make_tiny_potential() + potential(template.clone()) + ff = PotentialForceField(potential, template) # fp32 model + out = ff(template["atoms", "pos"].to(torch.float64)) + assert out.forces.dtype == torch.float32 + + +class TestCallableForceField: + """The escape hatch for non-TensorDict force providers.""" + + def test_wraps_a_plain_tuple_callable(self): + ff = CallableForceField(lambda pos: ((pos * pos).sum(), -2.0 * pos)) + pos = torch.randn(5, 3) + out = ff(pos) + assert isinstance(out, ForceOutput) + assert torch.equal(out.forces, -2.0 * pos) + + def test_applies_the_energy_scale(self): + ff = CallableForceField( + lambda pos: (torch.tensor(2.0), torch.ones_like(pos)), energy_scale=0.5 + ) + out = ff(torch.zeros(3, 3)) + assert float(out.energy) == 1.0 + assert torch.equal(out.forces, 0.5 * torch.ones(3, 3)) + + def test_delegates_the_policy_to_the_neighbor_list(self): + """The bound list's *policy* runs (``update``), not a forced ``rebuild``. + + The list owns the cadence (skin / every / delay / check); the force + field only passes on the positions being evaluated and lets it decline. + """ + recorder = _RecordingNeighbors() + ff = CallableForceField(lambda pos: (pos.sum(), pos), neighbors=recorder) + pos = torch.zeros(2, 3) + ff.rebuild_neighbors(pos) + assert len(recorder.updates) == 1 + assert torch.equal(recorder.updates[0], pos) + assert recorder.rebuilds == [] + + def test_without_neighbors_rebuild_is_a_noop(self): + ff = CallableForceField(lambda pos: (pos.sum(), pos)) + ff.rebuild_neighbors(torch.zeros(2, 3)) # must not raise + + +class TestPotentialContract: + """A potential that writes no forces must fail loudly, not integrate garbage.""" + + def test_missing_forces_raises(self): + template = make_pinet_template() + + class _EnergyOnly(torch.nn.Module): + def forward(self, td): + td["graphs", "energy"] = torch.zeros(1) + return td + + ff = PotentialForceField(_EnergyOnly(), template) + with pytest.raises(RuntimeError, match="wrote no"): + ff(template["atoms", "pos"]) + + +class _ShiftAwarePairPotential(torch.nn.Module): + """Pair potential reading ``edges.shifts`` — envelope exactly zero beyond cutoff.""" + + def __init__(self, cutoff: float) -> None: + super().__init__() + self.cutoff = float(cutoff) + + def forward(self, td: TensorDict) -> TensorDict: + pos = td["atoms", "pos"] + ei = td["edges", "edge_index"] + shifts = td["edges", "shifts"] + leaf = pos.detach().requires_grad_(True) + with torch.enable_grad(): + vec = leaf[ei[:, 1]] - leaf[ei[:, 0]] + shifts + r = torch.linalg.norm(vec, dim=-1) + pair = torch.where(r < self.cutoff, (r - self.cutoff) ** 2, torch.zeros_like(r)) + energy = pair.sum() + (grad,) = torch.autograd.grad(energy, leaf) + td["graphs", "energy"] = energy.detach().reshape(1) + td["atoms", "forces"] = -grad + return td + + +def _periodic_template(pos: torch.Tensor) -> TensorDict: + n = pos.shape[0] + return TensorDict( + { + "atoms": TensorDict( + { + "pos": pos, + "Z": torch.ones(n, dtype=torch.long), + "batch": torch.zeros(n, dtype=torch.long), + }, + batch_size=[n], + ) + }, + batch_size=[], + ) + + +class TestPeriodicPotentialForceField: + """The joined periodic component: potential + rebuilding fixed-capacity list.""" + + def test_dead_edge_padding_is_invisible(self): + """Different capacity factors (more padding) must give identical physics.""" + pos, cell = make_cubic_lattice() + outs = [] + for factor in (1.1, 3.0): + nl = NeighborList(cell=cell, cutoff=3.5, positions=pos, capacity_factor=factor) + ff = PeriodicPotentialForceField( + _ShiftAwarePairPotential(3.5), _periodic_template(pos), neighbors=nl + ) + outs.append(ff(pos)) + assert torch.equal(outs[0].energy, outs[1].energy) + assert torch.equal(outs[0].forces, outs[1].forces) + + def test_rebuild_is_visible_through_the_bound_buffers(self): + """The working batch holds the list's buffers by reference, so an + in-place rebuild changes the energy without any re-binding. + + Compression (not dilation) is the discriminating displacement: with an + envelope that is exactly zero beyond the cutoff, a stale *superset* + list gives the same energy — only pairs *entering* the cutoff that the + stale list has never seen can differ. + """ + pos, cell = make_cubic_lattice() + nl = NeighborList(cell=cell, cutoff=3.5, positions=pos, capacity_factor=8.0) + ff = PeriodicPotentialForceField( + _ShiftAwarePairPotential(3.5), _periodic_template(pos), neighbors=nl + ) + compressed = pos * 0.8 # second-neighbour pairs enter the cutoff + stale = ff(compressed) # list still from the original positions + ff.rebuild_neighbors(compressed) + fresh = ff(compressed) + assert nl.rebuild_count == 1 + assert not torch.equal(stale.energy, fresh.energy) + + def test_a_cast_keeps_the_rebuild_visible(self): + """``.to(dtype)`` must leave the by-reference tie alive. + + ``TensorDict.apply`` in the parent ``_apply`` produces *new* leaf + tensors and ``NeighborList.to`` rebinds ``shifts`` to a new tensor, so a + cast severs the tie twice over unless the owner re-binds afterwards — a + later rebuild would then update buffers the potential no longer sees, + and the PES would freeze silently. + + Asserted through the public energy rather than through ``_work`` on + purpose: this is the property a user has, and it must survive the + ownership move of the binding knowledge onto ``NeighborList.build``. + """ + pos, cell = make_cubic_lattice() + nl = NeighborList(cell=cell, cutoff=3.5, positions=pos, capacity_factor=8.0) + ff = PeriodicPotentialForceField( + _ShiftAwarePairPotential(3.5), _periodic_template(pos), neighbors=nl + ).to(torch.float32) + compressed = (pos * 0.8).to(torch.float32) + stale = ff(compressed) + ff.rebuild_neighbors(compressed) + fresh = ff(compressed) + assert nl.rebuild_count == 1 + assert not torch.equal(stale.energy, fresh.energy) + + +class TestHarmonicForceField: + def test_energy_and_force_are_consistent(self): + pos = torch.randn(6, 3, dtype=torch.float64, requires_grad=True) + out = HarmonicForceField(k=2.0).to(torch.float64)(pos) + (ref,) = torch.autograd.grad(out.energy, pos) + assert torch.allclose(out.forces, -ref) + + +class TestLennardJonesCutForceField: + """Periodic truncated-shifted LJ over the fixed-capacity neighbour list.""" + + _EPS, _SIGMA = 0.7, 1.1 + + def _dimer(self, d: float, *, box: float = 20.0, cutoff: float = 5.0, shift: bool = True): + pos = torch.tensor([[0.0, 0.0, 0.0], [d, 0.0, 0.0]], dtype=torch.float64) + cell = torch.eye(3, dtype=torch.float64) * box + nl = NeighborList(cell=cell, cutoff=cutoff, positions=pos) + ff = LennardJonesCutForceField( + epsilon=self._EPS, sigma=self._SIGMA, neighbors=nl, shift=shift + ).to(torch.float64) + return ff, pos + + def test_force_matches_autograd(self): + """Closed-form forces must equal -dE/dpos through the mask and the shifts.""" + pos, cell = make_cubic_lattice() + torch.manual_seed(0) + pos = pos + 0.3 * torch.randn_like(pos) + nl = NeighborList(cell=cell, cutoff=3.5, positions=pos) + ff = LennardJonesCutForceField(epsilon=0.7, sigma=2.5, neighbors=nl).to(torch.float64) + leaf = pos.clone().requires_grad_(True) + out = ff(leaf) + (ref,) = torch.autograd.grad(out.energy, leaf) + assert torch.allclose(out.forces, -ref, atol=1e-10), "lj/cut closed form != -dE/dx" + + def test_matches_all_pairs_in_the_open_limit(self): + """A cutoff spanning the whole cluster + shift=False is the all-pairs LJ.""" + torch.manual_seed(1) + pos = torch.randn(8, 3, dtype=torch.float64) * 1.5 + 15.0 # blob at box centre + cell = torch.eye(3, dtype=torch.float64) * 30.0 + nl = NeighborList(cell=cell, cutoff=14.0, positions=pos) + cut = LennardJonesCutForceField(epsilon=0.9, sigma=1.2, neighbors=nl, shift=False).to( + torch.float64 + ) + ref = LennardJonesForceField(epsilon=0.9, sigma=1.2).to(torch.float64) + out, expected = cut(pos), ref(pos) + assert torch.allclose(out.energy, expected.energy, atol=1e-10) + assert torch.allclose(out.forces, expected.forces, atol=1e-10) + + def test_shift_makes_energy_continuous_at_the_cutoff(self): + """Shifted: E→0 continuously at r_cut; unshifted: E→E_lj(r_cut) (the step).""" + r_cut = 3.0 + just_inside = r_cut - 1e-9 + shifted, pos = self._dimer(just_inside, cutoff=r_cut, shift=True) + unshifted, _ = self._dimer(just_inside, cutoff=r_cut, shift=False) + sr6 = (self._SIGMA / r_cut) ** 6 + e_at_cut = 4.0 * self._EPS * (sr6 * sr6 - sr6) + assert abs(float(shifted(pos).energy)) < 1e-8 + assert abs(float(unshifted(pos).energy) - e_at_cut) < 1e-8 + + def test_dead_edge_padding_is_invisible(self): + """Different capacity factors (more padding) must give identical physics.""" + pos, cell = make_cubic_lattice() + outs = [] + for factor in (1.1, 4.0): + nl = NeighborList(cell=cell, cutoff=3.5, positions=pos, capacity_factor=factor) + ff = LennardJonesCutForceField(epsilon=0.8, sigma=2.5, neighbors=nl).to(torch.float64) + outs.append(ff(pos)) + assert torch.equal(outs[0].energy, outs[1].energy) + assert torch.equal(outs[0].forces, outs[1].forces) + + def test_minimum_image_across_the_boundary(self): + """Atoms near opposite faces interact at the wrapped distance.""" + box, r_cut = 12.0, 3.0 + pos = torch.tensor([[0.6, 0.0, 0.0], [box - 0.6, 0.0, 0.0]], dtype=torch.float64) + cell = torch.eye(3, dtype=torch.float64) * box + nl = NeighborList(cell=cell, cutoff=r_cut, positions=pos) + ff = LennardJonesCutForceField(epsilon=self._EPS, sigma=self._SIGMA, neighbors=nl).to( + torch.float64 + ) + open_ff, open_pos = self._dimer(1.2, cutoff=r_cut) + assert torch.allclose(ff(pos).energy, open_ff(open_pos).energy, atol=1e-12) + + def test_rebuild_tracks_pair_departure(self): + """A pair leaving the cutoff vanishes from the PES after a rebuild.""" + ff, pos = self._dimer(1.5, cutoff=3.0) + assert float(ff(pos).energy) != 0.0 + apart = torch.tensor([[0.0, 0.0, 0.0], [5.0, 0.0, 0.0]], dtype=torch.float64) + ff.rebuild_neighbors(apart) + assert float(ff(apart).energy) == 0.0 + assert torch.equal(ff(apart).forces, torch.zeros_like(apart)) + + def test_cutoff_defaults_to_the_lists(self): + ff, _ = self._dimer(1.5, cutoff=3.0) + assert float(ff.cutoff_sq) == pytest.approx(9.0) + + def test_rejects_cutoff_beyond_the_list_horizon(self): + pos, cell = make_cubic_lattice() + nl = NeighborList(cell=cell, cutoff=3.0, positions=pos) + with pytest.raises(ValueError, match="horizon"): + LennardJonesCutForceField(epsilon=1.0, sigma=1.0, neighbors=nl, cutoff=4.0) + + def test_to_reaches_the_neighbor_buffers(self): + """``.to(dtype)`` must cast the list's shifts alongside the module buffers.""" + ff, pos = self._dimer(1.5) + ff.to(torch.float32) + assert ff.neighbors.shifts.dtype == torch.float32 + out = ff(pos.to(torch.float32)) + assert out.energy.dtype == torch.float32 + + +class TestForceFieldNeighborSeam: + """The capability flag and the policy semantics of the shared rebuild seam. + + ``rebuilds_neighbors`` answers *can this force field run a neighbour + policy* (a read-only property mirroring ``Integrator.removed_dof``, so the + integrator never duck-reads ``getattr(force, "neighbors", None)``); + ``rebuild_neighbors(pos)`` runs that policy at the positions being + evaluated and lets the list decline. + """ + + def _lj_cut(self, *, skin: float) -> tuple[LennardJonesCutForceField, torch.Tensor]: + """lj/cut over a policy-configured list on the 64-atom, 12 Å lattice.""" + pos, cell = make_cubic_lattice(n_side=4, spacing=3.0) + nl = NeighborList( + cell=cell, cutoff=3.5, positions=pos, skin=skin, every=1, delay=0, check=True + ) + ff = LennardJonesCutForceField(epsilon=0.7, sigma=2.5, neighbors=nl).to(torch.float64) + return ff, pos + + def test_a_listless_force_field_declares_no_policy(self): + """Nothing to refresh ⇒ ``False``, so the integrator's switch stays off.""" + template = _periodic_template(make_cubic_lattice()[0]) + assert ForceField().rebuilds_neighbors is False + assert HarmonicForceField(k=1.0).rebuilds_neighbors is False + assert LennardJonesForceField(epsilon=1.0, sigma=1.0).rebuilds_neighbors is False + assert ( + PotentialForceField(_ShiftAwarePairPotential(3.5), template).rebuilds_neighbors is False + ) + assert CallableForceField(lambda pos: (pos.sum(), pos)).rebuilds_neighbors is False + + def test_a_list_backed_force_field_declares_a_policy(self): + """A live list ⇒ ``True``; ``CallableForceField`` answers per instance.""" + pos, cell = make_cubic_lattice() + periodic = PeriodicPotentialForceField( + _ShiftAwarePairPotential(3.5), + _periodic_template(pos), + neighbors=NeighborList(cell=cell, cutoff=3.5, positions=pos), + ) + lj_cut, _ = self._lj_cut(skin=0.0) + bound = CallableForceField(lambda pos: (pos.sum(), pos), neighbors=_RecordingNeighbors()) + assert periodic.rebuilds_neighbors is True + assert lj_cut.rebuilds_neighbors is True + assert bound.rebuilds_neighbors is True + + def test_lj_cut_asks_the_policy_and_accepts_a_refusal(self): + """``rebuild_neighbors`` is a *question* now: inside the half-skin the + list declines, so a displacement smaller than ``skin/2`` must leave + ``rebuild_count`` at zero while the policy clock still ticks.""" + ff, pos = self._lj_cut(skin=2.0) # half-skin 1.0 A + moved = pos.clone() + moved[0, 0] += 0.1 # far inside the half-skin + ff.rebuild_neighbors(moved) + assert ff.neighbors.rebuild_count == 0, "the seam forced a build the policy declined" + assert ff.neighbors.ago == 1, "the policy clock did not advance" + + def test_the_periodic_potential_asks_the_same_policy(self): + """The TensorDict-side owner delegates to the same ``update`` gate.""" + pos, cell = make_cubic_lattice(n_side=4, spacing=3.0) + nl = NeighborList( + cell=cell, cutoff=3.5, positions=pos, skin=2.0, every=1, delay=0, check=True + ) + ff = PeriodicPotentialForceField( + _ShiftAwarePairPotential(3.5), _periodic_template(pos), neighbors=nl + ) + moved = pos.clone() + moved[0, 0] += 0.1 + ff.rebuild_neighbors(moved) + assert nl.rebuild_count == 0 + assert nl.ago == 1 + + def test_the_seam_never_forces_a_build(self): + """Recorded at the strategy: ``update`` yes, ``rebuild`` never.""" + recorder = _RecordingNeighbors() + ff = CallableForceField(lambda pos: (pos.sum(), pos), neighbors=recorder) + pos = torch.zeros(4, 3) + ff.rebuild_neighbors(pos) + ff.rebuild_neighbors(pos) + assert len(recorder.updates) == 2 + assert recorder.rebuilds == [] + + def test_a_forced_build_is_still_reachable_on_the_list(self): + """The escape hatch the docstring promises: ``neighbors.rebuild(pos)`` + builds unconditionally — inside the half-skin, where the policy would + decline — counts exactly once, and re-phases the policy clock.""" + ff, pos = self._lj_cut(skin=2.0) + moved = pos.clone() + moved[0, 0] += 0.1 # a displacement ``update`` would refuse to build on + ff.neighbors.rebuild(moved) + assert ff.neighbors.rebuild_count == 1 + assert ff.neighbors.ago == 0 diff --git a/tests/test_molix/test_md/test_integrators.py b/tests/test_molix/test_md/test_integrators.py new file mode 100644 index 0000000..cf23679 --- /dev/null +++ b/tests/test_molix/test_md/test_integrators.py @@ -0,0 +1,287 @@ +"""Tests for the BAOAB Langevin velocity-Verlet integrator component. + +Single-function correctness only (repo rule: no e2e under ``tests/``). +Long-horizon physics validation — NVE energy conservation, equipartition — +lives in ``benchmarks/verify_md_lj_nve.py``, not here. +""" + +import pytest +import torch + +from molix.md import ( + HarmonicForceField, + Integrator, + LangevinVerletIntegrator, + LennardJonesCutForceField, + MDState, + NeighborList, +) +from tests.test_molix.test_md.conftest import make_cubic_lattice + +_DTYPE = torch.float64 + + +def _ig(k: float = 1.0, **kw): + return LangevinVerletIntegrator(HarmonicForceField(k).to(_DTYPE), **kw) + + +def test_step_constants_match_the_closed_form(): + """The precomputed buffers are the BAOAB constants from the paper.""" + import math + + dt, gamma, kbt, mass = 0.05, 2.0, 1.5, 2.0 + ig = _ig(1.0, dt=dt, gamma=gamma, kbt=kbt, mass=mass) + assert float(ig.c1) == pytest.approx(math.exp(-gamma * dt)) + assert float(ig.c2) == pytest.approx(math.sqrt(1.0 - math.exp(-2.0 * gamma * dt))) + assert float(ig.sigma) == pytest.approx(math.sqrt(kbt / mass)) + assert float(ig.inv_mass) == pytest.approx(1.0 / mass) + + +def test_step_is_one_baoab_update(): + """``step`` applies exactly B-A-O-A-B with the precomputed constants.""" + torch.manual_seed(0) + ig = _ig(1.0, dt=0.05, gamma=2.0, kbt=1.5, mass=2.0) + state = ig.initial(torch.randn(4, 3, dtype=_DTYPE), torch.randn(4, 3, dtype=_DTYPE)) + noise = torch.randn(4, 3, dtype=_DTYPE) + + half = 0.5 * ig.dt + vel = state.vel + half * state.forces * ig.inv_mass # B (cached entry force) + pos = state.pos + half * vel # A + vel = ig.c1 * vel + ig.c2 * ig.sigma * noise # O + pos = pos + half * vel # A + force = -pos # k = 1: F = -x at the new position + vel = vel + half * force * ig.inv_mass # B + + out = ig.step(state, noise) + assert torch.equal(out.pos, pos) + assert torch.equal(out.vel, vel) + assert torch.equal(out.forces, force) + + +def test_draw_noise_is_seed_reproducible(): + ref = torch.zeros(5, 3, dtype=_DTYPE) + a = _ig(1.0, dt=0.05, gamma=3.0, kbt=1.0, mass=1.0, seed=42) + b = _ig(1.0, dt=0.05, gamma=3.0, kbt=1.0, mass=1.0, seed=42) + c = _ig(1.0, dt=0.05, gamma=3.0, kbt=1.0, mass=1.0, seed=43) + assert torch.equal(a.draw_noise(ref), b.draw_noise(ref)) # same seed -> identical + assert not torch.equal(a.draw_noise(ref), c.draw_noise(ref)) # different seed + + +def test_step_nve_matches_baoab_at_gamma_zero(): + """The elided O step is the float identity, so the fast path is exact.""" + torch.manual_seed(0) + pos = torch.randn(6, 3, dtype=_DTYPE) + vel = torch.randn(6, 3, dtype=_DTYPE) + integ = _ig(1.0, dt=0.01, gamma=0.0, kbt=0.0, mass=1.0).cast_state(_DTYPE) + state = integ.initial(pos, vel) + noise = torch.randn_like(vel) # multiplied by c2=0 — must not matter + via_baoab = integ.step(state, noise) + via_nve = integ.step_nve(state) + assert torch.equal(via_baoab.pos, via_nve.pos) + assert torch.equal(via_baoab.vel, via_nve.vel) + + +def test_advance_n_matches_manual_advance_loop(): + """``advance_n`` visits exactly the states a manual ``advance`` loop does.""" + k, mass, dt, gamma, kbt = 1.0, 1.0, 0.05, 3.0, 1.0 + pos0 = torch.randn(6, 3, dtype=_DTYPE) + vel0 = torch.randn(6, 3, dtype=_DTYPE) + + chunked = _ig(k, dt=dt, gamma=gamma, kbt=kbt, mass=mass, seed=11) + end_a = chunked.advance_n(chunked.initial(pos0.clone(), vel0.clone()), 5) + + ig = _ig(k, dt=dt, gamma=gamma, kbt=kbt, mass=mass, seed=11) + state = ig.initial(pos0.clone(), vel0.clone()) + for _ in range(5): + state = ig.advance(state) + assert torch.equal(end_a.pos, state.pos) + assert torch.equal(end_a.vel, state.vel) + + +class _MidpointEuler(Integrator): + """Minimal conforming subclass: implements only what the ABC demands. + + Guards the extension seam — a subclass that writes exactly ``advance`` (+ + ``rollout``) must inherit working ``initial`` / ``advance_n`` / + ``removed_dof`` defaults. The old ABC declared ``step``/``rollout`` but + the runner called the undeclared ``advance_n`` — a conforming subclass + crashed at runtime. + """ + + def __init__(self, force, dt: float): + super().__init__(force) + self.register_buffer("dt", torch.as_tensor(dt)) + + def advance(self, state: MDState) -> MDState: + vel = state.vel + self.dt * state.forces + pos = state.pos + self.dt * vel + out = self.eval_force(pos) + return MDState(pos, vel, out.forces, out.energy) + + def rollout(self, state: MDState, n_steps: int) -> MDState: + return self.advance_n(state, n_steps) + + +def test_conforming_subclass_inherits_the_abc_defaults(): + ig = _MidpointEuler(HarmonicForceField(1.0).to(_DTYPE), dt=0.01).to(_DTYPE) + state = ig.initial(torch.randn(4, 3, dtype=_DTYPE), torch.zeros(4, 3, dtype=_DTYPE)) + out = ig.advance_n(state, 3) # inherited default: loop over advance + assert out.pos.shape == (4, 3) + assert not torch.equal(out.pos, state.pos) + assert ig.removed_dof == 3 # inherited NVE default + + +def test_removed_dof_follows_the_thermostat(): + assert _ig(1.0, dt=0.01, gamma=0.0, kbt=0.0, mass=1.0).removed_dof == 3 + assert _ig(1.0, dt=0.01, gamma=2.0, kbt=1.0, mass=1.0).removed_dof == 0 + + +def test_eval_force_casts_into_the_state_dtype(): + """An fp32 force field driven by an fp64 state must not leak fp32 into it.""" + from molix.md import CallableForceField + + # Returns strictly fp32 whatever it is fed — the potential's own precision. + ff = CallableForceField(lambda pos: ((pos.float() ** 2).sum(), -2.0 * pos.float())) + ig = LangevinVerletIntegrator(ff, dt=0.01, gamma=0.0, kbt=0.0, mass=1.0).cast_state(_DTYPE) + state = ig.initial(torch.zeros(3, 3, dtype=_DTYPE), torch.zeros(3, 3, dtype=_DTYPE)) + assert state.forces.dtype == _DTYPE + assert state.energy.dtype == _DTYPE + advanced = ig.advance(state) + assert advanced.pos.dtype == _DTYPE + assert advanced.vel.dtype == _DTYPE + + +def test_cast_state_leaves_the_force_field_alone(): + """The MD-side cast must not touch the potential's parameters.""" + ff = HarmonicForceField(1.0) + ig = LangevinVerletIntegrator(ff, dt=0.01, gamma=0.0, kbt=0.0, mass=1.0) + ig.cast_state(torch.float64) + assert ig.dt.dtype == torch.float64 + assert ig.mass_col.dtype == torch.float64 + assert ff.k.dtype == torch.float32 + + +class _CountingHarmonic(HarmonicForceField): + def __init__(self, k: float = 1.0): + super().__init__(k) + self.calls = 0 + + def forward(self, pos): + self.calls += 1 + return super().forward(pos) + + +def test_force_caching_one_eval_per_step(): + """Force caching: exactly one force-field evaluation per step (+1 to seed).""" + ff = _CountingHarmonic(1.0).to(_DTYPE) + ig = LangevinVerletIntegrator(ff, dt=0.05, gamma=1.0, kbt=1.0, mass=1.0, seed=1) + state = ig.initial(torch.zeros(5, 3, dtype=_DTYPE), torch.zeros(5, 3, dtype=_DTYPE)) + for _ in range(10): + state = ig.advance(state) + assert ff.calls == 11 # 1 (initial) + 10 (one per step) + + +def test_mass_must_be_positive(): + with pytest.raises(ValueError, match="strictly positive"): + _ig(1.0, dt=0.01, gamma=0.0, kbt=0.0, mass=-1.0) + with pytest.raises(ValueError, match="strictly positive"): + _ig(1.0, dt=0.01, gamma=0.0, kbt=0.0, mass=torch.tensor([1.0, -2.0, 3.0])) + + +class _CountingNLForce(HarmonicForceField): + """Harmonic well that records the positions passed to rebuild_neighbors.""" + + def __init__(self, k: float = 1.0): + super().__init__(k) + self.rebuild_positions: list[torch.Tensor] = [] + + def rebuild_neighbors(self, pos: torch.Tensor) -> None: + self.rebuild_positions.append(pos.detach().clone()) + + +def _lj_cut_force() -> tuple[LennardJonesCutForceField, torch.Tensor]: + """lj/cut argon over a live list — a force field that owns a policy.""" + pos, cell = make_cubic_lattice(n_side=4, spacing=3.0) + nl = NeighborList(cell=cell, cutoff=3.5, positions=pos, skin=0.5, every=1, delay=0, check=True) + ff = LennardJonesCutForceField(epsilon=0.7, sigma=2.5, neighbors=nl).to(_DTYPE) + return ff, pos + + +class TestIntegratorRebuildSwitch: + """``Integrator.rebuild`` — a construction-time Python bool, not a counter. + + *Whether this integrator asks* the neighbour policy, derived from *whether + the force field can answer* (``ForceField.rebuilds_neighbors``). Static so + dynamo specialises the branch: ``rebuild=False`` leaves the body dead and + ``torch.compile(fullgraph=True)`` still traces one graph, ``rebuild=True`` + runs the list's policy eagerly between force calls. + """ + + def test_the_switch_is_derived_from_the_force_field(self): + """No kwarg ⇒ follow the force field, so no caller has to remember.""" + listless = Integrator(HarmonicForceField(1.0).to(_DTYPE)) + ff, _ = _lj_cut_force() + live = LangevinVerletIntegrator(ff, dt=0.5, gamma=0.0, kbt=0.0, mass=39.95) + assert listless.rebuild is False + assert live.rebuild is True + + def test_the_kwarg_overrides_the_derivation_both_ways(self): + """The frozen-list run and the forced-policy run are both reachable.""" + ff, _ = _lj_cut_force() + frozen = LangevinVerletIntegrator(ff, dt=0.5, gamma=0.0, kbt=0.0, mass=39.95, rebuild=False) + forced = LangevinVerletIntegrator( + HarmonicForceField(1.0).to(_DTYPE), dt=0.01, gamma=0.0, kbt=0.0, mass=1.0, rebuild=True + ) + assert frozen.rebuild is False + assert forced.rebuild is True + + def test_the_seam_fires_once_per_force_evaluation_at_those_positions(self): + """The policy must run at the positions ``F`` is evaluated at. + + Velocity-Verlet evaluates ``F`` at the *end-of-step* positions, so a + seam wired to step-start would leave the list one displacement behind + the positions entering ``F = -∇E`` — a systematic NVE energy leak + (surviving assertion of the deleted ``rebuild_every`` test). + """ + torch.manual_seed(0) + force = _CountingNLForce(1.0).to(_DTYPE) + ig = LangevinVerletIntegrator( + force, dt=0.01, gamma=0.0, kbt=0.0, mass=1.0, rebuild=True + ).cast_state(_DTYPE) + pos0 = torch.randn(4, 3, dtype=_DTYPE) + vel0 = torch.randn(4, 3, dtype=_DTYPE) * 0.1 + state = ig.initial(pos0, vel0) + # initial() → one force eval at pos0 + assert len(force.rebuild_positions) == 1 + assert torch.equal(force.rebuild_positions[0], pos0) + + state = ig.step_nve(state) + # second force eval at the *new* positions (not pos0) + assert len(force.rebuild_positions) == 2 + assert torch.equal(force.rebuild_positions[1], state.pos) + assert not torch.equal(force.rebuild_positions[1], pos0) + + def test_a_disabled_switch_never_touches_the_list(self): + """``rebuild=False`` over a list-backed force field freezes the list: + the policy clock must not even tick, or the compiled path would carry + the host sync it exists to avoid.""" + ff, pos = _lj_cut_force() + ig = LangevinVerletIntegrator( + ff, dt=0.5, gamma=0.0, kbt=0.0, mass=39.95, rebuild=False + ).cast_state(_DTYPE) + state = ig.initial(pos, torch.zeros_like(pos)) + for _ in range(3): + state = ig.step_nve(state) + assert ff.neighbors.rebuild_count == 0 + assert ff.neighbors.ago == 0 + + def test_the_modulo_counter_is_gone(self): + """One owner: the list. The integrator keeps no cadence state at all.""" + ig = _ig(1.0, dt=0.01, gamma=0.0, kbt=0.0, mass=1.0) + assert not hasattr(ig, "rebuild_every") + assert not hasattr(ig, "_force_eval_count") + + def test_a_conforming_subclass_derives_the_switch(self): + """``super().__init__(force)`` — the positional seam — still suffices.""" + ig = _MidpointEuler(HarmonicForceField(1.0).to(_DTYPE), dt=0.01) + assert ig.rebuild is False diff --git a/tests/test_molix/test_md/test_neighbor_graph_oracle.py b/tests/test_molix/test_md/test_neighbor_graph_oracle.py new file mode 100644 index 0000000..873323e --- /dev/null +++ b/tests/test_molix/test_md/test_neighbor_graph_oracle.py @@ -0,0 +1,340 @@ +"""NeighborList graph vs independent multi-image brute-force oracle. + +SUT: ``molix.md.NeighborList`` (the graph MACE consumes). Production cutoff +filter is ``0 < r <= r_c`` (see ``get_neighbor_pairs`` / binned path). + +Does **not** import production neighbor backends into the oracle module. +""" + +from __future__ import annotations + +from pathlib import Path + +import pytest +import torch + +from molix.md import NeighborList +from tests.test_molix.test_md.oracle_bruteforce_neighbors import ( + assert_graphs_equal, + bruteforce_edges, + neighborlist_edge_keys, + physical_dr_multiset, +) + +# --------------------------------------------------------------------------- +# helpers +# --------------------------------------------------------------------------- + + +def _list( + pos: torch.Tensor, + cell: torch.Tensor, + cutoff: float, + *, + skin: float = 0.0, + bin: float | None = None, +) -> NeighborList: + return NeighborList( + cell=cell, + cutoff=cutoff, + positions=pos, + skin=skin, + every=1, + delay=0, + check=True, + capacity_factor=2.0, + bin=bin, + ) + + +def _sut_keys(nl: NeighborList, cell: torch.Tensor) -> set: + return neighborlist_edge_keys(nl.edge_index, nl.shifts, nl.num_edges, cell, open_system=False) + + +def _compare(nl: NeighborList, pos: torch.Tensor, cell: torch.Tensor, cutoff: float, label: str): + ref = bruteforce_edges(pos, cell=cell, cutoff=cutoff, pbc=(True, True, True)) + sut = _sut_keys(nl, cell) + return assert_graphs_equal(ref, sut, label=label) + + +# --------------------------------------------------------------------------- +# ac-001 / oracle self-checks +# --------------------------------------------------------------------------- + + +class TestBruteforceOracle: + def test_excludes_only_true_self(self): + pos = torch.zeros(1, 3, dtype=torch.float64) + cell = torch.eye(3, dtype=torch.float64) * 5.0 + keys = bruteforce_edges(pos, cell=cell, cutoff=2.0, pbc=(True, True, True)) + assert (0, 0, 0, 0, 0) not in keys + # Self-image along ±e_x at distance L; need r_c >= L. + cell2 = torch.eye(3, dtype=torch.float64) * 3.0 + keys2 = bruteforce_edges(pos, cell=cell2, cutoff=3.0, pbc=(True, True, True)) + assert (0, 0, 1, 0, 0) in keys2 or (0, 0, -1, 0, 0) in keys2 + + def test_open_no_pbc(self): + pos = torch.tensor([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]], dtype=torch.float64) + keys = bruteforce_edges(pos, cell=None, cutoff=1.5, pbc=(False, False, False)) + assert keys == {(0, 1, 0, 0, 0), (1, 0, 0, 0, 0)} + + +# --------------------------------------------------------------------------- +# ac-002 open random (large cell ≈ open under MIC guard) +# --------------------------------------------------------------------------- + + +class TestOpenRandom: + def test_random_open_system_matches_oracle(self): + torch.manual_seed(0) + n = 20 + pos = torch.rand(n, 3, dtype=torch.float64) * 8.0 + 1.0 # stay away from edges + cell = torch.eye(3, dtype=torch.float64) * 20.0 # half-width 10 > r_c + cutoff = 3.5 + nl = _list(pos, cell, cutoff) + _compare(nl, pos, cell, cutoff, "open-random: ") + + +# --------------------------------------------------------------------------- +# ac-003 cutoff boundary ≤ +# --------------------------------------------------------------------------- + + +class TestCutoffBoundary: + @pytest.mark.parametrize( + "delta,expect_edge", + [ + (-1e-3, True), + (-1e-6, True), + (0.0, True), # r == r_c included + (1e-6, False), + (1e-3, False), + ], + ) + def test_edge_presence_at_cutoff(self, delta: float, expect_edge: bool): + r_c = 2.0 + # two atoms along x in a large box (no PBC contact) + pos = torch.tensor([[0.0, 5.0, 5.0], [r_c + delta, 5.0, 5.0]], dtype=torch.float64) + cell = torch.eye(3, dtype=torch.float64) * 20.0 + nl = _list(pos, cell, r_c) + keys = _sut_keys(nl, cell) + has = (0, 1, 0, 0, 0) in keys and (1, 0, 0, 0, 0) in keys + assert has is expect_edge, ( + f"cutoff convention is 0 < r <= r_c; at r=r_c+{delta:g} expected " + f"edge={expect_edge}, got has={has}; num_edges={nl.num_edges}" + ) + ref = bruteforce_edges(pos, cell=cell, cutoff=r_c) + assert ((0, 1, 0, 0, 0) in ref) is expect_edge + + +# --------------------------------------------------------------------------- +# ac-004 PBC wrap +# --------------------------------------------------------------------------- + + +class TestPeriodicWrap: + def test_atoms_across_box_face(self): + pos = torch.tensor([[0.1, 5.0, 5.0], [9.9, 5.0, 5.0]], dtype=torch.float64) + cell = torch.eye(3, dtype=torch.float64) * 10.0 + cutoff = 1.0 + nl = _list(pos, cell, cutoff) + cmp = _compare(nl, pos, cell, cutoff, "pbc-wrap: ") + assert cmp.n_ref >= 2 + # physical |dr| ≈ 0.2 + src, tgt = nl.edge_index[0, 0], nl.edge_index[0, 1] + dr = pos[int(tgt)] - pos[int(src)] + nl.shifts[0] + assert float(dr.norm()) == pytest.approx(0.2, abs=1e-9) + + +# --------------------------------------------------------------------------- +# ac-005a graph wrap invariance +# --------------------------------------------------------------------------- + + +class TestWrapInvarianceGraph: + def _base(self): + torch.manual_seed(1) + pos = torch.rand(12, 3, dtype=torch.float64) * 6.0 + 1.0 + cell = torch.eye(3, dtype=torch.float64) * 10.0 + cutoff = 2.5 + return pos, cell, cutoff + + def test_wrap_translate_physical_dr_multiset(self): + pos, cell, cutoff = self._base() + nl0 = _list(pos, cell, cutoff) + keys0 = _sut_keys(nl0, cell) + dr0 = physical_dr_multiset(keys0, pos, cell) + + # wrap atom 0 by +cell_x + pos_w = pos.clone() + pos_w[0] = pos_w[0] + cell[0] + nl_w = _list(pos_w, cell, cutoff) + keys_w = _sut_keys(nl_w, cell) + dr_w = physical_dr_multiset(keys_w, pos_w, cell) + assert dr0 == dr_w, "wrap one atom: physical dr multiset must match" + _compare(nl_w, pos_w, cell, cutoff, "wrap-atom: ") + + # translate whole structure by lattice vector + pos_t = pos + cell[1] + nl_t = _list(pos_t, cell, cutoff) + keys_t = _sut_keys(nl_t, cell) + dr_t = physical_dr_multiset(keys_t, pos_t, cell) + assert dr0 == dr_t + _compare(nl_t, pos_t, cell, cutoff, "translate-all: ") + + +# --------------------------------------------------------------------------- +# ac-005b E/F slow +# --------------------------------------------------------------------------- + + +def _matpes_weights() -> Path | None: + p = Path("/home/jicli594/work/mace_models") + if (p / "matpes_r2scan_config.json").is_file() and ( + p / "matpes_r2scan_cueq_state.pt" + ).is_file(): + return p + return None + + +@pytest.mark.slow +class TestWrapInvarianceEnergyForces: + def test_energy_forces_invariant_under_wrap(self): + weights = _matpes_weights() + if weights is None: + pytest.skip("MatPES weights not available") + from molix import config + from molpot.derivation.force import autograd_forces_from_energy + from molzoo.mace import MACEPotential + + config.set_precision("fp64") + model = ( + MACEPotential.from_checkpoint( + weights / "matpes_r2scan_config.json", + weights / "matpes_r2scan_cueq_state.pt", + use_fallback=True, + ) + .eval() + .to(dtype=torch.float64) + ) + r_c = float(model.cutoff_fn.r_cut) + # small water-like openish box: need L/2 >= r_c + L = 2.0 * r_c + 0.5 + torch.manual_seed(2) + n = 8 + # place atoms in interior + pos = (torch.rand(n, 3, dtype=torch.float64) * 0.4 + 0.3) * L + cell = torch.eye(3, dtype=torch.float64) * L + Z = torch.full((n,), 8, dtype=torch.long) # oxygen + + def ef(p: torch.Tensor) -> tuple[float, torch.Tensor]: + nl = _list(p, cell, r_c) + leaf = p.detach().requires_grad_(True) + batch = torch.zeros(n, dtype=torch.long) + with torch.enable_grad(): + e = model.energy_core( + leaf, + Z, + nl.edge_index[: nl.num_edges], + batch, + 1, + nl.shifts[: nl.num_edges], + ).sum() + f = autograd_forces_from_energy(e, leaf) + return float(e.detach()), f.detach() + + e0, f0 = ef(pos) + pos_w = pos.clone() + pos_w[0] = pos_w[0] + cell[0] + e1, f1 = ef(pos_w) + assert e0 == pytest.approx(e1, abs=1e-6, rel=1e-8) + # force on atom 0 after wrap: still same Cartesian F + assert torch.allclose(f0, f1, atol=1e-5, rtol=1e-6) + + +# --------------------------------------------------------------------------- +# ac-006 triclinic +# --------------------------------------------------------------------------- + + +class TestTriclinic: + def test_tilted_cell_matches_oracle(self): + torch.manual_seed(3) + pos = torch.rand(10, 3, dtype=torch.float64) * 4.0 + 0.5 + cell = torch.tensor( + [[8.0, 0.0, 0.0], [1.5, 7.5, 0.0], [0.5, 0.8, 7.0]], + dtype=torch.float64, + ) + cutoff = 2.0 + # ensure r_build ok + nl = _list(pos, cell, cutoff) + _compare(nl, pos, cell, cutoff, "triclinic: ") + + +# --------------------------------------------------------------------------- +# ac-007 multi-image domain (oracle multi-image + SUT under r_c <= L/2) +# --------------------------------------------------------------------------- + + +class TestMultiImage: + def test_oracle_finds_self_images_when_rc_exceeds_half_box(self): + pos = torch.zeros(1, 3, dtype=torch.float64) + cell = torch.eye(3, dtype=torch.float64) * 3.0 + keys = bruteforce_edges(pos, cell=cell, cutoff=3.0) + assert any(k[0] == 0 and k[1] == 0 and k[2:] != (0, 0, 0) for k in keys) + + def test_neighborlist_matches_oracle_near_half_width(self): + # NeighborList forbids r_build > L/2; stay just under. + L = 10.0 + cutoff = 4.9 # half-width = 5.0 + torch.manual_seed(4) + pos = torch.rand(15, 3, dtype=torch.float64) * (L - 1.0) + 0.5 + cell = torch.eye(3, dtype=torch.float64) * L + nl = _list(pos, cell, cutoff) + _compare(nl, pos, cell, cutoff, "near-half-width: ") + + +# --------------------------------------------------------------------------- +# ac-008 cutoff-crossing frames +# --------------------------------------------------------------------------- + + +class TestCutoffCrossing: + def test_per_frame_rebuild_no_hysteresis(self): + r_c = 2.0 + cell = torch.eye(3, dtype=torch.float64) * 20.0 + # approach from outside, enter, leave + distances = [2.5, 2.01, 1.99, 1.5, 1.99, 2.01, 2.5] + states = [] + for r in distances: + pos = torch.tensor([[0.0, 5.0, 5.0], [r, 5.0, 5.0]], dtype=torch.float64) + nl = _list(pos, cell, r_c) # fresh each frame + keys = _sut_keys(nl, cell) + has = (0, 1, 0, 0, 0) in keys + ref = bruteforce_edges(pos, cell=cell, cutoff=r_c) + assert_graphs_equal(ref, keys, label=f"crossing r={r}: ") + states.append(has) + # expected: outside False, inside True, no sticky True after leaving + assert states == [False, False, True, True, True, False, False] + + +# --------------------------------------------------------------------------- +# ac-012 dual backend +# --------------------------------------------------------------------------- + + +class TestDualBackend: + def test_bin_matches_default_on_wrap_case(self): + pos = torch.tensor([[0.1, 5.0, 5.0], [9.9, 5.0, 5.0]], dtype=torch.float64) + cell = torch.eye(3, dtype=torch.float64) * 10.0 + cutoff = 1.0 + nl_def = _list(pos, cell, cutoff, bin=None) + try: + nl_bin = _list(pos, cell, cutoff, bin=0.0) # auto bin size + except Exception as exc: + pytest.skip(f"binned backend unavailable: {exc}") + k_def = _sut_keys(nl_def, cell) + k_bin = _sut_keys(nl_bin, cell) + assert_graphs_equal(k_def, k_bin, label="bin-vs-default: ") + ref = bruteforce_edges(pos, cell=cell, cutoff=cutoff) + assert_graphs_equal(ref, k_bin, label="bin-vs-oracle: ") diff --git a/tests/test_molix/test_md/test_neighbors.py b/tests/test_molix/test_md/test_neighbors.py new file mode 100644 index 0000000..5887b3c --- /dev/null +++ b/tests/test_molix/test_md/test_neighbors.py @@ -0,0 +1,1780 @@ +"""Tests for molix.md.neighbors.""" + +import math +import re +from typing import NamedTuple + +import pytest +import torch +from tensordict import TensorDict + +import molix.data.tasks.neighbor +import molix.md +import molix.md.neighbors +from molix.md import ( + EV_PER_AMU_A2_FS2, + MD, + LennardJonesCutForceField, + MaxwellBoltzmann, + MDHook, + MDObservables, + MDRunner, +) +from molix.md.neighbors import NeighborList, NeighborStrategy +from tests.test_molix.test_md.conftest import make_cubic_lattice + + +def _triclinic() -> tuple[torch.Tensor, torch.Tensor]: + """The golden triclinic counterexample and ten atoms placed inside it. + + The cell ``[[10, 0, 0], [6, 8, 0], [0, 0, 10]]`` has ``V = 800 A^3`` and + perpendicular widths ``w_i = V / ||a_j x a_k|| = (8.0, 8.0, 10.0) A``, so + minimum-image completeness holds only up to ``min_i w_i / 2 = 4.000 A``. + Its shortest **row norm** is ``10.0 A``, which a row-norm guard reads as an + admissible cutoff of ``5.000 A`` — the gap between ``4.000`` and ``5.000`` + is the bug this cell pins. + + Positions are literal fractional coordinates mapped by ``frac @ cell`` (no + RNG): the closest minimum-image pairs sit at ``2.0 A`` and the next shell + at ``4.0 A``, so an accepted ``cutoff = 3.9 A`` yields a non-empty edge set + with no pair inside ``0.1 A`` of the cutoff. + + Returns: + ``(positions (10, 3), cell (3, 3))`` in Angstrom, ``float64``. + """ + cell = torch.tensor([[10.0, 0.0, 0.0], [6.0, 8.0, 0.0], [0.0, 0.0, 10.0]], dtype=torch.float64) + frac = torch.tensor( + [ + [0.05, 0.10, 0.10], + [0.25, 0.10, 0.10], + [0.45, 0.10, 0.10], + [0.65, 0.10, 0.10], + [0.05, 0.40, 0.45], + [0.25, 0.40, 0.45], + [0.45, 0.40, 0.45], + [0.05, 0.70, 0.80], + [0.25, 0.70, 0.80], + [0.45, 0.70, 0.80], + ], + dtype=torch.float64, + ) + return frac @ cell, cell + + +@pytest.fixture +def nlist(): + pos, cell = make_cubic_lattice() + return NeighborList(cell=cell, cutoff=3.5, positions=pos), pos + + +class TestNeighborList: + """Test the rebuilding fixed-capacity neighbour list.""" + + def test_renamed_symbol_is_the_md_export(self): + """The MD list is exported as ``NeighborList``; the old name is gone. + + There is no back-compat alias (``stage: experimental``, repo norm), so + the old name must be absent from the module *and* from ``__all__`` — a + ``PeriodicNeighborList = NeighborList`` shim would defeat this guard. + """ + from molix.md import NeighborList + + assert NeighborList is molix.md.neighbors.NeighborList + assert not hasattr(molix.md, "PeriodicNeighborList") + + names = list(molix.md.__all__) + assert "PeriodicNeighborList" not in names + assert "NeighborList" in names + # ``__all__`` stays alphabetized (notes, 2026-08-09) inside the + # PascalCase block that follows the CONSTANT_CASE block, so the entry + # sorts into the "N" run instead of keeping the old "P" slot. Pinned + # against ``NeighborStrategy`` (a permanent export) rather than the + # ``NeighborListHook`` this chain's link 07 deletes. + pascal = [name for name in names if not name.isupper()] + assert pascal == sorted(pascal) + assert names.index("NeighborList") < names.index("NeighborStrategy") + + def test_md_list_is_not_the_pipeline_task(self): + """Anti-shadow: two deliberate same-name types in different layers. + + ``molix.md.neighbors`` imports the pipeline ``SampleTask`` — the one + owner of kernel-output normalisation — into its own namespace. Once the + MD buffer owner takes the bare name, that import must be aliased to + ``NeighborListTask``, or the class definition rebinds the module-level + name and the constructor calls *itself* recursively. + """ + assert molix.md.NeighborList is not molix.data.tasks.neighbor.NeighborList + assert molix.md.neighbors.NeighborListTask is molix.data.tasks.neighbor.NeighborList + + def test_satisfies_the_neighbor_strategy_protocol(self, nlist): + nl, _ = nlist + assert isinstance(nl, NeighborStrategy) + + def test_capacity_exceeds_initial_edges(self, nlist): + """Headroom is what lets the edge count grow without reallocating.""" + nl, _ = nlist + assert nl.capacity > nl.num_edges + + def test_buffer_shapes_are_the_capacity(self, nlist): + """Shapes must be the capacity, not the live edge count — that is the + whole point: a CUDA graph sees constant shapes across rebuilds. The + edge buffer is ``(E, 2)`` per the repo-wide edge convention (``(2, N)`` + is reserved for ``bond_index`` as an anti-alias guard).""" + nl, _ = nlist + assert nl.edge_index.shape == (nl.capacity, 2) + assert nl.shifts.shape == (nl.capacity, 3) + + def test_rebuild_keeps_shapes_constant(self, nlist): + """Displacing every atom changes the edge set, never the shapes.""" + nl, pos = nlist + shapes = (nl.edge_index.shape, nl.shifts.shape) + torch.manual_seed(0) + nl.rebuild(pos + torch.randn_like(pos) * 0.3) + assert (nl.edge_index.shape, nl.shifts.shape) == shapes + + def test_rebuild_tracks_a_changed_neighbour_set(self, nlist): + """A real displacement must actually change the recorded edges.""" + nl, pos = nlist + before = nl.num_edges + nl.rebuild(pos * 1.15) # dilate: fewer pairs inside the cutoff + assert nl.num_edges != before + + def test_dead_edges_are_self_loops_on_atom_zero(self, nlist): + """Padding rows must not point at real atoms.""" + nl, _ = nlist + tail = nl.edge_index[nl.num_edges :] + assert torch.count_nonzero(tail) == 0 + + def test_dead_edge_shift_exceeds_the_cutoff(self, nlist): + """Beyond the cutoff every envelope is 0, which is what zeroes them.""" + nl, _ = nlist + tail = nl.shifts[nl.num_edges :] + assert bool((torch.linalg.norm(tail, dim=-1) > nl.cutoff).all()) + + def test_shifts_reconstruct_minimum_image_displacements(self, nlist): + """``pos[t] - pos[s] + shift`` must be the minimum-image vector, i.e. + no live edge is longer than the cutoff.""" + nl, pos = nlist + n = nl.num_edges + src, tgt = nl.edge_index[:n, 0], nl.edge_index[:n, 1] + vectors = pos[tgt] - pos[src] + nl.shifts[:n] + assert float(torch.linalg.norm(vectors, dim=-1).max()) <= nl.cutoff + 1e-9 + + def test_rebuild_count_increments(self, nlist): + """Diagnostics: callers need to know the cadence actually fired.""" + nl, pos = nlist + assert nl.rebuild_count == 0 + nl.rebuild(pos) + nl.rebuild(pos) + assert nl.rebuild_count == 2 + + def test_to_accepts_positional_dtype(self, nlist): + """``nl.to(torch.float32)`` must work — Tensor.to semantics, as documented.""" + nl, _ = nlist + out = nl.to(torch.float32) + assert out is nl + assert nl.shifts.dtype == torch.float32 + assert nl.cell.dtype == torch.float32 + assert nl.edge_index.dtype == torch.long # indices are never cast + + def test_rejects_cutoff_beyond_half_the_cell(self): + """Minimum image silently misses images past L/2 — refuse instead.""" + pos, cell = make_cubic_lattice() + with pytest.raises(ValueError, match="exceeds half the minimum perpendicular cell width"): + NeighborList(cell=cell, cutoff=5.0, positions=pos) + + def test_accepts_orthorhombic_cutoff_just_below_half_the_cell(self): + """Orthorhombic parity: for a cube ``w_i = ||a_i||``, so the bound stays + ``4.500 A`` on the 9 A cell and ``4.4 A`` must still build.""" + pos, cell = make_cubic_lattice() + nl = NeighborList(cell=cell, cutoff=4.4, positions=pos) + assert nl.cutoff == 4.4 + assert nl.num_edges > 0 + + def test_rejects_triclinic_cutoff_admitted_by_the_row_norm(self): + """The bug, asserted directly: the golden cell's shortest row norm is + ``10 A`` (row-norm bound ``5.000 A``) but its narrowest perpendicular + width is ``8 A`` (true bound ``4.000 A``), so ``5.0 A`` must be refused + instead of silently dropping pairs inside the cutoff.""" + pos, cell = _triclinic() + with pytest.raises(ValueError): + NeighborList(cell=cell, cutoff=5.0, positions=pos) + + def test_triclinic_rejection_names_the_perpendicular_bound(self): + """The measured bound must be observable through the public error, not + just the refusal: ``min_i w_i / 2 = 4.000 A`` for the golden cell.""" + pos, cell = _triclinic() + with pytest.raises(ValueError, match=r"4\.000 A"): + NeighborList(cell=cell, cutoff=4.5, positions=pos) + + def test_accepts_triclinic_cutoff_below_the_perpendicular_bound(self): + """Acceptance must mean "actually built", not "did not raise": below the + ``4.000 A`` bound the list constructs and reports real edges (closest + golden pairs are at ``2.0 A``).""" + pos, cell = _triclinic() + nl = NeighborList(cell=cell, cutoff=3.9, positions=pos) + assert nl.cutoff == 3.9 + assert nl.num_edges > 0 + + def test_rejects_a_singular_cell(self): + """A zero-volume cell has no finite width. ``V / area`` would be ``nan`` + and ``cutoff > nan`` is ``False`` — the guard must raise rather than let + a degenerate cell through the hole it opens.""" + cell = torch.tensor( + [[10.0, 0.0, 0.0], [10.0, 0.0, 0.0], [0.0, 0.0, 10.0]], dtype=torch.float64 + ) + pos = torch.tensor( + [[0.0, 0.0, 0.0], [2.0, 0.0, 0.0], [0.0, 2.0, 0.0], [0.0, 0.0, 2.0]], + dtype=torch.float64, + ) + with pytest.raises(ValueError): + NeighborList(cell=cell, cutoff=3.0, positions=pos) + + def test_rejects_a_non_3x3_cell(self): + """A batched ``(1, 3, 3)`` cell must fail with a clean ``ValueError`` at + the guard, not an opaque indexing error deeper in the kernel path.""" + pos, cell = make_cubic_lattice() + with pytest.raises(ValueError): + NeighborList(cell=cell.unsqueeze(0), cutoff=3.5, positions=pos) + + def test_overflow_raises_rather_than_truncating(self): + """A truncated neighbour list is a silently wrong energy.""" + pos, cell = make_cubic_lattice() + nl = NeighborList(cell=cell, cutoff=3.5, positions=pos, capacity_factor=1.0) + with pytest.raises(RuntimeError, match="overflow"): + nl.rebuild(pos * 0.5) # compress: many more pairs inside the cutoff + + +def _policy_list( + *, + skin: float = 0.0, + every: int = 1, + delay: int = 0, + check: bool = True, + cutoff: float = 3.5, + capacity_factor: float = 1.35, +) -> tuple[NeighborList, torch.Tensor]: + """A policy-configured list over the 4x4x4 lattice, plus its positions. + + The shared fixture of the policy suite: 64 atoms in a 12 A cube, whose + perpendicular half-width is 6.0 A and whose neighbour counts are exact + integers from crystallography (6 at 3.0 A, 12 at 4.2426 A, 8 at 5.196 A). + """ + pos, cell = make_cubic_lattice(n_side=4, spacing=3.0) + nl = NeighborList( + cell=cell, + cutoff=cutoff, + positions=pos, + skin=skin, + every=every, + delay=delay, + check=check, + capacity_factor=capacity_factor, + ) + return nl, pos + + +def _displaced(pos: torch.Tensor, distance: float, *, atom: int = 0) -> torch.Tensor: + """``pos`` with one atom translated ``distance`` A along x. + + Raw and unwrapped, as an MD trajectory drifts: the displacement test the + policy runs is a raw difference, never a minimum image. + """ + moved = pos.clone() + moved[atom, 0] += distance + return moved + + +class _Frame(NamedTuple): + """One observed step: the evaluated configuration and the list behind it.""" + + pos: torch.Tensor + edge_index: torch.Tensor + shifts: torch.Tensor + total: torch.Tensor + forces: torch.Tensor + + +class _FrameRecorder(MDHook): + """Capture ``obs.pos`` together with the live list it was evaluated against.""" + + def __init__(self, neighbors: NeighborList) -> None: + self._neighbors = neighbors + self.frames: list[_Frame] = [] + + def on_step_end(self, runner: MDRunner, step: int, obs: MDObservables) -> None: + """Snapshot the step; the live edges are the ones ``obs.forces`` used.""" + n = self._neighbors.num_edges + self.frames.append( + _Frame( + pos=obs.pos.detach().clone(), + edge_index=self._neighbors.edge_index[:n].clone(), + shifts=self._neighbors.shifts[:n].clone(), + total=obs.total.detach().clone(), + forces=obs.forces.detach().clone(), + ) + ) + + +def _run_lj_lattice( + *, skin: float, every: int = 1, delay: int = 0, check: bool = True, n_steps: int = 100 +) -> tuple[NeighborList, list[_Frame]]: + """NVE argon over the 64-atom lattice, driving the policy once per force eval. + + Deterministic CPU float64 throughout: seeded Maxwell-Boltzmann velocities, + gamma = 0, no wall clock, no filesystem, no network. Argon in (amu, A, fs): + eps = 0.0103 eV, sigma = 2.5 A, cutoff = 3.5 A, m = 39.95 amu, dt = 4 fs. + + No cadence knob is passed: link 07 makes the *force field* declare that it + owns a live list (``LennardJonesCutForceField.rebuilds_neighbors``), the + integrator derive its static switch from that, and ``rebuild_neighbors`` + land on :meth:`NeighborList.update` — so the policy runs once per force + evaluation, at the positions being evaluated, with no driver kwarg and no + ``_PolicyForceField`` preview subclass in the way. + + ``capacity_factor=2.5`` is measured, not defensive: at ``skin=0.5`` this run + reaches 600 live edges against the 519 rows the default 1.35 would allocate + from the initial 384, and the overflow guard is not what these tests pin. + """ + pos, cell = make_cubic_lattice(n_side=4, spacing=3.0) + neighbors = NeighborList( + cell=cell, + cutoff=3.5, + positions=pos, + skin=skin, + every=every, + delay=delay, + check=check, + capacity_factor=2.5, + ) + force = LennardJonesCutForceField( + epsilon=0.0103 / EV_PER_AMU_A2_FS2, # argon well depth, eV -> amu A^2/fs^2 + sigma=2.5, + neighbors=neighbors, + cutoff=3.5, + ) + recorder = _FrameRecorder(neighbors) + velocities = MaxwellBoltzmann(39.95, n_atoms=64).sample(300.0, seed=0) + md = MD( + force, + mass=39.95, + dt=4.0, + gamma=0.0, + dtype=torch.float64, + hooks=[recorder], + ) + md.set_potential_dtype(torch.float64) + md.run(pos, velocities, n_steps, chunk=1) + return neighbors, recorder.frames + + +def _reference_pairs( + pos: torch.Tensor, cell: torch.Tensor, cutoff: float +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """The exact O(N^2) minimum-image pair set within ``cutoff`` — the oracle. + + Computed from fractional coordinates with ``round`` to the nearest image + (exact for a cubic cell), so it holds for the unwrapped positions an MD run + drifts into. Shares no code with the neighbour list under test. + + Returns: + ``(source, target, distance)`` for every **ordered** pair inside + ``cutoff``, with ``distance = ||pos[target] - pos[source] + shift||``. + """ + fractional = pos @ torch.linalg.inv(cell) + delta = fractional.unsqueeze(0) - fractional.unsqueeze(1) # [i, j] = frac[j] - frac[i] + delta = delta - torch.round(delta) + distance = torch.linalg.norm(delta @ cell, dim=-1) + inside = (distance < cutoff) & ~torch.eye(pos.shape[0], dtype=torch.bool) + source, target = torch.nonzero(inside, as_tuple=True) + return source, target, distance[source, target] + + +class TestNeighborListPolicy: + """Verlet skin + LAMMPS ``neigh_modify every/delay/check`` rebuild policy. + + Reference: + ``lammps/lammps`` develop, ``src/neighbor.cpp`` — ``Neighbor::decide`` + (2408-2424), ``Neighbor::check_distance`` (2438-2490), ``Neighbor::init``. + K. Nordlund, *Introduction to molecular dynamics simulations*, lecture 3, + for the half-skin two-atom criterion. + """ + + # --- construction contract: r_build derivation, sizing, guards ---------- + + def test_build_radius_is_the_cutoff_plus_the_skin(self): + """``cutoff`` stays the *interaction* cutoff every consumer means, and + ``r_build = cutoff + skin`` is the derived build radius.""" + nl, _ = _policy_list(skin=1.5) + assert nl.cutoff == 3.5 + assert nl.skin == 1.5 + assert nl.r_build == pytest.approx(5.0, abs=1e-12) + + def test_build_radius_is_not_settable(self): + """Derived, never assigned: a writable ``r_build`` could drift out of + step with the capacity, the kernel radius and the half-width guard that + were all sized from it at construction.""" + nl, _ = _policy_list(skin=1.5) + with pytest.raises(AttributeError): + setattr(nl, "r_build", 6.0) + + def test_skin_extends_the_build_to_the_second_neighbour_shell(self): + """Crystallography, not a fit: at ``r_build = 5.0 A`` every atom of the + simple-cubic lattice sees 6 neighbours at 3.0 A and 12 at 3*sqrt(2) = + 4.2426 A (the 8 body-diagonal ones at 5.196 A stay out), so the + bidirectional list holds 64 * 18 = 1152 edges.""" + nl, _ = _policy_list(skin=1.5) + assert nl.num_edges == 1152 + + def test_zero_skin_builds_only_the_first_shell(self): + """The same lattice at ``r_build = cutoff = 3.5 A``: 6 neighbours each, + 64 * 6 = 384 edges. The 1152/384 = 3.0x ratio is the direct measurement + of the ``(1 + s/r_cut)^3`` growth the capacity has to absorb.""" + nl, _ = _policy_list(skin=0.0) + assert nl.num_edges == 384 + + def test_capacity_is_sized_from_the_build_radius(self): + """Sizing from ``cutoff`` instead would allocate 519 rows for a list + that starts with 1152 live edges — an overflow on the constructor's own + build, before a single MD step.""" + nl, _ = _policy_list(skin=1.5) + assert nl.capacity >= math.ceil(1.35 * 1152) + + def test_buffer_shapes_stay_the_capacity_under_a_skin(self): + """The skin must not cost the constant shapes a CUDA graph needs.""" + nl, _ = _policy_list(skin=1.5) + assert nl.edge_index.shape == (nl.capacity, 2) + assert nl.shifts.shape == (nl.capacity, 3) + + def test_a_fresh_list_reports_no_rebuild_history(self): + """Defaults are the pre-skin behaviour: ``skin=0``, and the + constructor's initial build is not counted as a rebuild.""" + pos, cell = make_cubic_lattice(n_side=4, spacing=3.0) + nl = NeighborList(cell=cell, cutoff=3.5, positions=pos) + assert (nl.skin, nl.ago, nl.rebuild_count, nl.ndanger) == (0.0, 0, 0, 0) + + def test_to_casts_the_displacement_reference(self): + """``_x_hold`` must follow ``shifts`` / ``cell`` through a cast. + + Asserted on the buffer dtype rather than through behaviour on purpose: + a float64 ``_x_hold`` differenced against float32 positions *promotes + silently* — no error, just a mixed-precision comparison nobody asked + for — so the dtype is the only observable. + """ + nl, pos = _policy_list(skin=1.5) + nl.to(torch.float32) + assert nl._x_hold.dtype == torch.float32 + assert nl.update(_displaced(pos.to(torch.float32), 1.0)) is True + + def test_rejects_a_skin_that_pushes_the_build_past_half_the_cell(self): + """The link-01 half-width guard, re-derived on ``r_build``: ``cutoff = + 3.5 A`` alone is admissible in the 12 A cube (half-width 6.0 A), but + ``skin = 3.0`` makes ``r_build = 6.5 A``, past which the kernel's + minimum-image reduction silently drops pairs inside the cutoff.""" + pos, cell = make_cubic_lattice(n_side=4, spacing=3.0) + NeighborList(cell=cell, cutoff=3.5, positions=pos) # the cutoff alone passes + with pytest.raises(ValueError, match="r_build"): + NeighborList(cell=cell, cutoff=3.5, positions=pos, skin=3.0) + + def test_rejects_a_skin_that_reaches_the_dead_edge_shift(self): + """Dead padding edges sit at ``DEAD_EDGE_CUTOFF_FACTOR * cutoff = 10x`` + the cutoff; a skin of ``9x`` puts them exactly on the build radius, + where they stop being inert and start being counted as real pairs. The + 60 A cell keeps the half-width guard — checked first — out of the way, + so this pins the dead-edge assertion specifically.""" + pos, _ = make_cubic_lattice(n_side=4, spacing=3.0) + cell = torch.eye(3, dtype=torch.float64) * 60.0 + with pytest.raises(ValueError, match="dead"): + NeighborList(cell=cell, cutoff=1.0, positions=pos, skin=9.0) + + @pytest.mark.parametrize( + ("skin", "every", "delay"), + [(-0.1, 1, 0), (0.0, 0, 0), (0.0, -1, 0), (0.0, 1, -1)], + ids=["negative-skin", "zero-every", "negative-every", "negative-delay"], + ) + def test_rejects_out_of_domain_policy_parameters(self, skin: float, every: int, delay: int): + """Each arm has a domain; a silently clamped one disables the gate.""" + pos, cell = make_cubic_lattice(n_side=4, spacing=3.0) + with pytest.raises(ValueError): + NeighborList(cell=cell, cutoff=3.5, positions=pos, skin=skin, every=every, delay=delay) + + def test_rejects_a_delay_that_is_not_a_multiple_of_every(self): + """LAMMPS ``Neighbor::init`` parity: with ``every=4, delay=10`` the + danger threshold ``max(every, delay) = 10`` is an ``ago`` the gate never + permits (10 % 4 != 0), so ``ndanger`` could never fire and the + correctness alarm would be silently dead.""" + pos, cell = make_cubic_lattice(n_side=4, spacing=3.0) + with pytest.raises(ValueError, match="multiple"): + NeighborList(cell=cell, cutoff=3.5, positions=pos, every=4, delay=10) + + def test_accepts_a_delay_that_is_a_multiple_of_every(self): + """The guard rejects a non-multiple, not a coarse gate as such.""" + nl, _ = _policy_list(every=4, delay=8) + assert (nl.every, nl.delay, nl.ago) == (4, 8, 0) + + def test_accepts_zero_delay_under_any_every(self): + """``0 % every == 0`` for every legal ``every``, so the default + ``delay=0`` is always LAMMPS-legal — the guard must not read a falsy + delay as unset and reject it.""" + nl, _ = _policy_list(every=7, delay=0) + assert (nl.every, nl.delay) == (7, 0) + + # --- gating arithmetic (LAMMPS ``Neighbor::decide``) -------------------- + + def test_the_gate_is_conjunctive_and_first_permits_ago_eight(self): + """``every=4, delay=8``: both arms must agree. A displacement four times + the half-skin cannot pull a rebuild forward — ``ago = 4`` clears + ``ago % every`` but not ``ago >= delay``, and nothing before 8 clears + both.""" + nl, pos = _policy_list(skin=1.0, every=4, delay=8) + moved = _displaced(pos, 2.0) + assert [nl.update(moved) for _ in range(8)] == [False] * 7 + [True] + assert (nl.rebuild_count, nl.ago) == (1, 0) + + def test_the_every_arm_alone_first_permits_ago_four(self): + """``delay=0`` leaves the cadence arm as the only gate.""" + nl, pos = _policy_list(skin=1.0, every=4, delay=0) + moved = _displaced(pos, 2.0) + assert [nl.update(moved) for _ in range(4)] == [False, False, False, True] + + def test_the_delay_arm_alone_first_permits_ago_ten(self): + """``every=1`` leaves the delay arm as the only gate.""" + nl, pos = _policy_list(skin=1.0, every=1, delay=10) + moved = _displaced(pos, 2.0) + assert [nl.update(moved) for _ in range(10)] == [False] * 9 + [True] + + def test_a_forced_rebuild_rephases_the_schedule(self): + """``ago`` counts from the last build, not from an absolute step index: + after two blocked updates a forced ``rebuild`` restarts the clock, so + the next permitted opportunity is four updates later, not two.""" + nl, pos = _policy_list(skin=1.0, every=4, delay=0) + moved = _displaced(pos, 2.0) + assert [nl.update(moved) for _ in range(2)] == [False, False] + nl.rebuild(pos) + assert nl.ago == 0 + assert [nl.update(moved) for _ in range(4)] == [False, False, False, True] + + def test_a_displacement_of_exactly_half_the_skin_does_not_rebuild(self): + """LAMMPS compares with a strict ``>``: at ``max_i d_i == s/2`` the + two-atom criterion ``d_(1) + d_(2) <= s`` still holds, so the list is + still provably complete and a rebuild would be wasted work.""" + nl, pos = _policy_list(skin=1.0) + assert nl.update(_displaced(pos, 0.5)) is False + assert nl.rebuild_count == 0 + + def test_a_displacement_just_past_half_the_skin_rebuilds(self): + """The other side of the same strict comparison, one ulp-scale step + away: 0.5 + 1e-9 A is no longer covered by the completeness proof.""" + nl, pos = _policy_list(skin=1.0) + assert nl.update(_displaced(pos, 0.5 + 1e-9)) is True + assert nl.rebuild_count == 1 + + def test_zero_skin_holds_the_list_when_nothing_moved(self): + """The degenerate limit is not "rebuild unconditionally": with + ``max_d2 == 0`` and a strict ``>``, the list is already exactly right, + so the bookkeeping says so.""" + nl, pos = _policy_list(skin=0.0) + assert nl.update(pos.clone()) is False + assert nl.rebuild_count == 0 + + def test_zero_skin_rebuilds_on_any_motion(self): + """``half_skin_sq = 0`` reproduces today's rebuild-every-ask behaviour + edge for edge: any nonzero displacement rebuilds.""" + nl, pos = _policy_list(skin=0.0) + assert nl.update(_displaced(pos, 1e-6)) is True + assert nl.rebuild_count == 1 + + def test_an_unchecked_policy_rebuilds_on_cadence_alone(self): + """``check=False`` buys speed by dropping the distance test entirely: + every permitted opportunity rebuilds even though nothing has moved.""" + nl, pos = _policy_list(skin=1.0, every=3, check=False) + frozen = pos.clone() + assert [nl.update(frozen) for _ in range(6)] == [False, False, True] * 2 + assert nl.rebuild_count == 2 + + def test_ndanger_counts_a_rebuild_at_the_first_permitted_opportunity(self): + """``every=1, delay=0`` puts ``_danger_ago`` at 1, so a rebuild that + fires immediately may already have been overdue on the step before — + which is exactly what the counter is for.""" + nl, pos = _policy_list(skin=1.0) + assert nl.update(_displaced(pos, 1.0)) is True + assert nl.ndanger == 1 + assert nl.update(_displaced(pos, 2.0)) is True + assert nl.ndanger == 2 + + def test_ndanger_stays_zero_when_the_rebuild_was_not_overdue(self): + """``every=2, delay=6`` puts ``_danger_ago`` at 6. The half-skin is + crossed only *after* that first opportunity, so the rebuild lands at + ``ago = 8`` (7 % 2 != 0) and nothing was missed.""" + nl, pos = _policy_list(skin=1.0, every=2, delay=6) + small, large = _displaced(pos, 0.2), _displaced(pos, 1.0) + assert [nl.update(small) for _ in range(6)] == [False] * 6 + assert [nl.update(large) for _ in range(2)] == [False, True] + assert (nl.rebuild_count, nl.ndanger) == (1, 0) + + def test_ndanger_fires_at_the_lammps_threshold(self): + """The same coarse gate, but the half-skin is already crossed at the + first permitted opportunity: the rebuild lands at ``ago = 6 = + max(every, delay)``, ``neighbor.cpp:2488`` verbatim, on a LAMMPS-legal + configuration.""" + nl, pos = _policy_list(skin=1.0, every=2, delay=6) + moved = _displaced(pos, 1.0) + assert [nl.update(moved) for _ in range(6)] == [False] * 5 + [True] + assert (nl.rebuild_count, nl.ndanger) == (1, 1) + + def test_a_full_cell_translation_is_refused_as_unwrapped(self): + """Frozen shifts plus a raw displacement test are correct only while + positions stay unwrapped. A 12 A jump — one full cell vector, past the + 6.0 A half-width — means mid-run wrapping, a changed cell, or a blown-up + trajectory; all three are fatal and none is recoverable.""" + nl, pos = _policy_list(skin=1.0) + with pytest.raises(RuntimeError, match="unwrapped"): + nl.update(_displaced(pos, 12.0)) + + def test_an_unchecked_policy_skips_the_unwrapped_guard(self): + """Documented consequence, pinned so the trade stays visible: the guard + lives inside the displacement branch, so ``check=False`` switches off + the invariant alarm along with the criterion.""" + nl, pos = _policy_list(skin=1.0, check=False) + assert nl.update(_displaced(pos, 12.0)) is True + + # --- protocol ----------------------------------------------------------- + + def test_a_skinned_list_satisfies_the_widened_protocol(self): + """Widening ``NeighborStrategy`` must not push its own implementation + out of the contract.""" + nl, _ = _policy_list(skin=1.5) + assert isinstance(nl, NeighborStrategy) + + def test_a_list_without_the_policy_members_fails_the_protocol(self): + """The widening is what lets ``forcefield.py`` drop its + ``getattr(neighbors, "cutoff", None)`` duck-read, so the protocol has to + actually *require* ``cutoff`` / ``skin`` / ``update``: a stub carrying + only the buffer members must no longer pass.""" + + class _BufferOnly: + edge_index: torch.Tensor = torch.zeros(1, 2, dtype=torch.long) + shifts: torch.Tensor = torch.zeros(1, 3) + num_edges: int = 0 + capacity: int = 1 + + def rebuild(self, positions: torch.Tensor) -> None: + """No-op build.""" + + def to( + self, + device: torch.device | str | torch.dtype | None = None, + dtype: torch.dtype | None = None, + ) -> "_BufferOnly": + return self + + assert not isinstance(_BufferOnly(), NeighborStrategy) + + def test_the_interaction_cutoff_is_the_bar_not_the_build_radius(self): + """A skinned list holds pairs out to ``r_build = 5.0 A``, but they are + only *complete* out to ``cutoff`` between rebuilds — so a 5.0 A + interaction cutoff over this list is precisely the silent truncation + that check exists to prevent.""" + nl, _ = _policy_list(skin=1.5) + with pytest.raises(ValueError): + LennardJonesCutForceField(epsilon=1.0, sigma=2.5, neighbors=nl, cutoff=5.0) + + # --- trajectory falsification (NVE argon over the lattice) ------------- + + def test_every_pair_within_the_cutoff_stays_in_the_live_list(self): + """PRIMARY falsification: a gated list must never miss a pair. + + At every step the exact minimum-image O(N^2) pair set within the + interaction cutoff — recomputed from the very configuration the forces + were evaluated at — must be a subset of the live list, and each pair's + list-reconstructed distance ``||pos[t] - pos[s] + shift||`` must match + the reference. The distance half is not redundant: a stale *shift* on a + surviving index pair is the frozen-shift failure mode, and an + index-subset check alone would sail straight past it. + """ + _, cell = make_cubic_lattice(n_side=4, spacing=3.0) + _, frames = _run_lj_lattice(skin=0.5) + n_atoms = 64 + for step, frame in enumerate(frames): + source, target, reference = _reference_pairs(frame.pos, cell, 3.5) + rows = torch.full((n_atoms * n_atoms,), -1, dtype=torch.long) + live = frame.edge_index + rows[live[:, 0] * n_atoms + live[:, 1]] = torch.arange(live.shape[0]) + found = rows[source * n_atoms + target] + missing = int((found < 0).sum()) + assert missing == 0, f"step {step}: {missing} pairs inside the cutoff are not listed" + reconstructed = torch.linalg.norm( + frame.pos[target] - frame.pos[source] + frame.shifts[found], dim=-1 + ) + torch.testing.assert_close(reconstructed, reference, atol=1e-9, rtol=0) + + def test_a_skinned_policy_matches_rebuilding_every_force_evaluation(self): + """Policy equivalence: the gated ``skin=0.5`` run and a ``skin=0.0`` run + that rebuilds at every force evaluation must trace the same physics. + Compared at ``atol=1e-10, rtol=0`` rather than bitwise on purpose — the + masked-zero skin edges reorder the ``index_add_`` accumulation.""" + _, gated = _run_lj_lattice(skin=0.5) + _, every_eval = _run_lj_lattice(skin=0.0) + assert len(gated) == len(every_eval) == 100 + for step, (a, b) in enumerate(zip(gated, every_eval, strict=True)): + torch.testing.assert_close(a.total, b.total, atol=1e-10, rtol=0, msg=f"step {step}") + torch.testing.assert_close(a.forces, b.forces, atol=1e-10, rtol=0, msg=f"step {step}") + + def test_the_standard_run_reports_no_dangerous_builds(self): + """``skin=0.5`` gives a half-skin of 0.25 A against ~0.01 A of motion + per step, so no rebuild is ever overdue. ``ndanger`` is the cheapest + correctness alarm available and must stay silent on a sane run — while + the gate itself stays alive (``rebuild_count > 0``).""" + neighbors, _ = _run_lj_lattice(skin=0.5) + assert neighbors.ndanger == 0 + assert neighbors.rebuild_count > 0 + + def test_rebuild_count_falls_as_the_skin_grows(self): + """The point of the skin, measured over the same trajectory. The strict + inequality at the ends is what catches a dead or inverted gate — plain + non-increasing monotonicity is also satisfied by a policy that never + rebuilds at all.""" + counts = [_run_lj_lattice(skin=skin)[0].rebuild_count for skin in (0.0, 0.25, 0.5, 1.0)] + assert counts == sorted(counts, reverse=True) + assert counts[-1] < counts[0] + + +def _batch(pos: torch.Tensor, cell: torch.Tensor | None = None) -> TensorDict: + """A single-system MD working batch over ``pos``, optionally carrying a cell. + + The same shape ``_periodic_template`` builds in ``test_forcefield.py`` — the + batch :class:`~molix.md.forcefield.PeriodicPotentialForceField` binds into — + plus the optional ``("graphs", "cell")`` the bind path validates. + + ``graphs`` carries ``batch_size=[]`` on purpose: the container must not be + the thing that decides the cell's leading dimension, or ``(3, 3)`` and + ``(1, 3, 3)`` could not both be stored and the accept/refuse decision under + test would be made by the fixture instead of by ``build``. + + Args: + pos: Positions ``(N, 3)`` in Angstrom. + cell: Optional cell vectors in Angstrom, any shape — the point of + several of these tests is that ``build`` judges the shape. + + Returns: + A batch ``TensorDict`` with root ``batch_size=[]``. + """ + n = int(pos.shape[0]) + data: dict[str, TensorDict] = { + "atoms": TensorDict( + { + "pos": pos, + "Z": torch.ones(n, dtype=torch.long), + "batch": torch.zeros(n, dtype=torch.long), + }, + batch_size=[n], + ) + } + if cell is not None: + data["graphs"] = TensorDict({"cell": cell}, batch_size=[]) + return TensorDict(data, batch_size=[]) + + +class TestNeighborListBind: + """``build(batch)`` / ``update(batch)`` — the TensorDict side of the list. + + The list owns ``edges`` once bound: ``build`` refreshes the buffers at + ``batch["atoms", "pos"]``, validates that the batch describes the system the + list was constructed for, writes the *live* buffers into ``batch["edges"]`` + **by reference** and hands the same batch back so it composes with the + repo-wide ``forward(td) -> td`` convention. + """ + + # --- happy path: bind ------------------------------------------------- + + def test_build_returns_the_argument_itself(self): + """``potential(nl.build(batch))`` only composes if the return is the + argument — a copy would leave the caller holding an unbound batch.""" + nl, pos = _policy_list(skin=1.5, capacity_factor=4.0) + batch = _batch(pos) + assert nl.build(batch) is batch + + def test_build_binds_the_live_buffers_by_reference(self): + """The load-bearing property of the whole link: **identity**, not equality. + + A container that copied on assignment would give the potential a + snapshot, so every later in-place rebuild would be invisible and the PES + would silently freeze. This assertion is that alarm. + """ + nl, pos = _policy_list(skin=1.5, capacity_factor=4.0) + batch = nl.build(_batch(pos)) + assert batch["edges", "edge_index"] is nl.edge_index + assert batch["edges", "shifts"] is nl.shifts + + def test_bound_edges_span_the_capacity_not_the_live_count(self): + """Fixed-capacity buffers go in whole: the tail of dead padding edges is + part of the contract (constant shapes for a CUDA graph), so the bound + namespace is sized by ``capacity``, never by ``num_edges``.""" + nl, pos = _policy_list(skin=1.5, capacity_factor=4.0) + batch = nl.build(_batch(pos)) + assert batch["edges"].batch_size == torch.Size([nl.capacity]) + + def test_bound_edges_hold_exactly_the_index_and_the_shifts(self): + """The bind writes the two keys the periodic potentials read and nothing + else — a derived ``edge_diff`` / ``edge_dist`` written here would be a + stale straight-through value the moment an atom moves.""" + nl, pos = _policy_list(skin=1.5, capacity_factor=4.0) + batch = nl.build(_batch(pos)) + assert set(batch["edges"].keys()) == {"edge_index", "shifts"} + + def test_build_is_never_counted_as_a_rebuild(self): + """``rebuild_count`` means "rebuilds driven during the run". ``build`` is + a binding operation — it also runs on every ``.to()`` re-sync — so + counting it would make a dtype cast look like physics.""" + nl, pos = _policy_list(skin=1.5, capacity_factor=4.0) + assert nl.rebuild_count == 0 + nl.build(_batch(pos)) + assert nl.rebuild_count == 0 + nl.rebuild(pos) + nl.build(_batch(pos)) + assert nl.rebuild_count == 1 + + def test_build_restarts_the_policy_clock(self): + """The buffers are fresh after a bind, so the ``ago`` clock the + ``every`` / ``delay`` gate runs on has to start there too — otherwise the + next ``update`` measures a staleness that was just eliminated.""" + nl, pos = _policy_list(skin=1.0, every=4, capacity_factor=4.0) + frozen = pos.clone() + assert [nl.update(frozen) for _ in range(2)] == [False, False] + assert nl.ago == 2 + nl.build(_batch(pos)) + assert nl.ago == 0 + + def test_build_uses_the_positions_in_the_batch(self): + """The bind is also a build: it must land on the batch's geometry, not + re-emit the constructor's. Measured against a list constructed directly + at the dilated lattice — the same kernel, so any difference is the bind + having built at the wrong positions.""" + pos, cell = make_cubic_lattice(n_side=4, spacing=3.0) + nl, _ = _policy_list(skin=1.5, capacity_factor=4.0) + assert nl.num_edges == 1152 # crystallography, link 04 + dilated = pos * 1.05 + reference = NeighborList( + cell=cell, cutoff=3.5, positions=dilated, skin=1.5, capacity_factor=4.0 + ) + nl.build(_batch(dilated)) + assert reference.num_edges != 1152 + assert nl.num_edges == reference.num_edges + + # --- liveness of the tie ---------------------------------------------- + + def test_a_forced_rebuild_is_visible_through_the_bound_batch(self): + """An in-place rebuild must reach the batch with no re-binding at all. + + The value read is not redundant with the identity check: a stale + *shift* on a surviving index pair is the frozen-shift failure mode, and + it is invisible to an index comparison alone. + """ + pos, cell = make_cubic_lattice(n_side=4, spacing=3.0) + nl, _ = _policy_list(skin=0.0, capacity_factor=4.0) + batch = nl.build(_batch(pos)) + before = nl.num_edges + compressed = pos * 0.8 # second-neighbour pairs enter the cutoff + nl.rebuild(compressed) + reference = NeighborList(cell=cell, cutoff=3.5, positions=compressed, capacity_factor=4.0) + assert nl.num_edges != before + assert nl.num_edges == reference.num_edges + assert batch["edges", "edge_index"] is nl.edge_index + assert batch["edges", "shifts"] is nl.shifts + n = nl.num_edges + assert torch.equal(batch["edges", "edge_index"][:n], reference.edge_index[:n]) + assert torch.equal(batch["edges", "shifts"][:n], reference.shifts[:n]) + + def test_a_policy_rebuild_is_visible_through_the_bound_batch(self): + """Same liveness, driven through the batch entry point instead: the + per-step idiom is ``nl.update(batch)``, so that is the path that has to + keep the tie alive.""" + pos, cell = make_cubic_lattice(n_side=4, spacing=3.0) + nl, _ = _policy_list(skin=0.0, capacity_factor=4.0) + batch = nl.build(_batch(pos)) + before = nl.num_edges + compressed = pos * 0.8 + assert nl.update(_batch(compressed)) is True + reference = NeighborList(cell=cell, cutoff=3.5, positions=compressed, capacity_factor=4.0) + assert nl.num_edges != before + assert nl.num_edges == reference.num_edges + assert batch["edges", "edge_index"] is nl.edge_index + assert batch["edges", "shifts"] is nl.shifts + n = nl.num_edges + assert torch.equal(batch["edges", "edge_index"][:n], reference.edge_index[:n]) + assert torch.equal(batch["edges", "shifts"][:n], reference.shifts[:n]) + + def test_a_bare_cast_severs_the_tie(self): + """The documented consequence, pinned so the trade stays visible. + + ``to(dtype)`` rebinds ``shifts`` to a new tensor, so a batch bound + beforehand keeps pointing at the old one. Auto-re-binding from inside + ``to()`` was rejected — it would make the list hold a reference to a + batch it does not own — so the **owner** re-binds; the force-field half + of this trade is pinned in ``test_forcefield.py``. + """ + nl, pos = _policy_list(skin=1.5, capacity_factor=4.0) + batch = nl.build(_batch(pos)) + nl.to(torch.float32) + assert batch["edges", "shifts"] is not nl.shifts + assert batch["edges", "shifts"].dtype == torch.float64 + assert nl.shifts.dtype == torch.float32 + + # --- the list owns ``edges`` once bound -------------------------------- + + def test_build_replaces_a_stale_edges_namespace_wholesale(self): + """Whatever was under ``edges`` is dropped, not merged. + + A precomputed ``edge_diff`` / ``edge_dist`` pair — exactly what + ``PotentialForceField._STALE_EDGE_KEYS`` strips at construction — would + be used straight through by a potential and freeze the PES, and a + surviving shorter ``edge_index`` would disagree with the capacity. + """ + nl, pos = _policy_list(skin=1.5, capacity_factor=4.0) + batch = _batch(pos) + batch["edges"] = TensorDict( + { + "edge_index": torch.zeros(3, 2, dtype=torch.long), + "edge_diff": torch.zeros(3, 3, dtype=torch.float64), + "edge_dist": torch.zeros(3, dtype=torch.float64), + }, + batch_size=[3], + ) + nl.build(batch) + assert set(batch["edges"].keys()) == {"edge_index", "shifts"} + assert batch["edges"].batch_size == torch.Size([nl.capacity]) + + # --- validation: the batch must describe *this* system ----------------- + + def test_build_without_positions_names_the_missing_key(self): + """``KeyError(('atoms', 'pos'))`` from three frames deep is the classic + two-tier data-contract confusion; the bind names the key it wanted.""" + nl, pos = _policy_list(skin=1.5, capacity_factor=4.0) + n = int(pos.shape[0]) + batch = TensorDict( + {"atoms": TensorDict({"Z": torch.ones(n, dtype=torch.long)}, batch_size=[n])}, + batch_size=[], + ) + with pytest.raises(ValueError, match="pos") as excinfo: + nl.build(batch) + assert "atoms" in str(excinfo.value) + + def test_build_rejects_a_different_atom_count(self): + """A different ``N`` is a different system: ``_x_hold.copy_`` would raise + somewhere unhelpful, and the capacity was sized for the original.""" + nl, _ = _policy_list(skin=1.5, capacity_factor=4.0) + eight_atoms, _ = make_cubic_lattice(n_side=2, spacing=3.0) + with pytest.raises(ValueError, match=r"\b8\b") as excinfo: + nl.build(_batch(eight_atoms)) + assert "64" in str(excinfo.value) + + def test_build_rejects_recast_positions(self): + """No silent cast: a float32 ``pos`` differenced against a float64 + ``_x_hold`` promotes silently — a mixed-precision comparison nobody + asked for — so the bind names both sides and the owner casts.""" + nl, pos = _policy_list(skin=1.5, capacity_factor=4.0) + with pytest.raises(ValueError, match="float32") as excinfo: + nl.build(_batch(pos.to(torch.float32))) + assert "float64" in str(excinfo.value) + + def test_build_rejects_positions_on_another_device(self): + """The ``meta`` device stands in for a real second device: validation + precedes any kernel call, so this needs no CUDA in CI.""" + nl, pos = _policy_list(skin=1.5, capacity_factor=4.0) + elsewhere = torch.empty_like(pos, device="meta") + with pytest.raises(ValueError, match="meta") as excinfo: + nl.build(_batch(elsewhere)) + assert "cpu" in str(excinfo.value) + + def test_build_accepts_the_constructor_cell(self): + """A batch may carry its cell; agreeing with the list's is the point of + checking it. The constructor cell stays the **owner** — the batch's copy + is validated, never adopted.""" + pos, cell = make_cubic_lattice(n_side=4, spacing=3.0) + nl, _ = _policy_list(skin=1.5, capacity_factor=4.0) + before = nl.cell.clone() + nl.build(_batch(pos, cell)) + assert nl.num_edges == 1152 + assert torch.equal(nl.cell, before) + + def test_build_accepts_a_single_system_batched_cell(self): + """``(1, 3, 3)`` is what a collated batch's ``graphs`` namespace holds for + one system, so the bind must read through the leading batch dim.""" + pos, cell = make_cubic_lattice(n_side=4, spacing=3.0) + nl, _ = _policy_list(skin=1.5, capacity_factor=4.0) + before = nl.cell.clone() + nl.build(_batch(pos, cell.unsqueeze(0))) + assert nl.num_edges == 1152 + assert torch.equal(nl.cell, before) + + def test_build_rejects_a_disagreeing_cell(self): + """A cell the list did not build against silently invalidates every + frozen shift, so a disagreement is refused rather than adopted — 0.01 A + is far below anything physical and far above the fp32 round-trip + tolerance the check allows.""" + pos, cell = make_cubic_lattice(n_side=4, spacing=3.0) + nl, _ = _policy_list(skin=1.5, capacity_factor=4.0) + before = nl.cell.clone() + perturbed = cell.clone() + perturbed[0, 0] += 0.01 + with pytest.raises(ValueError, match="cell"): + nl.build(_batch(pos, perturbed)) + assert torch.equal(nl.cell, before) + + def test_build_rejects_a_multi_system_cell(self): + """This list is single-system: one cell, one ``_x_hold``, one capacity. + A ``B > 1`` batch has no meaningful semantics here and is refused by + name rather than silently reduced to its first row.""" + pos, cell = make_cubic_lattice(n_side=4, spacing=3.0) + nl, _ = _policy_list(skin=1.5, capacity_factor=4.0) + with pytest.raises(ValueError, match=r"\b2\b") as excinfo: + nl.build(_batch(pos, cell.repeat(2, 1, 1))) + assert "cell" in str(excinfo.value) + + # --- one ``update``, two input types ----------------------------------- + + def test_the_batch_and_tensor_paths_decide_alike(self): + """The dispatch is a type test at the top of one method, so both inputs + must trace the same policy exactly: same decisions, same bookkeeping, + same buffers. A schedule that both rebuilds and holds, so neither arm of + the comparison is vacuous.""" + nl_tensor, pos = _policy_list(skin=1.0, capacity_factor=4.0) + nl_batch, _ = _policy_list(skin=1.0, capacity_factor=4.0) + schedule = [_displaced(pos, 0.1 * k) for k in range(1, 13)] + by_tensor = [nl_tensor.update(step) for step in schedule] + by_batch = [nl_batch.update(_batch(step)) for step in schedule] + assert set(by_tensor) == {True, False} # the schedule exercises both arms + assert by_batch == by_tensor + assert (nl_batch.ago, nl_batch.rebuild_count, nl_batch.ndanger) == ( + nl_tensor.ago, + nl_tensor.rebuild_count, + nl_tensor.ndanger, + ) + assert nl_batch.num_edges == nl_tensor.num_edges + assert torch.equal(nl_batch.edge_index, nl_tensor.edge_index) + assert torch.equal(nl_batch.shifts, nl_tensor.shifts) + + def test_update_validates_a_batch_exactly_as_build_does(self): + """Shared validation, not a second copy: a batch whose ``pos`` was + re-cast must fail loud on the hot path too, instead of promoting + silently against ``_x_hold`` for the rest of the run.""" + nl, pos = _policy_list(skin=1.0, capacity_factor=4.0) + with pytest.raises(ValueError, match="float32") as excinfo: + nl.update(_batch(pos.to(torch.float32))) + assert "float64" in str(excinfo.value) + + def test_the_dispatch_has_no_twin_entry_points(self): + """One method, two accepted types — ``update_td`` / ``update_pos`` twins + are rejected by design: callers would have to know which one their + driver holds, and the policy state would live behind two doors.""" + assert not hasattr(molix.md.neighbors, "update_td") + assert not hasattr(molix.md.neighbors, "update_pos") + assert not hasattr(NeighborList, "update_td") + assert not hasattr(NeighborList, "update_pos") + + # --- protocol ----------------------------------------------------------- + + def test_a_list_without_build_fails_the_widened_protocol(self): + """``PeriodicPotentialForceField`` calls ``build`` through the + ``NeighborStrategy`` annotation, so the protocol has to actually + *require* it: a stub carrying every link-04 member but no ``build`` must + no longer pass.""" + + class _PolicyWithoutBuild: + edge_index: torch.Tensor = torch.zeros(1, 2, dtype=torch.long) + shifts: torch.Tensor = torch.zeros(1, 3) + num_edges: int = 0 + capacity: int = 1 + cutoff: float = 3.5 + skin: float = 0.0 + + def rebuild(self, positions: torch.Tensor) -> None: + """No-op build.""" + + def update(self, positions: TensorDict | torch.Tensor) -> bool: + return False + + def to( + self, + device: torch.device | str | torch.dtype | None = None, + dtype: torch.dtype | None = None, + ) -> "_PolicyWithoutBuild": + return self + + assert not isinstance(_PolicyWithoutBuild(), NeighborStrategy) + + +# --------------------------------------------------------------------------- +# The binned (cell-list) build path: fixtures, edge keys, comparison contract +# --------------------------------------------------------------------------- + +#: ``(source, target, shift)`` with the shift rounded to 6 decimals — one live +#: edge, reduced to something hashable. Used both directed (as emitted) and +#: canonicalised low->high (see :func:`_canonical_keys`). +_EdgeKey = tuple[int, int, tuple[float, ...]] + + +def _directed_keys(nl: NeighborList) -> list[_EdgeKey]: + """The live edges as hashable ``(source, target, shift)`` keys, as emitted. + + Only ``[0, num_edges)`` is read: the tail is dead padding, not edges. + + Shift components are integer combinations of the cell rows — exact values + like ``12.0`` or ``-6.0`` — so rounding at 6 decimals absorbs the ~1e-13 + disagreement between two different minimum-image reductions without ever + landing near a rounding boundary. + + Args: + nl: A built list. + + Returns: + One key per live edge, in buffer order. Duplicates here are real + duplicates: the same directed edge emitted twice. + """ + n = nl.num_edges + sources = nl.edge_index[:n, 0].tolist() + targets = nl.edge_index[:n, 1].tolist() + shifts = nl.shifts[:n].tolist() + return [ + (source, target, tuple(round(component, 6) for component in shift)) + for source, target, shift in zip(sources, targets, shifts, strict=True) + ] + + +def _canonical_keys(nl: NeighborList) -> list[_EdgeKey]: + """The live edges as orientation-free ``(low, high, shift)`` keys. + + Edge **order is not part of the contract** between the two build backends — + the binned path emits bin-sorted edges, the kernel upper-triangle-sorted + ones — so equality is compared as a *set*. The key must still separate + periodic images, which is what the shift carries: for an edge ``(s, t, D)`` + the reverse edge is ``(t, s, -D)`` (``edge_diff = pos[t] - pos[s] + D`` + flips sign wholesale), so orienting the shift low->high makes the key + independent of which way the edge was emitted, while keeping ``(i, j)`` + across the ``+x`` face distinct from ``(i, j)`` across the ``-x`` one. + + Note: + Both paths emit a **full bidirectional** list, so every undirected pair + contributes *two* live edges that share one canonical key: the unique + canonical keys number ``num_edges / 2``, not ``num_edges``. The + duplicate-free clause is therefore asserted on the *directed* keys (see + :func:`_assert_paths_agree`), which is where a wrapped stencil's double + emission would actually show up. + + Args: + nl: A built list. + + Returns: + One key per live edge, in buffer order. + """ + keys: list[_EdgeKey] = [] + for source, target, shift in _directed_keys(nl): + if source < target: + keys.append((source, target, shift)) + else: + keys.append((target, source, tuple(-component for component in shift))) + return keys + + +def _closest_approach_to(pos: torch.Tensor, cell: torch.Tensor, radius: float) -> float: + """``min_{i != j} |r_ij - radius|`` over minimum-image pairs, in Angstrom. + + The float-tie precondition of every binned-vs-kernel comparison. The two + backends reduce to the minimum image differently — fractional rounding + against the kernel's sequential subtraction — so a pair sitting *on* + ``r_build`` could fall on either side of the closed ``r <= r_build`` filter + for reasons that have nothing to do with the stencil. Every fixture asserts + it keeps clear of that radius, so a failure is always a real disagreement. + + Computed independently of both paths, from fractional rounding, which is + exact here: ``r_build <= min_i w_i / 2`` bounds the fractional offset of + every in-range image by 1/2 (the lemma the binned path rests on). + + Args: + pos: Positions ``(N, 3)`` in Angstrom, wrapped or not. + cell: Cell vectors ``(3, 3)`` in Angstrom, one per row. + radius: The radius to measure the closest approach to, in Angstrom. + + Returns: + The smallest ``|r_ij - radius|`` over all ordered ``i != j`` pairs. + """ + vectors = cell.detach().to(torch.float64) + fractional = pos.detach().to(torch.float64) @ torch.linalg.inv(vectors) + delta = fractional.unsqueeze(0) - fractional.unsqueeze(1) + distance = torch.linalg.norm((delta - torch.round(delta)) @ vectors, dim=-1) + off_diagonal = ~torch.eye(pos.shape[0], dtype=torch.bool) + return float((distance[off_diagonal] - radius).abs().min()) + + +def _assert_paths_agree(binned: NeighborList, kernel: NeighborList, pos: torch.Tensor) -> None: + """The full equality contract between the binned build and the kernel oracle. + + Four clauses, all of them load-bearing: + + 1. **Precondition** — no pair within 1e-9 A of ``r_build``, so no clause + below can be decided by a float tie (:func:`_closest_approach_to`). + 2. **Counts** — ``num_edges`` equal. Set equality alone cannot see a + duplicated edge; a wrapped stencil that emitted every pair twice would + double this. + 3. **Duplicate-free** — every *directed* key occurs once on each path, i.e. + ``len(set(directed)) == num_edges``, and the canonical keys come out at + exactly ``num_edges / 2`` because both paths emit a full bidirectional + list (``(s,t,D)`` and ``(t,s,-D)`` share one canonical key). This is the + intrinsic check the aliasing fixture is built to trip, and it does not + depend on the oracle. + 4. **Set equality** of the canonical keys — catches *missing* edges, which + counts alone cannot. + + Args: + binned: The list built with a ``bin`` grid. + kernel: A list over the same geometry and parameters with ``bin=None``. + pos: The positions both were built at ``(N, 3)`` in Angstrom. + """ + margin = _closest_approach_to(pos, binned.cell, binned.r_build) + assert margin > 1e-9, ( + f"fixture precondition violated: a pair sits {margin:.3e} A from r_build " + f"{binned.r_build} A, where the closed r <= r_build filter is a float coin-flip" + ) + assert binned.num_edges == kernel.num_edges + directed_binned, directed_kernel = _directed_keys(binned), _directed_keys(kernel) + assert len(set(directed_binned)) == binned.num_edges, "the binned path emitted a duplicate edge" + assert len(set(directed_kernel)) == kernel.num_edges, "the kernel path emitted a duplicate edge" + canonical_binned = set(_canonical_keys(binned)) + assert 2 * len(canonical_binned) == binned.num_edges, "the binned list is not bidirectional" + assert canonical_binned == set(_canonical_keys(kernel)) + + +def _pair_distances(nl: NeighborList, pos: torch.Tensor) -> list[tuple[int, int, float]]: + """Live edges as a sorted ``(low, high, |displacement|)`` multiset, in Angstrom. + + Reconstructed the way a consumer does it — ``pos[t] - pos[s] + shift`` — + so it reads the *stored* positions through the *stored* shifts. That makes + it the observable that catches a build-time coordinate wrap leaking into + either of them. + + Args: + nl: A built list. + pos: The positions it was built at ``(N, 3)`` in Angstrom. + + Returns: + One entry per live edge, sorted (a multiset, duplicates kept). + """ + n = nl.num_edges + source, target = nl.edge_index[:n, 0], nl.edge_index[:n, 1] + distance = torch.linalg.norm(pos[target] - pos[source] + nl.shifts[:n], dim=-1) + return sorted( + (min(s, t), max(s, t), round(d, 6)) + for s, t, d in zip(source.tolist(), target.tolist(), distance.tolist(), strict=True) + ) + + +def _jittered_box() -> tuple[torch.Tensor, torch.Tensor]: + """512 atoms in a 24 A cube, jittered off the lattice — the pruning regime. + + The 8x8x8 simple-cubic lattice at 3.0 A displaced by +/-0.4 A uniform under + ``torch.manual_seed(0)``. At ``r_build = 5.0 A`` the grid is 9x9x9 with + ``k_i = 2``, so the stencil searches 125 of 729 bins: real pruning, and no + lattice symmetry left for a broken stencil to hide behind. The closest any + pair comes to ``r_build`` is 3.8e-4 A (measured), clear of the 1e-9 tie band. + + Returns: + ``(positions (512, 3), cell (3, 3))`` in Angstrom, ``float64``. + """ + pos, cell = make_cubic_lattice(n_side=8, spacing=3.0) + torch.manual_seed(0) + return pos + (torch.rand(pos.shape, dtype=torch.float64) * 2.0 - 1.0) * 0.4, cell + + +def _triclinic_forty() -> tuple[torch.Tensor, torch.Tensor]: + """40 atoms in the golden triclinic cell — where the ``|f_i| <= 1/2`` lemma is tested. + + The cell ``[[10,0,0],[6,8,0],[0,0,10]]`` has ``V = 800 A^3`` and + perpendicular widths ``w = (8, 8, 10) A`` against row norms that are all + ``10 A``, so a grid sized on the wrong quantity is immediately visible in + ``n_bins``. At ``cutoff = 3.0``, ``skin = 0.5`` the build radius ``3.5 A`` + stays strictly inside the ``min_i w_i / 2 = 4.000 A`` guard, which is what + makes fractional-rounding minimum image and the kernel's sequential + reduction provably the same image. + + Fractional coordinates are **literals**: generated offline once with + ``torch.manual_seed(1); torch.rand(40, 3, dtype=torch.float64)``, rounded to + 6 decimals and pasted, so no RNG runs at test time and the geometry cannot + drift with a torch RNG change. As pasted, the closest approach to + ``r_build`` is 1.1e-3 A (measured), clear of the 1e-9 tie band. + + Returns: + ``(positions (40, 3), cell (3, 3))`` in Angstrom, ``float64``. + """ + cell = torch.tensor([[10.0, 0.0, 0.0], [6.0, 8.0, 0.0], [0.0, 0.0, 10.0]], dtype=torch.float64) + frac = torch.tensor( + [ + [0.061053, 0.224555, 0.234253], + [0.177099, 0.556068, 0.109444], + [0.460913, 0.708365, 0.579776], + [0.496667, 0.510375, 0.329538], + [0.718206, 0.384511, 0.089797], + [0.117456, 0.640239, 0.196767], + [0.512447, 0.711838, 0.924872], + [0.999699, 0.892730, 0.876720], + [0.844972, 0.154448, 0.170536], + [0.984198, 0.812706, 0.435849], + [0.414321, 0.428408, 0.757762], + [0.922513, 0.964327, 0.176018], + [0.953894, 0.313379, 0.454398], + [0.295552, 0.187507, 0.243258], + [0.349296, 0.444072, 0.406873], + [0.285938, 0.803593, 0.321766], + [0.363903, 0.298510, 0.663531], + [0.255167, 0.414372, 0.839555], + [0.741833, 0.286491, 0.792859], + [0.500116, 0.897740, 0.105125], + [0.580914, 0.986660, 0.131524], + [0.239137, 0.304684, 0.515845], + [0.451441, 0.492893, 0.530066], + [0.264720, 0.167118, 0.548191], + [0.237952, 0.537363, 0.442156], + [0.645389, 0.537569, 0.224480], + [0.663186, 0.843878, 0.010876], + [0.280679, 0.930112, 0.543798], + [0.812327, 0.774969, 0.730758], + [0.992421, 0.728189, 0.232834], + [0.999747, 0.554004, 0.420049], + [0.541916, 0.864175, 0.431247], + [0.121250, 0.895592, 0.878425], + [0.912789, 0.968760, 0.415001], + [0.409411, 0.688470, 0.679978], + [0.641520, 0.401901, 0.487456], + [0.956891, 0.517200, 0.953366], + [0.854016, 0.955512, 0.083597], + [0.168356, 0.188330, 0.938444], + [0.354260, 0.202702, 0.506931], + ], + dtype=torch.float64, + ) + return frac @ cell, cell + + +def _unwrapped_lattice() -> tuple[torch.Tensor, torch.Tensor]: + """The 64-atom lattice with 8 atoms pushed out of the box by whole cell vectors. + + Physically the identical system — a lattice vector is a symmetry — but the + coordinates are no longer inside ``[0, L)``, which is exactly the state an + unwrapped MD trajectory drifts into. The binned path wraps *fractionally* to + index its bins; this fixture is what pins that the wrap is an indexing + device only and never reaches the stored positions or the shifts. + + Returns: + ``(positions (64, 3), cell (3, 3))`` in Angstrom, ``float64``. + """ + pos, cell = make_cubic_lattice(n_side=4, spacing=3.0) + moved = pos.clone() + rows = (0, 1, 2, 0, 1, 2, 0, 1) + signs = (1.0, 1.0, 1.0, -1.0, -1.0, -1.0, 1.0, -1.0) + for atom, (row, sign) in enumerate(zip(rows, signs, strict=True)): + moved[atom] = moved[atom] + sign * cell[row] + return moved, cell + + +def _coincident_lattice() -> tuple[torch.Tensor, torch.Tensor]: + """The 64-atom lattice carrying both flavours of zero-distance pair. + + ``pos[1]`` is moved onto ``pos[0]`` — coincident in real space — and + ``pos[2]`` onto ``pos[3] + a_1``, i.e. coincident only *through* a lattice + vector: raw separation 12 A, minimum image the zero vector. The compiled + kernel rejects both (``distances > 0``), so the binned path's ``r > 0`` + filter has to reject exactly the same two and nothing else. + + Every atom still sits on a lattice site, so the pair distances stay + ``{0, 3.0, 4.2426, 5.196, ...} A`` and the closest approach to + ``r_build = 5.0 A`` remains 0.196 A. + + Returns: + ``(positions (64, 3), cell (3, 3))`` in Angstrom, ``float64``. + """ + pos, cell = make_cubic_lattice(n_side=4, spacing=3.0) + moved = pos.clone() + moved[1] = pos[0] + moved[2] = pos[3] + cell[0] + return moved, cell + + +def _binned_pair( + pos: torch.Tensor, + cell: torch.Tensor, + *, + bin: float = 0.0, + cutoff: float = 3.5, + skin: float = 1.5, + capacity_factor: float = 1.35, +) -> tuple[NeighborList, NeighborList]: + """Two lists over one geometry: the binned build and its ``bin=None`` oracle. + + The compiled path is not modified by this link, which is precisely what + lets it stand as the reference for the new one. + + Args: + pos: Positions ``(N, 3)`` in Angstrom. + cell: Cell vectors ``(3, 3)`` in Angstrom. + bin: Requested perpendicular bin thickness in Angstrom; ``0.0`` selects + the automatic ``r_build / 2``. + cutoff: Interaction cutoff in Angstrom. + skin: Verlet skin in Angstrom. + capacity_factor: Buffer headroom, as in the constructor. + + Returns: + ``(binned, kernel)`` — both built at ``pos``. + """ + binned = NeighborList( + cell=cell, + cutoff=cutoff, + positions=pos, + skin=skin, + bin=bin, + capacity_factor=capacity_factor, + ) + kernel = NeighborList( + cell=cell, + cutoff=cutoff, + positions=pos, + skin=skin, + bin=None, + capacity_factor=capacity_factor, + ) + return binned, kernel + + +def _policy_schedule(pos: torch.Tensor) -> list[torch.Tensor]: + """One atom drifting 0.1 A per step along x, twelve steps. + + Against ``skin = 1.0`` (half-skin 0.5 A) this crosses the rebuild criterion + twice, so both arms of the gate are exercised and neither half of a + comparison over it is vacuous. Every configuration keeps its closest pair + 0.034 A away from ``r_build = 4.5 A`` (measured), so the per-rebuild edge-set + comparison never rides a float tie. + + Args: + pos: The reference positions ``(N, 3)`` in Angstrom. + + Returns: + Twelve position tensors, each ``(N, 3)`` in Angstrom. + """ + return [_displaced(pos, 0.1 * step) for step in range(1, 13)] + + +class TestNeighborListBinned: + """The pure-torch binned (cell-list) build path behind ``bin=``. + + ``bin=None`` (the default) keeps the compiled O(N^2) kernel and is therefore + available as the *oracle*: for every fixture the two backends must return + the identical set of ``(source, target, shift)`` triples. Edge order is not + part of that contract (see :func:`_canonical_keys`); completeness, + duplicate-freedom and the ``0 < r <= r_build`` filter are. + + Reference: + Allen, M. P.; Tildesley, D. J. *Computer Simulation of Liquids*, 2nd + ed.; Oxford University Press, 2017. + https://doi.org/10.1093/oso/9780198803195.001.0001 — cell lists, the + 27-cell stencil and minimum-image validity. + + Thompson, A. P. et al. *Comput. Phys. Commun.* **271** (2022) 108171, + https://doi.org/10.1016/j.cpc.2021.108171; binning policy in + ``lammps/lammps`` develop ``src/nbin_standard.cpp`` + (``binsize_optimal = 0.5 * cutneighmax``), which is the ``bin=0.0`` + automatic size adopted here. + """ + + # --- grid derivation: n_bins is the only public window onto the stencil --- + + def test_the_auto_bin_is_half_the_build_radius(self): + """``bin=0.0`` requests ``b = r_build / 2`` (LAMMPS ``nbin_standard``). + + On the 12 A cube at ``r_build = 5.0 A`` that is 2.5 A, and + ``n_i = floor(w_i / b) = floor(4.8) = 4``, giving effective bins of + 3.0 A and ``k_i = ceil(5.0 / 3.0) = 2``. So ``2k + 1 = 5 > 4``: the raw + stencil wraps onto the same bin twice, which is the aliasing regime the + equivalence fixtures below are built to trip. + """ + pos, cell = make_cubic_lattice(n_side=4, spacing=3.0) + nl = NeighborList(cell=cell, cutoff=3.5, positions=pos, skin=1.5, bin=0.0) + assert nl.n_bins == (4, 4, 4) + assert isinstance(nl.n_bins, tuple) + + def test_the_auto_grid_scales_with_the_cell(self): + """Same 2.5 A request in a 24 A cube: ``floor(24 / 2.5) = 9`` per axis. + + The count is what makes the path O(N): 9^3 = 729 bins searched 125 at a + time, against the 4^3 = 64 bins of the 12 A cube where the stencil still + covers everything. + """ + pos, cell = make_cubic_lattice(n_side=8, spacing=3.0) + nl = NeighborList(cell=cell, cutoff=3.5, positions=pos, skin=1.5, bin=0.0) + assert nl.n_bins == (9, 9, 9) + + def test_the_auto_grid_follows_the_build_radius_not_the_cutoff(self): + """``skin=0.0`` moves ``r_build`` to 3.5 A, so the auto bin is 1.75 A and + ``floor(12 / 1.75) = 6``. A grid derived from ``cutoff`` would report the + same ``(6, 6, 6)`` here *only because* the skin is zero — which is why + the skinned case above pins ``(4, 4, 4)`` and this one pins the move.""" + pos, cell = make_cubic_lattice(n_side=4, spacing=3.0) + nl = NeighborList(cell=cell, cutoff=3.5, positions=pos, skin=0.0, bin=0.0) + assert nl.n_bins == (6, 6, 6) + + def test_an_explicit_bin_thickness_sizes_the_grid(self): + """``bin=5.0`` in the 12 A cube: ``floor(12 / 5) = 2`` bins of 6.0 A, + ``k_i = ceil(5.0 / 6.0) = 1``. The explicit knob is a *requested* + thickness — the effective one is ``w_i / n_i``, never smaller.""" + pos, cell = make_cubic_lattice(n_side=4, spacing=3.0) + nl = NeighborList(cell=cell, cutoff=3.5, positions=pos, skin=1.5, bin=5.0) + assert nl.n_bins == (2, 2, 2) + + def test_a_bin_as_wide_as_the_cell_degenerates_to_one_bin(self): + """``max(1, floor(w_i / b))`` — a bin request at or beyond the cell width + must clamp to a single bin per axis rather than produce a zero-bin grid + (a division by zero in the flat-id arithmetic). One bin per axis is the + graceful all-pairs degeneration: correct, just not faster.""" + pos, cell = make_cubic_lattice(n_side=4, spacing=3.0) + nl = NeighborList(cell=cell, cutoff=3.5, positions=pos, skin=1.5, bin=12.0) + assert nl.n_bins == (1, 1, 1) + + def test_the_grid_is_sized_on_perpendicular_widths_not_row_norms(self): + """The discriminating golden — triclinic sizing, asserted directly. + + For ``[[10,0,0],[6,8,0],[0,0,10]]`` the perpendicular widths are + ``w = (8, 8, 10) A`` while **all three row norms are 10 A**. At + ``r_build = 3.5 A`` the auto bin is 1.75 A, so sizing on ``w`` gives + ``(floor(8/1.75), floor(8/1.75), floor(10/1.75)) = (4, 4, 5)`` where an + implementation sizing on ``||a_i||`` would report ``(5, 5, 5)``. That + wrong grid makes the bins thinner than requested along the sheared axes + and the stencil incomplete — a silently short neighbour list. + """ + pos, cell = _triclinic_forty() + nl = NeighborList(cell=cell, cutoff=3.0, positions=pos, skin=0.5, bin=0.0) + assert nl.n_bins == (4, 4, 5) + + def test_the_default_backend_reports_no_grid(self): + """``bin=None`` — passed or omitted — is the untouched kernel path. + + ``n_bins is None`` is the observable that says "no grid was derived", + and the crystallographic 1152 edges say the build itself is bit-for-bit + the pre-link behaviour. The explicit form is constructed first so this + also pins that ``None`` is *accepted*, not just defaulted to. + """ + pos, cell = make_cubic_lattice(n_side=4, spacing=3.0) + explicit = NeighborList(cell=cell, cutoff=3.5, positions=pos, skin=1.5, bin=None) + omitted = NeighborList(cell=cell, cutoff=3.5, positions=pos, skin=1.5) + assert explicit.bin is None + assert explicit.n_bins is None + assert omitted.n_bins is None + assert explicit.num_edges == omitted.num_edges == 1152 + + # --- construction validation: refuse before allocating anything --------- + + def test_a_negative_bin_is_refused(self): + """A negative thickness is a typo, and ``0.0`` already means "choose for + me" — so the message has to name both, or the reader's next guess is + that ``-1`` was the way to ask for the automatic size.""" + pos, cell = make_cubic_lattice(n_side=4, spacing=3.0) + with pytest.raises(ValueError, match=r"\bbin\b") as excinfo: + NeighborList(cell=cell, cutoff=3.5, positions=pos, skin=1.5, bin=-1.0) + assert "0.0" in str(excinfo.value) + + def test_a_bin_far_below_the_build_radius_is_refused(self): + """A 0.05 A bin at ``r_build = 5.0 A`` is a hang, not a fine grid. + + It gives 240 bins per axis, ``k_i = ceil(5.0 / 0.05) = 100`` and a + ``201^3 ~ 8.1e6``-offset stencil — a Python loop that never finishes. + Refused at construction against the half-width cap of 8, with the + measured numbers in the message: the implied 100 and the cap it broke. + """ + pos, cell = make_cubic_lattice(n_side=4, spacing=3.0) + with pytest.raises(ValueError, match=r"\bbin\b") as excinfo: + NeighborList(cell=cell, cutoff=3.5, positions=pos, skin=1.5, bin=0.05) + message = str(excinfo.value) + assert "100" in message + assert re.search(r"\b8\b", message), f"the stencil cap is not named in: {message}" + assert "0.0" in message + + # --- equivalence with the kernel oracle --------------------------------- + + def test_the_binned_path_matches_the_kernel_on_the_aliasing_lattice(self): + """PRIMARY falsification, in the regime that breaks a naive stencil. + + 64 atoms, ``n_bins = (4, 4, 4)``, ``k_i = 2``: the raw offsets + ``-2..2 (mod 4)`` visit bins 2 and 3 twice each, so an implementation + that does not reduce the stencil to *distinct residues* emits a large + share of the pairs twice — which the duplicate-free and count clauses + catch, and set equality alone would not. + """ + pos, cell = make_cubic_lattice(n_side=4, spacing=3.0) + binned, kernel = _binned_pair(pos, cell) + _assert_paths_agree(binned, kernel, pos) + + @pytest.mark.parametrize(("skin", "edges"), [(0.0, 384), (1.5, 1152)], ids=["bare", "skinned"]) + def test_the_binned_path_builds_at_the_build_radius(self, skin: float, edges: int): + """``skin`` sets the radius, ``bin`` only the search strategy. + + Crystallography, not a fit: the simple-cubic lattice has 6 neighbours at + 3.0 A and 12 more at ``3*sqrt(2) = 4.2426 A`` (the 8 body diagonals at + 5.196 A stay out of both radii), so 64*6 = 384 edges at ``r_build = + 3.5 A`` and 64*18 = 1152 at 5.0 A. A binned path that built at ``cutoff`` + would report 384 in both rows, leaving the Verlet skin dead while every + energy still looked plausible. + """ + pos, cell = make_cubic_lattice(n_side=4, spacing=3.0) + binned, kernel = _binned_pair(pos, cell, skin=skin) + assert binned.num_edges == edges + assert kernel.num_edges == edges + + def test_the_binned_path_matches_the_kernel_on_a_jittered_dense_box(self): + """512 atoms off-lattice: the stencil actually prunes, 125 bins of 729. + + No lattice symmetry is left to make a half-wrong stencil look right, and + no count golden is available — the kernel is the oracle and the only + claim is that the two agree edge for edge. + """ + pos, cell = _jittered_box() + binned, kernel = _binned_pair(pos, cell) + _assert_paths_agree(binned, kernel, pos) + + def test_the_binned_path_matches_the_kernel_on_a_triclinic_cell(self): + """Where fractional rounding has to reproduce the kernel's reduction. + + In a sheared cell the two minimum-image algorithms are visibly different + procedures; they agree only because ``r_build <= min_i w_i / 2`` bounds + every in-range image's fractional offset by 1/2. This fixture is the + measurement of that lemma. + """ + pos, cell = _triclinic_forty() + binned, kernel = _binned_pair(pos, cell, cutoff=3.0, skin=0.5) + _assert_paths_agree(binned, kernel, pos) + + def test_the_binned_path_matches_the_kernel_on_unwrapped_positions(self): + """Eight atoms sitting a whole cell vector outside the box. + + The binned path wraps fractionally to index its bins. If that wrapped + copy leaked into the displacement or the shift, this fixture — the same + physical system as the plain lattice — would disagree with the kernel. + """ + pos, cell = _unwrapped_lattice() + binned, kernel = _binned_pair(pos, cell) + _assert_paths_agree(binned, kernel, pos) + + def test_a_whole_cell_translation_leaves_the_pair_distances_unchanged(self): + """The other half of the wrap invariant, read through the shifts. + + Translating atoms by lattice vectors is a symmetry, so the multiset of + ``(low, high, ||pos[t] - pos[s] + shift||)`` must be *identical* to the + untranslated lattice's. Set equality against the kernel cannot see this: + it would also hold if both paths reconstructed the same wrong geometry. + """ + pos, cell = make_cubic_lattice(n_side=4, spacing=3.0) + moved, _ = _unwrapped_lattice() + binned = NeighborList(cell=cell, cutoff=3.5, positions=pos, skin=1.5, bin=0.0) + translated = NeighborList(cell=cell, cutoff=3.5, positions=moved, skin=1.5, bin=0.0) + assert _pair_distances(translated, moved) == _pair_distances(binned, pos) + + def test_both_paths_drop_the_zero_distance_pairs(self): + """``0 < r <= r_build`` filter parity, both flavours of ``r == 0``. + + Atom 1 is coincident with atom 0 in real space; atom 2 is coincident + with atom 3 only through ``a_1``, so its raw separation is 12 A and its + minimum image is the zero vector. The compiled backends reject both + (``distances > 0`` in the C++ path, ``distance2 == 0`` in the CUDA one), + so a binned path filtering on ``r <= r_build`` alone would build two + self-cancelling edges the oracle does not have — and, worse, a division + by zero anywhere a unit vector is taken. + """ + pos, cell = _coincident_lattice() + binned, kernel = _binned_pair(pos, cell) + _assert_paths_agree(binned, kernel, pos) + pairs = {(low, high) for low, high, _ in _canonical_keys(binned)} + assert (0, 1) not in pairs + assert (2, 3) not in pairs + + def test_the_edge_set_is_independent_of_the_bin_size(self): + """``bin`` is a cost knob, never a physics knob. + + The sweep spans every regime the derivation has: 0.0 and 2.5 A give the + 9x9x9 pruning grid, 5.0 A a coarse 4x4x4 one, and 24.0 A the single-bin + full degeneration. A bin size that changes the edge set is a broken + stencil derivation, full stop — so all four are compared against the one + kernel oracle rather than against each other. + """ + pos, cell = _jittered_box() + kernel = NeighborList(cell=cell, cutoff=3.5, positions=pos, skin=1.5, bin=None) + for thickness in (0.0, 2.5, 5.0, 24.0): + binned = NeighborList(cell=cell, cutoff=3.5, positions=pos, skin=1.5, bin=thickness) + _assert_paths_agree(binned, kernel, pos) + + # --- orthogonality with the rest of the class --------------------------- + + def test_the_rebuild_policy_is_untouched_by_the_binned_backend(self): + """``skin`` decides the radius, ``bin`` the strategy, ``update`` the moment. + + Over one scripted drift the gate's decisions and all three counters must + be identical with and without a grid: the policy reads displacements and + a clock, neither of which the build backend touches. + """ + pos, cell = make_cubic_lattice(n_side=4, spacing=3.0) + binned, kernel = _binned_pair(pos, cell, skin=1.0, capacity_factor=4.0) + schedule = _policy_schedule(pos) + by_binned = [binned.update(step) for step in schedule] + by_kernel = [kernel.update(step) for step in schedule] + assert set(by_kernel) == {True, False} # the schedule exercises both arms + assert by_binned == by_kernel + assert (binned.ago, binned.rebuild_count, binned.ndanger) == ( + kernel.ago, + kernel.rebuild_count, + kernel.ndanger, + ) + + def test_the_live_edges_after_a_policy_rebuild_match_the_kernel(self): + """The same schedule, but checking *what* was built, not *when*. + + A rebuild driven through ``update`` goes down the same dispatch as the + constructor's initial build, so the equivalence has to survive it — a + path that only agreed on the first build would leave the run drifting + away from the oracle one rebuild at a time. + """ + pos, cell = make_cubic_lattice(n_side=4, spacing=3.0) + binned, kernel = _binned_pair(pos, cell, skin=1.0, capacity_factor=4.0) + compared = 0 + for moved in _policy_schedule(pos): + rebuilt = binned.update(moved) + assert kernel.update(moved) is rebuilt + if rebuilt: + _assert_paths_agree(binned, kernel, moved) + compared += 1 + assert compared == 2 # the schedule crosses the half-skin twice + + def test_a_cast_carries_the_grid_and_the_inverse_cell(self): + """``to(dtype)`` must move the binned state with the buffers. + + The stencil and the cached inverse cell are the two tensors nothing else + in the class owns, so they are exactly what a ``to`` that only knows + about ``edge_index`` / ``shifts`` / ``cell`` / ``_x_hold`` leaves behind + — in the wrong dtype (a float64 inverse cell against float32 positions + promotes *silently*) or on the wrong device. The rebuilt 1152 is the + cheapest observable that both survived. + """ + pos, cell = make_cubic_lattice(n_side=4, spacing=3.0) + nl = NeighborList(cell=cell, cutoff=3.5, positions=pos, skin=1.5, bin=0.0) + nl.to(torch.float32) + nl.rebuild(pos.to(torch.float32)) + assert nl.n_bins == (4, 4, 4) + assert nl.num_edges == 1152 + assert nl.edge_index.dtype == torch.long # indices are never cast + + def test_the_protocol_does_not_learn_about_bins(self): + """The build backend is an implementation detail of *this* class. + + ``NeighborStrategy`` stays as link 04 left it: a binned list satisfies + it unchanged, and ``bin`` / ``n_bins`` must not appear in the protocol, + which would be creep that every other strategy — including the stubs in + this file — would then have to carry. + """ + pos, cell = make_cubic_lattice(n_side=4, spacing=3.0) + nl = NeighborList(cell=cell, cutoff=3.5, positions=pos, skin=1.5, bin=0.0) + assert isinstance(nl, NeighborStrategy) + assert "bin" not in NeighborStrategy.__annotations__ + assert "n_bins" not in NeighborStrategy.__annotations__ diff --git a/tests/test_molix/test_md/test_runner.py b/tests/test_molix/test_md/test_runner.py new file mode 100644 index 0000000..8bc9a64 --- /dev/null +++ b/tests/test_molix/test_md/test_runner.py @@ -0,0 +1,235 @@ +"""Tests for the hook-driven MD runner (molix.md.runner). + +Single-function correctness only. ``MDRunner.run`` is exercised with an +analytic harmonic well just far enough to observe its own contract (hook +dispatch order, step accounting, typed return); every hook class is driven +**directly** with synthetic :class:`MDObservables` / :class:`MDState` — no +trajectory loop stands between a hook test and the method it tests. +""" + +import torch + +import molix.md +import molix.md.runner +from molix.md import ( + HarmonicForceField, + LangevinVerletIntegrator, + MDCheckpointHook, + MDHook, + MDObservables, + MDRunner, + MDState, + TrajectoryHook, +) + +_DTYPE = torch.float64 + + +def _integrator(seed: int = 0): + return LangevinVerletIntegrator( + HarmonicForceField(1.0).to(_DTYPE), dt=0.01, gamma=1.0, kbt=1.0, mass=1.0, seed=seed + ) + + +def _obs(n_atoms: int = 4, fill: float = 1.0) -> MDObservables: + """A synthetic observation — hooks only read fields, so values are free.""" + return MDObservables( + pos=torch.full((n_atoms, 3), fill, dtype=_DTYPE), + vel=torch.full((n_atoms, 3), 2.0 * fill, dtype=_DTYPE), + forces=torch.full((n_atoms, 3), -fill, dtype=_DTYPE), + potential=torch.tensor(fill, dtype=_DTYPE), + kinetic=torch.tensor(2.0 * fill, dtype=_DTYPE), + total=torch.tensor(3.0 * fill, dtype=_DTYPE), + temperature=torch.tensor(300.0, dtype=_DTYPE), + ) + + +class _RecordingHook(MDHook): + """Records lifecycle call order and the per-step observables.""" + + def __init__(self): + self.events: list[str] = [] + self.steps: list[int] = [] + self.observed: list = [] + + def on_run_start(self, runner): + self.events.append("start") + + def on_step_end(self, runner, step, obs): + self.events.append("step") + self.steps.append(step) + self.observed.append(obs) + + def on_run_end(self, runner): + self.events.append("end") + + +class TestMDRunner: + def test_fires_lifecycle_and_advances_step(self): + rec = _RecordingHook() + runner = MDRunner(_integrator(), mass=1.0, hooks=[rec]) + out = runner.run(torch.zeros(4, 3, dtype=_DTYPE), torch.zeros(4, 3, dtype=_DTYPE), 3) + + assert rec.events == ["start", "step", "step", "step", "end"] + assert rec.steps == [1, 2, 3] + assert isinstance(out, MDState) + assert out.pos.shape == (4, 3) + assert out.forces.shape == (4, 3) + + def test_observables_are_typed_and_consistent(self): + rec = _RecordingHook() + runner = MDRunner(_integrator(), mass=1.0, hooks=[rec]) + runner.run(torch.randn(4, 3, dtype=_DTYPE), torch.randn(4, 3, dtype=_DTYPE), 2) + + last = rec.observed[-1] + assert torch.allclose(last.total, last.potential + last.kinetic) + assert last.forces.shape == (4, 3) + + def test_dof_follows_the_integrator_convention(self): + """The DoF convention lives on ``Integrator.removed_dof`` — the runner + must not duck-type thermostat internals (the old ``getattr(integrator, + "gamma")`` sniffing silently mis-reported temperature for custom + integrators).""" + assert _integrator().removed_dof == 0 # gamma > 0 + nve = LangevinVerletIntegrator( + HarmonicForceField(1.0).to(_DTYPE), dt=0.01, gamma=0.0, kbt=0.0, mass=1.0 + ) + assert nve.removed_dof == 3 + + def test_step_start_precedes_each_advance(self): + """``on_step_start`` fires before the chunk it precedes — the ordering + a hook needs to observe or stage per-step state ahead of the force + evaluations inside the advance. (Neighbour rebuilds are **not** such a + user: the list's policy runs inside ``Integrator.eval_force``, at the + positions being evaluated.)""" + events: list[tuple[str, int]] = [] + + class _Interleaved(MDHook): + def on_step_start(self, runner, step, state): + events.append(("start", step)) + + def on_step_end(self, runner, step, obs): + events.append(("end", step)) + + runner = MDRunner(_integrator(), mass=1.0, hooks=[_Interleaved()]) + runner.run(torch.zeros(2, 3, dtype=_DTYPE), torch.zeros(2, 3, dtype=_DTYPE), 4, chunk=2) + assert events == [("start", 0), ("end", 2), ("start", 2), ("end", 4)] + + def test_misaligned_hook_cadence_is_rejected(self): + """A declared cadence that chunking would silently skip must raise.""" + import pytest + + hook = TrajectoryHook("/tmp/_chunk_misalign.pt", stride=3, write_xyz=False) + runner = MDRunner(_integrator(), mass=1.0, hooks=[hook]) + with pytest.raises(ValueError, match="multiple of chunk"): + runner.run(torch.zeros(2, 3, dtype=_DTYPE), torch.zeros(2, 3, dtype=_DTYPE), 6, chunk=2) + + def test_hook_priority_orders_firing(self): + """Lower priority fires earlier; ties keep registration order.""" + order: list[str] = [] + + class _Tagged(MDHook): + def __init__(self, tag): + self.tag = tag + + def on_run_start(self, runner): + order.append(self.tag) + + runner = MDRunner( + _integrator(), + mass=1.0, + hooks=[(_Tagged("late"), 100), (_Tagged("first"), 0), _Tagged("default")], + ) + runner.run(torch.zeros(2, 3, dtype=_DTYPE), torch.zeros(2, 3, dtype=_DTYPE), 1) + assert order == ["first", "late", "default"] + + +class TestTrajectoryHook: + """Driven directly: synthetic observables in, files out.""" + + def _drive(self, hook: TrajectoryHook, n_steps: int) -> None: + for step in range(1, n_steps + 1): + hook.on_step_end(None, step, _obs(fill=float(step))) + hook.on_run_end(None) + + def test_persists_pt_and_xyz(self, tmp_path): + numbers = torch.tensor([1, 6, 8, 1]) + out_pt = tmp_path / "traj.pt" + self._drive(TrajectoryHook(out_pt, stride=2, numbers=numbers, write_xyz=True), 10) + + payload = torch.load(out_pt, weights_only=False) + # stride=2 over steps 1..10: kept at 2,4,6,8,10 -> 5 frames. + assert payload["pos"].shape == (5, 4, 3) + assert payload["forces"].shape == (5, 4, 3) + assert payload["temp"].shape == (5,) + assert payload["stride"] == 2 + assert torch.equal(payload["Z"], numbers) + assert float(payload["pos"][0, 0, 0]) == 2.0 # first kept step is 2 + + lines = out_pt.with_suffix(".xyz").read_text().splitlines() + # Each extended-XYZ frame: 1 count line + 1 comment + N atom lines. + assert lines[0] == "4" + assert "energy=" in lines[1] and "temperature=" in lines[1] + assert lines[2].split()[0] == "H" # Z=1 -> H + + def test_shard_flush_matches_single_buffer(self, tmp_path): + """Shard-flushing (bounded host memory) yields the same file as one + in-memory buffer, and cleans up its shard files.""" + big = tmp_path / "big.pt" # one buffer (flush_every >> frames) + self._drive(TrajectoryHook(big, flush_every=1000), 10) + small = tmp_path / "small.pt" # flush every 2 frames -> 5 shards + self._drive(TrajectoryHook(small, flush_every=2), 10) + + a = torch.load(big, weights_only=True) + b = torch.load(small, weights_only=True) + assert b["pos"].shape == (10, 4, 3) + for key in ("pos", "vel", "forces", "pe", "temp"): + assert torch.equal(a[key], b[key]), key + assert not list(tmp_path.glob("small.part*.pt")) # shards removed + + def test_without_numbers_skips_xyz(self, tmp_path): + out_pt = tmp_path / "traj.pt" + self._drive(TrajectoryHook(out_pt, stride=1, write_xyz=True), 4) # no numbers + assert out_pt.exists() + assert not out_pt.with_suffix(".xyz").exists() + + def test_declares_its_stride_as_cadence(self): + assert TrajectoryHook("x.pt", stride=3).cadence == 3 + + +def test_the_step_start_rebuild_hook_is_gone(): + """``NeighborListHook`` is deleted, not deprecated (``stage: experimental``). + + It rebuilt at the *start-of-step* positions while velocity-Verlet evaluates + ``F`` at the end-of-step ones, so ``F`` was not ``-∇E`` of the surface the + list defined — a one-signed leak measured ~30× worse at ``every=5`` than at + ``every=1``. The seam is ``Integrator.eval_force``; the migration is + ``NeighborList(skin=, every=, delay=, check=)``. A back-compat alias would + defeat this guard, so the name must be absent from the module, the package + namespace and ``__all__``. + """ + assert not hasattr(molix.md.runner, "NeighborListHook") + assert not hasattr(molix.md, "NeighborListHook") + assert "NeighborListHook" not in molix.md.__all__ + + +class TestMDCheckpointHook: + def test_writes_a_restartable_payload_on_cadence(self, tmp_path): + ck = tmp_path / "state.pt" + hook = MDCheckpointHook(ck, every=5) + hook.on_step_end(None, 3, _obs()) # off-cadence: no write + assert not ck.exists() + hook.on_step_end(None, 5, _obs(fill=7.0)) + saved = torch.load(ck, weights_only=True) + assert saved["step"] == 5 + assert torch.equal(saved["pos"], torch.full((4, 3), 7.0, dtype=_DTYPE)) + assert torch.equal(saved["vel"], torch.full((4, 3), 14.0, dtype=_DTYPE)) + + def test_step_offset_keeps_one_step_axis(self, tmp_path): + """A resumed segment reports absolute steps, not segment-local ones.""" + ck = tmp_path / "state.pt" + MDCheckpointHook(ck, every=5, step_offset=100).on_step_end(None, 5, _obs()) + assert torch.load(ck, weights_only=True)["step"] == 105 + + def test_declares_its_interval_as_cadence(self, tmp_path): + assert MDCheckpointHook(tmp_path / "s.pt", every=7).cadence == 7 diff --git a/tests/test_molix/test_md/test_types.py b/tests/test_molix/test_md/test_types.py new file mode 100644 index 0000000..1d868a0 --- /dev/null +++ b/tests/test_molix/test_md/test_types.py @@ -0,0 +1,44 @@ +"""Tests for molix.md.types — typed pytree contracts.""" + +import torch +from torch.utils._pytree import tree_flatten, tree_unflatten + +from molix.md import ForceOutput, MDObservables, MDState + + +class TestMDState: + """MDState must behave as a transparent pytree for torch.compile/func.""" + + def test_is_a_pytree(self): + s = MDState(torch.randn(4, 3), torch.randn(4, 3), torch.randn(4, 3), torch.tensor(1.0)) + leaves, spec = tree_flatten(s) + assert len(leaves) == 4 + rebuilt = tree_unflatten(leaves, spec) + assert isinstance(rebuilt, MDState) + assert torch.equal(rebuilt.pos, s.pos) and torch.equal(rebuilt.energy, s.energy) + + def test_forces_field_is_plural(self): + """The per-atom force tensor is named ``forces`` everywhere — the + singular ``force`` / plural ``forces`` split was a naming-drift bug.""" + assert "forces" in MDState._fields + assert "force" not in MDState._fields + + +class TestForceOutput: + def test_is_a_pytree(self): + fo = ForceOutput(torch.tensor(1.0), torch.randn(4, 3)) + leaves, spec = tree_flatten(fo) + assert isinstance(tree_unflatten(leaves, spec), ForceOutput) + + +class TestMDObservables: + def test_fields_cover_the_hook_contract(self): + assert MDObservables._fields == ( + "pos", + "vel", + "forces", + "potential", + "kinetic", + "total", + "temperature", + ) diff --git a/tests/test_molix/test_md_dynamics.py b/tests/test_molix/test_md_dynamics.py deleted file mode 100644 index dd949b6..0000000 --- a/tests/test_molix/test_md_dynamics.py +++ /dev/null @@ -1,175 +0,0 @@ -"""Tests for the paired-trajectory protocol + TrajectoryArtifact (tiny PiNet).""" - -import pytest -import torch - -from molix.md import ( - HAS_ASE, - PotentialForceField, - TrajectoryArtifact, - build_paired_trajectory, - make_pinet_calculator, -) -from molix.quant import Quantizer -from molzoo.pinet import PiNetPotential -from tests.conftest import make_graph_batch - -_DEVICE = torch.device("cpu") - - -def _tiny_potential() -> PiNetPotential: - torch.manual_seed(0) - return ( - PiNetPotential( - atom_types=[1, 6, 7, 8], - r_max=4.0, - n_basis=3, - pp_nodes=[8, 8], - pi_nodes=[8, 8], - ii_nodes=[8, 8], - depth=2, - rank=3, - hidden_dim=16, - ) - .to(_DEVICE) - .eval() - ) - - -def _template(): - pos = torch.tensor( - [[0.0, 0.0, 0.0], [1.1, 0.1, 0.0], [0.3, 1.2, 0.2], [1.4, 1.1, -0.1]], - dtype=torch.float32, - device=_DEVICE, - ) - z = torch.tensor([1, 6, 7, 8], dtype=torch.long, device=_DEVICE) - edge_index = torch.tensor( - [[0, 1], [1, 0], [0, 2], [2, 0], [1, 3], [3, 1], [2, 3], [3, 2]], - dtype=torch.long, - device=_DEVICE, - ) - batch = torch.zeros(4, dtype=torch.long, device=_DEVICE) - return make_graph_batch(pos, z, edge_index, batch) - - -def _warmed_ref_and_quant(template): - ref = _tiny_potential() - quant = _tiny_potential() - ref(template.clone(), compute_forces=False) # warmup lazy params - quant(template.clone(), compute_forces=False) - quant.load_state_dict(Quantizer("int4").quantize_state_dict(ref.state_dict())) - return ref, quant - - -def _run(ref, quant, template, n_steps=5): - pos0 = template["atoms", "pos"].clone() - vel0 = torch.zeros_like(pos0) - return build_paired_trajectory( - ref, - quant, - template, - pos0, - vel0, - n_steps, - dt=0.001, - gamma=1.0, - kbt=0.0257, - mass=12.0, - seed=3, - condition={"scheme": "int4", "dataset": "qm9"}, - ) - - -def test_paired_trajectory_schema(): - template = _template() - ref, quant = _warmed_ref_and_quant(template) - art = _run(ref, quant, template, n_steps=5) - - assert isinstance(art, TrajectoryArtifact) - n = 4 - assert art.pos.shape == (5, n, 3) - assert art.vel.shape == (5, n, 3) - assert art.energy.shape == (5,) - assert art.f_ref.shape == (5, n, 3) - assert art.f_quant.shape == (5, n, 3) - assert art.df.shape == (5, n, 3) - assert torch.allclose(art.df, art.f_quant - art.f_ref) - assert art.metadata["dof"] == 3 * n - assert art.metadata["integrator"].startswith("velocity-verlet") - assert art.metadata["scheme"] == "int4" - - -def test_to_dict_has_all_keys(): - template = _template() - ref, quant = _warmed_ref_and_quant(template) - d = _run(ref, quant, template).to_dict() - for key in ("pos", "vel", "energy", "f_ref", "f_quant", "df", "metadata"): - assert key in d - - -def test_null_paired_trajectory_zero_df(): - template = _template() - ref = _tiny_potential() - ref(template.clone(), compute_forces=False) - twin = _tiny_potential() - twin(template.clone(), compute_forces=False) - twin.load_state_dict(ref.state_dict()) # identical -> ΔF == 0 - pos0 = template["atoms", "pos"].clone() - vel0 = torch.zeros_like(pos0) - art = build_paired_trajectory( - ref, - twin, - template, - pos0, - vel0, - 4, - dt=0.001, - gamma=1.0, - kbt=0.0257, - mass=12.0, - seed=1, - ) - assert art.df.abs().max().item() == pytest.approx(0.0, abs=1e-10) - - -def test_force_seam_tracks_live_geometry(): - """PotentialForceField must recompute edge geometry from the live positions, - not a frozen template ``edge_diff`` — regression for the constant-PES bug - where swapping only ``pos`` left energy/force pinned to the initial geometry. - """ - template = _template() - ref = _tiny_potential() - ref(template.clone(), compute_forces=False) # warmup lazy params - ff = PotentialForceField(ref, template) - pos0 = template["atoms", "pos"] - out0 = ff(pos0) - torch.manual_seed(1) - pos1 = pos0 + 0.3 * torch.randn_like(pos0) # non-rigid displacement - out1 = ff(pos1) - assert (out1.energy - out0.energy).abs().item() > 1e-6, "energy frozen at initial geometry" - assert (out1.forces - out0.forces).abs().max().item() > 1e-6, ( - "forces frozen at initial geometry" - ) - - -def test_paired_trajectory_energy_varies(): - """End-to-end: the PES must be sampled, so energy is not constant over steps.""" - template = _template() - ref, quant = _warmed_ref_and_quant(template) - art = _run(ref, quant, template, n_steps=8) - assert art.energy.std().item() > 1e-9, "energy constant — PES not sampled" - - -@pytest.mark.skipif(not HAS_ASE, reason="ASE not installed") -def test_ase_calculator_matches_force_seam(): - import ase - - template = _template() - ref, quant = _warmed_ref_and_quant(template) - calc = make_pinet_calculator(ref, template) - pos = template["atoms", "pos"] - atoms = ase.Atoms(numbers=[1, 6, 7, 8], positions=pos.cpu().numpy()) - atoms.calc = calc - forces_ase = torch.as_tensor(atoms.get_forces(), dtype=pos.dtype) - out = ref(template.clone(), compute_forces=True) - assert torch.allclose(forces_ase, out["forces"].detach(), atol=1e-4) diff --git a/tests/test_molix/test_md_integrators.py b/tests/test_molix/test_md_integrators.py deleted file mode 100644 index 12433a7..0000000 --- a/tests/test_molix/test_md_integrators.py +++ /dev/null @@ -1,117 +0,0 @@ -"""Tests for the BAOAB Langevin velocity-Verlet integrator component. - -Decoupled from PiNet via :class:`HarmonicForceField` (F = -k x), so energy -conservation (NVE) and equipartition (Langevin) are provable analytically. -Exercises the component API: ForceField → MDState → Integrator. -""" - -import pytest -import torch - -from molix.md import HarmonicForceField, LangevinVerletIntegrator - -_DTYPE = torch.float64 - - -def _ig(k: float = 1.0, **kw): - return LangevinVerletIntegrator(HarmonicForceField(k).to(_DTYPE), **kw) - - -def _total_energy(state, mass, k): - return (0.5 * mass * (state.vel**2).sum() + 0.5 * k * (state.pos**2).sum()).item() - - -def test_nve_conserves_energy(): - torch.manual_seed(0) - k, mass, dt = 1.0, 1.0, 0.01 - pos = torch.randn(8, 3, dtype=_DTYPE) - vel = torch.randn(8, 3, dtype=_DTYPE) - ig = _ig(k, dt=dt, gamma=0.0, kbt=0.0, mass=mass) - state = ig.initial(pos, vel) - e0 = _total_energy(state, mass, k) - drift = 0.0 - for _ in range(2000): - state = ig.advance(state) - drift = max(drift, abs(_total_energy(state, mass, k) - e0) / abs(e0)) - assert drift < 1e-3, f"NVE energy drift too large: {drift}" - - -def test_langevin_reaches_equipartition(): - torch.manual_seed(1) - k, mass, dt, gamma, kbt = 1.0, 1.0, 0.05, 5.0, 2.0 - n = 60 - ig = _ig(k, dt=dt, gamma=gamma, kbt=kbt, mass=mass, seed=7) - state = ig.initial(torch.zeros(n, 3, dtype=_DTYPE), torch.zeros(n, 3, dtype=_DTYPE)) - dof = n * 3 - ke_samples = [] - for i in range(20000): - state = ig.advance(state) - if i >= 10000: # discard equilibration - ke_samples.append((0.5 * mass * (state.vel**2).sum()).item()) - measured_kbt = (sum(ke_samples) / len(ke_samples)) / (0.5 * dof) # =0.5*dof*kbt - assert abs(measured_kbt - kbt) / kbt < 0.05, f"equipartition off: {measured_kbt} vs {kbt}" - - -def test_noise_is_seed_reproducible(): - pos0 = torch.zeros(5, 3, dtype=_DTYPE) - vel0 = torch.zeros(5, 3, dtype=_DTYPE) - - def run(seed): - ig = _ig(1.0, dt=0.05, gamma=3.0, kbt=1.0, mass=1.0, seed=seed) - return ig.run(pos0.clone(), vel0.clone(), 50)["pos"] - - assert torch.allclose(run(42), run(42)) # same seed -> identical - assert not torch.allclose(run(42), run(43)) # different seed -> different - - -def test_run_records_trajectory_shapes(): - ig = _ig(1.0, dt=0.01, gamma=1.0, kbt=1.0, mass=1.0) - out = ig.run(torch.zeros(4, 3, dtype=_DTYPE), torch.zeros(4, 3, dtype=_DTYPE), 10) - assert out["pos"].shape == (10, 4, 3) - assert out["vel"].shape == (10, 4, 3) - assert out["energy"].shape == (10,) - - -def test_run_matches_manual_advance_loop(): - """``run`` records exactly the states a manual ``initial`` + ``advance`` loop visits.""" - k, mass, dt, gamma, kbt = 1.0, 1.0, 0.05, 3.0, 1.0 - pos0 = torch.randn(6, 3, dtype=_DTYPE) - vel0 = torch.randn(6, 3, dtype=_DTYPE) - - out = _ig(k, dt=dt, gamma=gamma, kbt=kbt, mass=mass, seed=11).run( - pos0.clone(), vel0.clone(), 30 - ) - - ig = _ig(k, dt=dt, gamma=gamma, kbt=kbt, mass=mass, seed=11) - state = ig.initial(pos0.clone(), vel0.clone()) - for i in range(30): - state = ig.advance(state) - assert torch.equal(out["pos"][i], state.pos) - assert torch.equal(out["vel"][i], state.vel) - - -class _CountingHarmonic(HarmonicForceField): - def __init__(self, k: float = 1.0): - super().__init__(k) - self.calls = 0 - - def forward(self, pos): - self.calls += 1 - return super().forward(pos) - - -def test_force_caching_one_eval_per_step(): - """Force caching: exactly one force-field evaluation per step (+1 to seed).""" - ff = _CountingHarmonic(1.0).to(_DTYPE) - ig = LangevinVerletIntegrator(ff, dt=0.05, gamma=1.0, kbt=1.0, mass=1.0, seed=1) - state = ig.initial(torch.zeros(5, 3, dtype=_DTYPE), torch.zeros(5, 3, dtype=_DTYPE)) - for _ in range(10): - state = ig.advance(state) - assert ff.calls == 11 # 1 (initial) + 10 (one per step) - - -def test_mass_must_be_positive(): - with pytest.raises(ValueError, match="strictly positive"): - _ig(1.0, dt=0.01, gamma=0.0, kbt=0.0, mass=-1.0) - with pytest.raises(ValueError, match="strictly positive"): - _ig(1.0, dt=0.01, gamma=0.0, kbt=0.0, mass=torch.tensor([1.0, -2.0, 3.0])) diff --git a/tests/test_molix/test_md_runner.py b/tests/test_molix/test_md_runner.py deleted file mode 100644 index 41bb953..0000000 --- a/tests/test_molix/test_md_runner.py +++ /dev/null @@ -1,134 +0,0 @@ -"""Tests for the hook-driven MD runner (molix.md.runner). - -The integrator is an analytic harmonic well (F = -k x) so the runner can be -exercised without PiNet. Coverage: hook lifecycle ordering + global_step -advance, physics handed via the ``outputs`` channel (never the reserved state -namespaces), and the reference :class:`TrajectoryHook` capture / persistence. -""" - -import torch - -from molix.core.hook import BaseHook -from molix.md import HarmonicForceField, LangevinVerletIntegrator, MDRunner, TrajectoryHook - -_DTYPE = torch.float64 - - -def _integrator(seed: int = 0): - return LangevinVerletIntegrator( - HarmonicForceField(1.0).to(_DTYPE), dt=0.01, gamma=1.0, kbt=1.0, mass=1.0, seed=seed - ) - - -class _RecordingHook(BaseHook): - """Records lifecycle call order and the per-step outputs payload.""" - - def __init__(self): - self.events: list[str] = [] - self.steps: list[int] = [] - self.outputs: list[dict] = [] - - def on_train_start(self, trainer, state): - self.events.append("start") - - def on_train_batch_end(self, trainer, state, batch, outputs): - self.events.append("batch") - self.steps.append(state["global_step"]) - self.outputs.append(outputs) - - def on_train_end(self, trainer, state): - self.events.append("end") - - -def test_runner_fires_lifecycle_and_advances_step(): - rec = _RecordingHook() - runner = MDRunner(_integrator(), mass=1.0, hooks=[rec]) - out = runner.run(torch.zeros(4, 3, dtype=_DTYPE), torch.zeros(4, 3, dtype=_DTYPE), 5) - - assert rec.events == ["start", "batch", "batch", "batch", "batch", "batch", "end"] - assert rec.steps == [1, 2, 3, 4, 5] - assert out["state"]["global_step"] == 5 - assert out["pos"].shape == (4, 3) - assert out["force"].shape == (4, 3) - - -def test_runner_outputs_carry_physics_not_state_namespaces(): - rec = _RecordingHook() - runner = MDRunner(_integrator(), mass=1.0, hooks=[rec]) - runner.run(torch.randn(4, 3, dtype=_DTYPE), torch.randn(4, 3, dtype=_DTYPE), 3) - - last = rec.outputs[-1] - for key in ("pos", "vel", "forces", "potential", "kinetic", "total", "temperature"): - assert key in last - assert torch.allclose(last["total"], last["potential"] + last["kinetic"]) - # Physics must NOT leak into the reserved TrainState namespaces. - for ns in ("train", "eval", "performance", "gpu"): - assert not runner.state.get(ns) - - -def test_runner_temperature_is_positive_under_thermostat(): - rec = _RecordingHook() - runner = MDRunner(_integrator(seed=2), mass=1.0, hooks=[rec]) - runner.run(torch.zeros(20, 3, dtype=_DTYPE), torch.zeros(20, 3, dtype=_DTYPE), 50) - temps = torch.stack([o["temperature"] for o in rec.outputs]) - assert (temps > 0).all() - - -def test_trajectory_hook_persists_pt_and_xyz(tmp_path): - numbers = torch.tensor([1, 6, 8, 1]) - out_pt = tmp_path / "traj.pt" - hook = TrajectoryHook(out_pt, stride=2, numbers=numbers, write_xyz=True) - runner = MDRunner(_integrator(), mass=1.0, hooks=[hook]) - runner.run(torch.zeros(4, 3, dtype=_DTYPE), torch.zeros(4, 3, dtype=_DTYPE), 10) - - assert out_pt.exists() - payload = torch.load(out_pt, weights_only=False) - # stride=2 over 10 steps (global_step 1..10): kept at steps 2,4,6,8,10 -> 5 frames. - assert payload["pos"].shape == (5, 4, 3) - assert payload["forces"].shape == (5, 4, 3) - assert payload["temp"].shape == (5,) - assert payload["stride"] == 2 - assert torch.equal(payload["Z"], numbers) - - xyz = out_pt.with_suffix(".xyz") - assert xyz.exists() - lines = xyz.read_text().splitlines() - # Each extended-XYZ frame: 1 count line + 1 comment + N atom lines. - assert lines[0] == "4" - assert "Etot=" in lines[1] and "T=" in lines[1] - assert lines[2].split()[0] == "H" # Z=1 -> H - - -def test_trajectory_hook_shard_flush_matches_single_buffer(tmp_path): - """Shard-flushing (bounded host memory) yields the same trajectory as one - in-memory buffer, and cleans up its shard files.""" - p0 = torch.zeros(3, 3, dtype=_DTYPE) - v0 = torch.randn(3, 3, dtype=_DTYPE) - - big = tmp_path / "big.pt" # one buffer (flush_every >> frames) - MDRunner(_integrator(seed=7), mass=1.0, hooks=[TrajectoryHook(big, flush_every=1000)]).run( - p0.clone(), v0.clone(), 10 - ) - small = tmp_path / "small.pt" # flush every 2 frames -> 5 shards - MDRunner(_integrator(seed=7), mass=1.0, hooks=[TrajectoryHook(small, flush_every=2)]).run( - p0.clone(), v0.clone(), 10 - ) - - a = torch.load(big, weights_only=True) - b = torch.load(small, weights_only=True) - assert b["pos"].shape == (10, 3, 3) - for key in ("pos", "vel", "forces", "pe", "temp"): - assert torch.equal(a[key], b[key]), key - assert not list(tmp_path.glob("small.part*.pt")) # shards removed - - -def test_trajectory_hook_without_numbers_skips_xyz(tmp_path): - out_pt = tmp_path / "traj.pt" - hook = TrajectoryHook(out_pt, stride=1, write_xyz=True) # no numbers -> no xyz - runner = MDRunner(_integrator(), mass=1.0, hooks=[hook]) - runner.run(torch.zeros(3, 3, dtype=_DTYPE), torch.zeros(3, 3, dtype=_DTYPE), 4) - - assert out_pt.exists() - assert not out_pt.with_suffix(".xyz").exists() - payload = torch.load(out_pt, weights_only=False) - assert payload["pos"].shape == (4, 3, 3) diff --git a/tests/test_molix/test_nn/__init__.py b/tests/test_molix/test_nn/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/test_molix/test_nn/test_init.py b/tests/test_molix/test_nn/test_init.py new file mode 100644 index 0000000..a2d738d --- /dev/null +++ b/tests/test_molix/test_nn/test_init.py @@ -0,0 +1,85 @@ +"""Tests for the ``molix.nn`` public import surface (``src/molix/nn/__init__.py``). + +That file is a pure re-export module, so the only contract it owns is the *set* +of names it publishes and the fact that each one resolves. Two rules are pinned +here: + +* ``NeighborList`` is **gone**. ``molix.nn.locality.NeighborList`` was an + ``nn.Module`` wrapper over :func:`molix.F.locality.get_neighbor_pairs` with + zero in-tree call sites, and it collided by name with two live classes + (:class:`molix.data.tasks.neighbor.NeighborList`, which builds the pipeline's + edge tensors, and :class:`molix.md.NeighborList`). Spec + ``md-neighborlist-skin-02-prune`` deletes it; the name must stay free. +* ``__all__`` stays alphabetized (``.claude/notes/notes.md:243``). The deletion + rewrites the literal anyway, so the sorted order is pinned in the same value. + +Following the one-assertion-family-per-method pattern of +``tests/test_molzoo/test_imports.py``. No physics here — this module computes +nothing, so there is no domain tier. +""" + +from __future__ import annotations + +import importlib.util +import inspect + +import molix.nn + +#: The full public surface, hard-coded rather than derived. Written out as a +#: literal so it pins **two** things at once: that ``NeighborList`` is no longer +#: exported, and that the remaining four names are in alphabetical order. A +#: ``sorted(...)``-based assertion would only catch the second. +EXPECTED_ALL = ["BatchAggregation", "KeyedMLP", "KeyedMLPSpec", "ScatterSum"] + + +class TestNnExports: + """The names ``molix.nn`` publishes, and the ones it must not.""" + + def test___all___is_the_pinned_literal(self) -> None: + """``__all__`` is exactly the four sorted names — value, not just set.""" + assert molix.nn.__all__ == EXPECTED_ALL + + def test_every_export_resolves(self) -> None: + """Each ``__all__`` entry is actually bound on the module. + + Guards the failure mode a re-sorted list invites: an entry kept in + ``__all__`` whose ``from .x import y`` line was dropped, which + ``from molix.nn import *`` would only report at the caller's site. + """ + unresolved: list[str] = [name for name in molix.nn.__all__ if not hasattr(molix.nn, name)] + assert unresolved == [] + + def test_no_stray_public_attribute(self) -> None: + """The public non-module attributes are exactly ``set(__all__)``. + + Sub-module names (``mlp``, ``scatter``) are bound on the package as a + side effect of ``from .x import y`` and are filtered with + :func:`inspect.ismodule`; ``molix.nn`` defines no ``__dir__``, so a bare + ``set(dir(molix.nn)) == set(molix.nn.__all__)`` would fail on those and + be a false red. What is left after the filter is the surface a caller + can bind, and it must not exceed what is declared. + """ + public_attributes = { + name + for name in vars(molix.nn) + if not name.startswith("_") and not inspect.ismodule(getattr(molix.nn, name)) + } + assert public_attributes == set(molix.nn.__all__) + + def test_neighborlist_not_reintroduced(self) -> None: + """Neither ``__all__`` nor the module namespace carries ``NeighborList``. + + Two distinct regressions: re-adding the export trips the first clause, + while a bare ``from .locality import NeighborList`` with no ``__all__`` + edit trips only the second. + """ + assert "NeighborList" not in molix.nn.__all__ + assert not hasattr(molix.nn, "NeighborList") + + def test_locality_module_gone(self) -> None: + """``molix.nn.locality`` is not importable — the file itself is gone. + + The name check above passes for a module that still ships but is merely + unexported; this asserts the deletion, not just the hidden re-export. + """ + assert importlib.util.find_spec("molix.nn.locality") is None diff --git a/tests/test_molix/test_nn/test_mlp.py b/tests/test_molix/test_nn/test_mlp.py index f4daeb1..459270a 100644 --- a/tests/test_molix/test_nn/test_mlp.py +++ b/tests/test_molix/test_nn/test_mlp.py @@ -1,11 +1,29 @@ """Unit tests for KeyedMLP with dict-based inputs.""" +from collections.abc import Iterator + import pytest import torch +from molix import config from molix.nn.mlp import KeyedMLP, KeyedMLPSpec +@pytest.fixture +def fp64() -> Iterator[None]: + """Run the case under the global fp64 precision, restoring the previous one. + + :class:`KeyedMLP` bakes ``config["ftype"]`` into its parameters at + construction time (the contract documented in :mod:`molix.config`), so + the precision has to be switched *before* the module is built and + handed back afterwards. + """ + previous = config["ftype"] + config.set_precision("fp64") + yield + config.set_precision("fp64" if previous == torch.float64 else "fp32") + + class TestKeyedMLPSpec: def test_spec_basic(self): spec = KeyedMLPSpec( @@ -171,3 +189,42 @@ def test_serialization_roundtrip(self): out2 = restored(dict(sample)) assert out1["edge_weights"].shape == out2["edge_weights"].shape + + @pytest.mark.parametrize( + "hidden_dims", + [[32], [32, 16], [8, 8, 8]], + ids=["one-hidden", "two-hidden", "three-hidden"], + ) + def test_all_parameters_honour_the_fp64_precision(self, fp64, hidden_dims): + """Every parameter is fp64 when the MLP is built under fp64. + + ``config["ftype"]`` is the single source of truth for the working + precision; a layer that ignores it leaves the module + mixed-precision and its first forward dies on a dtype-mismatched + matmul. The depth sweep covers the first / intermediate / final + linear layers, which are constructed on three separate code paths. + """ + mlp = KeyedMLP( + input_key="rbf", + output_key="weights", + in_dim=8, + hidden_dims=hidden_dims, + out_dim=4, + ) + + assert {p.dtype for p in mlp.parameters()} == {torch.float64} + + def test_forward_runs_under_the_fp64_precision(self, fp64): + """A module built at fp64 consumes fp64 inputs and emits fp64.""" + mlp = KeyedMLP( + input_key="rbf", + output_key="weights", + in_dim=8, + hidden_dims=[32, 16], + out_dim=4, + ) + + out = mlp({"rbf": torch.ones(10, 8, dtype=torch.float64)}) + + assert out["weights"].shape == (10, 4) + assert out["weights"].dtype == torch.float64 diff --git a/tests/test_molix/test_profiler/test_dataset.py b/tests/test_molix/test_profiler/test_dataset.py new file mode 100644 index 0000000..669da06 --- /dev/null +++ b/tests/test_molix/test_profiler/test_dataset.py @@ -0,0 +1,243 @@ +"""Tests for :class:`molix.profiler.dataset.DatasetProfiler`. + +Every fixture is built from **literal** sample dicts (no RNG), so the sizes +asserted below are analytic constants of the fixture, not measured values. +The only timing assertions are ``> 0`` liveness checks — no wall-clock +magnitude is pinned. +""" + +from __future__ import annotations + +from pathlib import Path + +import pytest +import torch + +from molix.data.cache import PackedCache, _flatten +from molix.data.dataset import CachedDataset +from molix.profiler.dataset import ( + DatasetProfiler, + DatasetResult, + FieldSpec, + TargetStat, + _flatten_leaves, +) + +# --------------------------------------------------------------------------- +# Literal fixtures — no RNG, so every expectation below is analytic +# --------------------------------------------------------------------------- + +_N_RECORDS = 100 +_ATOM_MEAN = 3.5 # atom counts cycle 2,3,4,5 → 350 atoms over 100 records +_EDGE_MEAN = 7.0 # bidirectional ring → 2N edges per record → 700 total +_MAX_ATOMS = 5 +_MAX_EDGES = 10 +_AVG_NUM_NEIGHBORS = 2.0 # 700 edges / 350 atoms, exactly + + +def _record(i: int, n: int | None = None) -> dict: + """One literal sample dict: a ring molecule of ``n`` (default ``2 + i % 4``) atoms. + + Keys follow the raw-sample tier of the two-tier data contract: + ``Z`` ``(N,)``, ``pos`` ``(N, 3)``, ``edge_index`` ``(2N, 2)``, + ``edge_dist`` ``(2N,)``, and the scalar target ``targets.U0`` ``(1,)``. + + Args: + i: Record index; drives the atom count and the target value. + n: Fixed atom count, overriding the ``2 + i % 4`` cycle. + + Returns: + A flat sample dict suitable for :meth:`PackedCache.save`. + """ + n_atoms = 2 + i % 4 if n is None else n + src = torch.arange(n_atoms, dtype=torch.long) + dst = (src + 1) % n_atoms + edge_index = torch.cat([torch.stack([src, dst], dim=1), torch.stack([dst, src], dim=1)], dim=0) + return { + "Z": torch.full((n_atoms,), 6, dtype=torch.long), + "pos": torch.arange(3 * n_atoms, dtype=torch.float32).reshape(n_atoms, 3), + "edge_index": edge_index, + "edge_dist": torch.full((2 * n_atoms,), 1.5, dtype=torch.float32), + "targets": {"U0": torch.tensor([float(i)], dtype=torch.float32)}, + } + + +def _write_cache(tmp_path: Path, samples: list[dict]) -> Path: + """Pack *samples* into a cache file under *tmp_path* and return the sink path.""" + sink = tmp_path / "cache.pt" + PackedCache(sink).save(samples) + return sink + + +class _CountingCachedDataset(CachedDataset): + """:class:`CachedDataset` that counts ``__getitem__`` calls. + + Subclassing keeps the packed-pointer fast path (``atom_counts`` / + ``max_atoms`` / ``packed_view``) intact, so the counter observes only + the profiler's sampled slow path. + """ + + def __init__(self, sink: Path) -> None: + super().__init__(sink) + self.n_getitem = 0 + + def __getitem__(self, idx: int) -> dict: + self.n_getitem += 1 + return super().__getitem__(idx) + + +@pytest.fixture +def ring_dataset(tmp_path: Path) -> CachedDataset: + """Cache-backed dataset over the 100 literal ring records.""" + return CachedDataset(_write_cache(tmp_path, [_record(i) for i in range(_N_RECORDS)])) + + +# --------------------------------------------------------------------------- +# DatasetProfiler +# --------------------------------------------------------------------------- + + +class TestDatasetProfiler: + """Config in ``__init__``, data in ``run()``; exact cache stats + sampled access.""" + + def test_run_uses_exact_cache_counts(self, ring_dataset): + """Size stats come from the packed pointers — all 100 records, not the 5 sampled.""" + result = DatasetProfiler(n_samples=5).run(ring_dataset) + + assert isinstance(result, DatasetResult) + assert result.counts_exact is True + assert result.n_total == _N_RECORDS + assert result.max_atoms == ring_dataset.max_atoms == _MAX_ATOMS + assert result.max_edges == ring_dataset.max_edges == _MAX_EDGES + assert result.avg_num_neighbors == ring_dataset.avg_num_neighbors + assert result.avg_num_neighbors == pytest.approx(_AVG_NUM_NEIGHBORS, rel=1e-12) + assert result.atom_stats.mean == float(ring_dataset.atom_counts.double().mean()) + assert result.atom_stats.mean == pytest.approx(_ATOM_MEAN, rel=1e-12) + assert result.edge_stats is not None + assert result.edge_stats.mean == pytest.approx(_EDGE_MEAN, rel=1e-12) + + def test_run_samples_only_n_samples_for_access(self, tmp_path): + """The access loop is bounded by ``n_samples`` (+ warmup), never ``n_total``.""" + ds = _CountingCachedDataset(_write_cache(tmp_path, [_record(i) for i in range(_N_RECORDS)])) + + result = DatasetProfiler(n_samples=5, n_warmup=3).run(ds) + + assert result.n_sampled == 5 + assert result.cold_access_ms > 0.0 + assert result.access_ms.mean_ms > 0.0 + # at most one cold access + 3 discarded warmups + 5 measured accesses + assert 5 <= ds.n_getitem <= 1 + 3 + 5 + + def test_run_fields_come_from_packed_schema(self, ring_dataset): + """Field layout is read off ``payload["schema"]`` verbatim, hence exact.""" + result = DatasetProfiler(n_samples=5).run(ring_dataset) + schema = ring_dataset.packed_view().payload["schema"] + + assert result.fields_exact is True + assert all(isinstance(f, FieldSpec) for f in result.fields) + assert {(f.key, f.axis, f.dtype, f.extra_shape) for f in result.fields} == { + (key, *spec) for key, spec in schema.items() + } + + def test_run_on_plain_sequence_falls_back(self): + """A ``list[dict]`` has no packed pointers: degrade audibly, never raise.""" + records = [_record(i, n=4) for i in range(10)] + + result = DatasetProfiler(n_samples=5).run(records) + + assert isinstance(result, DatasetResult) + assert result.counts_exact is False + assert result.fields_exact is False + assert result.n_total == 10 + assert result.n_sampled == 5 + assert result.atom_stats.mean == pytest.approx(4.0, rel=1e-12) + assert result.fields # inferred from the sampled records + + def test_run_without_edge_ptr_warns_not_raises(self, tmp_path): + """A cache built without per-edge keys degrades to ``edge_stats=None`` + [WARN].""" + samples = [ + { + "Z": torch.full((4,), 8, dtype=torch.long), + "pos": torch.arange(12, dtype=torch.float32).reshape(4, 3), + } + for _ in range(6) + ] + ds = CachedDataset(_write_cache(tmp_path, samples)) + with pytest.raises(ValueError): + _ = ds.edge_counts # precondition: the exact fast path really does raise + + result = DatasetProfiler(n_samples=3).run(ds) + + assert result.edge_stats is None + assert result.warnings + + def test_run_rejects_empty_dataset(self): + """Empty input is a user error, and the message must say what to pass instead.""" + with pytest.raises(ValueError, match="CachedDataset"): + DatasetProfiler().run([]) + + def test_init_rejects_degenerate_config(self): + """``n_samples <= 0`` / ``stride <= 0`` are rejected at construction.""" + with pytest.raises(ValueError, match="n_samples"): + DatasetProfiler(n_samples=0) + with pytest.raises(ValueError, match="stride"): + DatasetProfiler(stride=0) + + +# --------------------------------------------------------------------------- +# DatasetResult +# --------------------------------------------------------------------------- + + +class TestDatasetResult: + """The report is a sectioned frame; diagnostics are ``[WARN]`` lines, never raises.""" + + def test_print_report_sections(self, ring_dataset, capsys): + """All five sections and the 72-char rule are present.""" + DatasetProfiler(n_samples=5).run(ring_dataset).print_report() + + out = capsys.readouterr().out + assert "─" * 72 in out + for section in ("Size", "Access", "Footprint", "Fields", "Targets"): + assert section in out + + def test_print_report_flags_nonfinite_targets(self, tmp_path, capsys): + """A NaN target column is counted and warned about, not raised on.""" + samples = [_record(i, n=4) for i in range(8)] + samples[3]["targets"]["U0"] = torch.tensor([float("nan")], dtype=torch.float32) + ds = CachedDataset(_write_cache(tmp_path, samples)) + + result = DatasetProfiler(n_samples=8).run(ds) + result.print_report() + + assert all(isinstance(t, TargetStat) for t in result.targets) + assert any(t.n_nonfinite == 1 for t in result.targets) + # uniform sizes + cache-backed input ⇒ the NaN column is the only warning source + assert "[WARN]" in capsys.readouterr().out + + +# --------------------------------------------------------------------------- +# _flatten_leaves +# --------------------------------------------------------------------------- + + +def test_flatten_leaves_does_not_raise_on_reserved_key(): + """Deliberate divergence from :func:`molix.data.cache._flatten`, which validates. + + The packing-time flattener rejects a sample key colliding with a reserved + packed-cache key; the profiler's read-only walker must survive any malformed + sample, because diagnostics never raise. + """ + sample = { + "schema": torch.zeros(2), # collides with a reserved packed-cache key + "Z": torch.ones(2, dtype=torch.long), + "targets": {"U0": torch.tensor([1.0])}, + } + + with pytest.raises(ValueError): + _flatten(sample) + + leaves = _flatten_leaves(sample) + + assert set(leaves) == {"schema", "Z", "targets.U0"} + assert leaves["targets.U0"] is sample["targets"]["U0"] diff --git a/tests/test_molix/test_profiler/test_mock.py b/tests/test_molix/test_profiler/test_mock.py new file mode 100644 index 0000000..2cc08c9 --- /dev/null +++ b/tests/test_molix/test_profiler/test_mock.py @@ -0,0 +1,53 @@ +"""Tests for :mod:`molix.profiler.mock` — the ``seed`` reproducibility promise. + +Both generators document ``seed`` as making their output reproducible. That +must cover the *drawn shapes* (atom / edge counts), not only the tensor +values: any golden built on mock data is otherwise unstable within a single +process. +""" + +from __future__ import annotations + +from tensordict import TensorDict + +from molix.profiler import MockBatch, MockSource + + +def _shapes(batch: TensorDict) -> tuple[tuple[int, ...], ...]: + """Shape signature of one mock batch. + + Args: + batch: Nested batch from :meth:`MockBatch.__call__`. + + Returns: + Shapes of ``atoms.Z`` ``(N,)``, ``atoms.pos`` ``(N, 3)``, + ``edges.edge_index`` ``(E, 2)`` and ``graphs.num_atoms`` ``(B,)``. + """ + return ( + tuple(batch["atoms", "Z"].shape), + tuple(batch["atoms", "pos"].shape), + tuple(batch["edges", "edge_index"].shape), + tuple(batch["graphs", "num_atoms"].shape), + ) + + +class TestMockSource: + """``seed`` fixes the per-sample atom counts, not just the tensor values.""" + + def test_seed_makes_atom_counts_reproducible(self): + """Two same-seed sources built in one process yield identical atom counts.""" + first = MockSource(n_samples=8, n_atoms=(5, 20), seed=0) + second = MockSource(n_samples=8, n_atoms=(5, 20), seed=0) + + assert [len(first[i]["Z"]) for i in range(8)] == [len(second[i]["Z"]) for i in range(8)] + + +class TestMockBatch: + """``seed`` fixes the drawn shapes of every call, not just the tensor values.""" + + def test_seed_makes_shapes_reproducible(self): + """Two same-seed factories emit the same shape sequence over repeated calls.""" + first = MockBatch(n_atoms=(8, 32), n_edges=(16, 64), seed=0) + second = MockBatch(n_atoms=(8, 32), n_edges=(16, 64), seed=0) + + assert [_shapes(first()) for _ in range(4)] == [_shapes(second()) for _ in range(4)] diff --git a/tests/test_molix/test_profiler/test_trainer_profiler.py b/tests/test_molix/test_profiler/test_trainer_profiler.py index 742d37b..bc9f439 100644 --- a/tests/test_molix/test_profiler/test_trainer_profiler.py +++ b/tests/test_molix/test_profiler/test_trainer_profiler.py @@ -65,3 +65,24 @@ def test_custom_batch_size_respected(self): result = TrainerProfiler(device="cpu").run(n_steps=100, n_warmup=5, batch=batch, top=5) assert result.steps_per_sec > 0 assert not torch.cuda.is_available() or result.device == "cpu" + + def test_data_description_derived_from_supplied_batch(self): + """A caller-supplied batch must not be reported as the default MockBatch. + + ``run(batch=...)`` skips default construction, so the description has + to come from the supplied batch itself. Wording is the implementer's + choice; the contract pinned here is (a) the default literal is gone and + (b) the real atom count is visible in the report line. + """ + batch = MockBatch(n_atoms=7, n_edges=9, n_graphs=3, device="cpu", seed=0)() + result = TrainerProfiler(device="cpu").run(n_steps=5, n_warmup=0, batch=batch, top=3) + desc = result.data_description + assert "MockBatch(n_atoms=32" not in desc # default literal must not leak + assert "7" in desc # the supplied batch's actual atom count + + def test_data_description_reports_default_batch_shape(self): + """With ``batch=None`` the description still describes the built default.""" + result = TrainerProfiler(device="cpu").run(n_steps=5, n_warmup=0, top=3) + desc = result.data_description + assert "32" in desc # default n_atoms + assert "128" in desc # default n_edges diff --git a/tests/test_molix/test_quant.py b/tests/test_molix/test_quant.py index f18592a..d731093 100644 --- a/tests/test_molix/test_quant.py +++ b/tests/test_molix/test_quant.py @@ -191,6 +191,8 @@ def _tiny_potential() -> PiNetPotential: depth=2, rank=3, hidden_dim=16, + # Monomorphic since b85d12f: force derivation is fixed here. + compute_forces=True, ) .to(_DEVICE) .eval() @@ -217,8 +219,8 @@ def test_null_control_identical_weights_zero_delta(): model = _tiny_potential() twin = _tiny_potential() batch = _tiny_batch() - model(batch.clone(), compute_forces=False) # materialise lazy params - twin(batch.clone(), compute_forces=False) + model(batch.clone()) # materialise lazy params + twin(batch.clone()) twin.load_state_dict(model.state_dict()) # identical weights s = ForceDelta.between(model, twin, batch).summary() assert s["F_rms"] == pytest.approx(0.0, abs=1e-10) @@ -229,8 +231,8 @@ def test_int4_ptq_perturbs_forces(): model = _tiny_potential() quant = _tiny_potential() batch = _tiny_batch() - model(batch.clone(), compute_forces=False) - quant(batch.clone(), compute_forces=False) + model(batch.clone()) + quant(batch.clone()) quant.load_state_dict(Quantizer("int4").quantize_state_dict(model.state_dict())) s = ForceDelta.between(model, quant, batch).summary() assert s["F_rms"] > 0.0 diff --git a/tests/test_molix/test_schema.py b/tests/test_molix/test_schema.py new file mode 100644 index 0000000..9f59f5b --- /dev/null +++ b/tests/test_molix/test_schema.py @@ -0,0 +1,40 @@ +"""Tests for molix.schema — the post-collate batch-schema key contract.""" + +import torch +from tensordict import TensorDict + +from molix.schema import ATOMIC_ENERGY_KEY, ENERGY_KEY, FORCES_KEY, POS_KEY, has_energy, has_forces + + +def _batch(*, energy: bool = False, forces: bool = False) -> TensorDict: + atoms = {"pos": torch.zeros(4, 3)} + if forces: + atoms["forces"] = torch.zeros(4, 3) + td = TensorDict({"atoms": TensorDict(atoms, batch_size=[4])}, batch_size=[]) + if energy: + td["graphs"] = TensorDict({"energy": torch.zeros(1)}, batch_size=[1]) + return td + + +def test_keys_match_the_documented_schema(): + """The tuple keys are the CLAUDE.md post-collate contract, verbatim.""" + assert ENERGY_KEY == ("graphs", "energy") + assert ATOMIC_ENERGY_KEY == ("atoms", "energy") + assert FORCES_KEY == ("atoms", "forces") + assert POS_KEY == ("atoms", "pos") + + +def test_has_energy_and_has_forces_walk_the_nesting(): + assert not has_energy(_batch()) + assert not has_forces(_batch()) + assert has_energy(_batch(energy=True)) + assert has_forces(_batch(forces=True)) + + +def test_molpot_protocol_reexports_the_same_objects(): + """molpot.derivation.protocol must alias, not restate, the schema keys.""" + from molpot.derivation import protocol + + assert protocol.ENERGY_KEY is ENERGY_KEY + assert protocol.FORCES_KEY is FORCES_KEY + assert protocol.has_forces is has_forces diff --git a/tests/test_molix/test_units.py b/tests/test_molix/test_units.py new file mode 100644 index 0000000..e6d4993 --- /dev/null +++ b/tests/test_molix/test_units.py @@ -0,0 +1,24 @@ +"""Tests for molix.units — the single source of physical constants.""" + +from molix.units import DEAD_EDGE_CUTOFF_FACTOR, EV_PER_AMU_A2_FS2, KB_AMU_A_FS, KB_EV_PER_K + + +def test_boltzmann_constant_is_codata(): + assert KB_EV_PER_K == 8.617333262e-5 + + +def test_kb_amu_a_fs_is_derived_not_restated(): + """The (amu, Å, fs) k_B must be the eV/K value through the unit bridge.""" + assert KB_AMU_A_FS == KB_EV_PER_K / EV_PER_AMU_A2_FS2 + + +def test_dead_edge_factor_clears_any_cutoff_envelope(): + """A dead edge must land strictly outside the cutoff, with margin.""" + assert DEAD_EDGE_CUTOFF_FACTOR > 1.0 + + +def test_quant_shares_the_single_source(): + """molix.quant's class attribute must alias molix.units, not restate it.""" + from molix.quant import EffectiveTemperature + + assert EffectiveTemperature.KB_EV_PER_K == KB_EV_PER_K diff --git a/tests/test_molpot/test_composition/__init__.py b/tests/test_molpot/test_composition/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/test_molpot/test_composition/conftest.py b/tests/test_molpot/test_composition/conftest.py index aeb2ef6..eaddc7d 100644 --- a/tests/test_molpot/test_composition/conftest.py +++ b/tests/test_molpot/test_composition/conftest.py @@ -31,7 +31,7 @@ from tensordict import TensorDict from molix.config import config -from molpot.composition import Sonata, build_sonata +from molpot.composition import Sonata from molzoo import Allegro # --------------------------------------------------------------------------- @@ -127,7 +127,7 @@ def sonata_pipeline() -> Sonata: avg_num_neighbors=12.0, expose_tensor_track=True, ) - sonata = build_sonata( + sonata = Sonata.from_encoder( encoder, sigma=1.0, dl=2.0, @@ -137,11 +137,6 @@ def sonata_pipeline() -> Sonata: constrain_total_charge=True, avg_num_neighbors=12.0, ) - # cuequivariance_torch.Linear ignores `config.ftype` and creates - # float32 weights regardless; cast the whole module tree to - # float64 after construction so the head's `cuet.Linear` - # collapse paths line up with the float64 batch inputs. - sonata = sonata.double() sonata.eval() yield sonata finally: diff --git a/tests/test_molpot/test_composition/test_byteff_heads.py b/tests/test_molpot/test_composition/test_byteff_heads.py index c6b7432..cc20bc7 100644 --- a/tests/test_molpot/test_composition/test_byteff_heads.py +++ b/tests/test_molpot/test_composition/test_byteff_heads.py @@ -2,12 +2,16 @@ from __future__ import annotations +from collections.abc import Iterator + import pytest import torch +from molix import config from molpot.composition.heads import ( ChargeHead, ChargeTransferParameterHead, + LJParameterHead, RepulsionParameterHead, TSScalingHead, ) @@ -24,6 +28,39 @@ def batch(): return torch.tensor([0, 0, 0, 1, 1], dtype=torch.long) +@pytest.fixture +def fp64() -> Iterator[None]: + """Run the case under the global fp64 precision, restoring the previous one. + + Every head in :mod:`molpot.composition.heads` bakes ``config["ftype"]`` + into its parameters at construction time (the contract documented in + :mod:`molix.config`), so the precision has to be switched *before* the + head is built and handed back afterwards. + """ + previous = config["ftype"] + config.set_precision("fp64") + yield + config.set_precision("fp64" if previous == torch.float64 else "fp32") + + +# --------------------------------------------------------------------------- +# LJParameterHead +# --------------------------------------------------------------------------- + + +class TestLJParameterHead: + def test_all_parameters_honour_the_fp64_precision(self, fp64): + """Every parameter is fp64 when the head is built under fp64. + + ``config["ftype"]`` is the single source of truth for the working + precision; a layer that ignores it leaves the head mixed-precision + and its first forward dies on a dtype-mismatched matmul. + """ + head = LJParameterHead(feature_dim=16) + + assert {p.dtype for p in head.parameters()} == {torch.float64} + + # --------------------------------------------------------------------------- # RepulsionParameterHead # --------------------------------------------------------------------------- @@ -50,6 +87,12 @@ def test_min_floor(self): assert torch.all(out["eps_rep"] >= 0.5) assert torch.all(out["lam_rep"] >= 0.3) + def test_all_parameters_honour_the_fp64_precision(self, fp64): + """Every parameter is fp64 when the head is built under fp64.""" + head = RepulsionParameterHead(feature_dim=16) + + assert {p.dtype for p in head.parameters()} == {torch.float64} + # --------------------------------------------------------------------------- # ChargeTransferParameterHead @@ -71,6 +114,12 @@ def test_outputs_positive(self, node_features): assert torch.all(out["eps_ct"] > 0) assert torch.all(out["lam_ct"] > 0) + def test_all_parameters_honour_the_fp64_precision(self, fp64): + """Every parameter is fp64 when the head is built under fp64.""" + head = ChargeTransferParameterHead(feature_dim=16) + + assert {p.dtype for p in head.parameters()} == {torch.float64} + # --------------------------------------------------------------------------- # ChargeHead @@ -110,6 +159,21 @@ def test_grad_flows(self, batch): out["charge"].sum().backward() assert x.grad is not None + def test_all_parameters_honour_the_fp64_precision(self, fp64): + """Every parameter is fp64 when the head is built under fp64.""" + head = ChargeHead(feature_dim=16) + + assert {p.dtype for p in head.parameters()} == {torch.float64} + + def test_forward_runs_under_the_fp64_precision(self, fp64, batch): + """A head built at fp64 conserves charge in fp64 arithmetic.""" + head = ChargeHead(feature_dim=16, total_charge=1.0) + + out = head(torch.ones(5, 16, dtype=torch.float64), batch=batch) + + assert out["charge"].shape == (5,) + assert out["charge"].dtype == torch.float64 + # --------------------------------------------------------------------------- # TSScalingHead @@ -149,6 +213,24 @@ def test_buffers_on_device(self, ts_head): assert ts_head.alpha_free is not None assert ts_head.r_star_free is not None + def test_all_parameters_and_buffers_honour_the_fp64_precision(self, fp64): + """Every parameter is fp64 when the head is built under fp64. + + The free-atom reference tables are caller-supplied buffers, so they + follow the dtype of the tensors handed in; passing fp64 references + pins the whole module to one precision. + """ + num_elements = 10 + head = TSScalingHead( + feature_dim=16, + c6_free=torch.rand(num_elements, dtype=torch.float64) * 10, + alpha_free=torch.rand(num_elements, dtype=torch.float64) * 5, + r_star_free=torch.rand(num_elements, dtype=torch.float64) * 2 + 1.0, + ) + + assert {p.dtype for p in head.parameters()} == {torch.float64} + assert {b.dtype for b in head.buffers() if b.is_floating_point()} == {torch.float64} + # --------------------------------------------------------------------------- # MultiHead diff --git a/tests/test_molpot/test_composition/test_classical_mm.py b/tests/test_molpot/test_composition/test_classical_mm.py new file mode 100644 index 0000000..ab05c06 --- /dev/null +++ b/tests/test_molpot/test_composition/test_classical_mm.py @@ -0,0 +1,258 @@ +"""Tests for ClassicalMMComposer (learnable-classical-ff-03).""" + +from __future__ import annotations + +import ast +import math +from pathlib import Path + +import torch +from tensordict import TensorDict + +from molpot.composition.classical_mm import ClassicalMMComposer +from molpot.composition.heads import ChargeHead, LJParameterHead +from molpot.composition.mm_heads import ( + AngleParamHead, + BondParamHead, + ImproperParamHead, + ProperTorsionParamHead, +) +from molpot.composition.multihead import MultiHead +from molpot.derivation import ForceDerivation +from molpot.ir import CLASS_I_CANONICAL, BondBag, PotentialIR + + +def _two_atom_bond_batch( + *, + r: float = 1.5, + dtype: torch.dtype = torch.float64, +) -> TensorDict: + pos = torch.tensor([[0.0, 0.0, 0.0], [r, 0.0, 0.0]], dtype=dtype) + return TensorDict( + { + "atoms": TensorDict( + { + "pos": pos, + "Z": torch.tensor([1, 1], dtype=torch.long), + "batch": torch.zeros(2, dtype=torch.long), + }, + batch_size=[2], + ), + "bonds": TensorDict( + { + "atomi": torch.tensor([0], dtype=torch.long), + "atomj": torch.tensor([1], dtype=torch.long), + }, + batch_size=[1], + ), + }, + batch_size=[], + ) + + +class _ConstantBondHead(torch.nn.Module): + """Mock head returning fixed k, r0 (ignores features).""" + + def __init__(self, k: float, r0: float): + super().__init__() + self.k = k + self.r0 = r0 + + def forward(self, features: torch.Tensor) -> dict[str, torch.Tensor]: + n = features.shape[0] + return { + "k": torch.full((n,), self.k, dtype=features.dtype, device=features.device), + "r0": torch.full((n,), self.r0, dtype=features.dtype, device=features.device), + } + + +# --------------------------------------------------------------------------- +# parameterize → PotentialIR (ac-006) +# --------------------------------------------------------------------------- + + +class TestClassicalMMComposerParameterize: + def test_parameterize_builds_class_i_ir(self): + bond_head = BondParamHead(feature_dim=4, hidden_dim=8) + angle_head = AngleParamHead(feature_dim=4, hidden_dim=8) + atom_head = MultiHead( + { + "lj": LJParameterHead(feature_dim=4, hidden_dim=8), + "q": ChargeHead(feature_dim=4, hidden_dim=8), + } + ) + composer = ClassicalMMComposer( + bond_head=bond_head, + angle_head=angle_head, + atom_head=atom_head, + ) + batch = TensorDict( + { + "atoms": TensorDict( + { + "pos": torch.randn(3, 3), + "Z": torch.tensor([6, 1, 1]), + "batch": torch.zeros(3, dtype=torch.long), + }, + batch_size=[3], + ), + "bonds": TensorDict( + { + "atomi": torch.tensor([0, 0]), + "atomj": torch.tensor([1, 2]), + }, + batch_size=[2], + ), + "angles": TensorDict( + { + "atomi": torch.tensor([1]), + "atomj": torch.tensor([0]), + "atomk": torch.tensor([2]), + }, + batch_size=[1], + ), + }, + batch_size=[], + ) + features = { + "bonds": torch.randn(2, 4), + "angles": torch.randn(1, 4), + "atoms": torch.randn(3, 4), + } + ir = composer.parameterize(features, batch) + assert isinstance(ir, PotentialIR) + assert ir.unit_system == "class_i_canonical" + assert ir.bonds is not None + assert ir.bonds.k.shape == (2,) + assert ir.bonds.r0.shape == (2,) + assert ir.angles is not None + assert ir.angles.k.shape == (1,) + assert ir.lj is not None + assert ir.lj.epsilon.shape == (3,) + assert ir.charges is not None + assert ir.charges.q.shape == (3,) + assert ir.scaling is not None + # Units tag consistency + assert CLASS_I_CANONICAL["energy"] == "kcal/mol" + + +# --------------------------------------------------------------------------- +# Energy golden (ac-007) +# --------------------------------------------------------------------------- + + +class TestClassicalMMComposerEnergy: + def test_bond_harmonic_analytic_golden(self): + k, r0, r = 2.0, 1.0, 1.5 + expected = 0.5 * k * (r - r0) ** 2 # 0.25 + composer = ClassicalMMComposer(bond_head=_ConstantBondHead(k=k, r0=r0)) + batch = _two_atom_bond_batch(r=r) + features = {"bonds": torch.zeros(1, 1, dtype=torch.float64)} + ir = composer.parameterize(features, batch) + energy = composer.energy(ir, batch, pos=batch["atoms", "pos"]) + assert energy.shape == () + assert math.isclose(float(energy), expected, rel_tol=1e-5, abs_tol=1e-5) + + def test_forward_chains_parameterize_and_energy(self): + k, r0, r = 4.0, 1.2, 1.4 + expected = 0.5 * k * (r - r0) ** 2 + composer = ClassicalMMComposer(bond_head=_ConstantBondHead(k=k, r0=r0)) + batch = _two_atom_bond_batch(r=r) + features = {"bonds": torch.zeros(1, 1, dtype=torch.float64)} + out = composer(batch, features) + assert math.isclose(float(out["energy"]), expected, rel_tol=1e-5, abs_tol=1e-5) + assert isinstance(out["ir"], PotentialIR) + + def test_optional_injected_evaluator(self): + """Injected evaluator receives IR + batch and returns energy sum.""" + bag = BondBag(k=torch.tensor([1.0]), r0=torch.tensor([1.0])) + + def evaluator(ir: PotentialIR, batch, *, pos=None) -> torch.Tensor: + assert ir.bonds is not None + return torch.tensor(42.0, dtype=torch.float64) + + composer = ClassicalMMComposer( + bond_head=_ConstantBondHead(k=1.0, r0=1.0), + evaluator=evaluator, + ) + batch = _two_atom_bond_batch() + ir = PotentialIR(bonds=bag, unit_system="class_i_canonical") + e = composer.energy(ir, batch, pos=batch["atoms", "pos"]) + assert float(e) == 42.0 + + +# --------------------------------------------------------------------------- +# Forces via ForceDerivation only (ac-008) +# --------------------------------------------------------------------------- + + +class TestClassicalMMComposerForces: + def test_energy_differentiable_wrt_pos_via_force_derivation(self): + composer = ClassicalMMComposer(bond_head=_ConstantBondHead(k=2.0, r0=1.0)) + batch = _two_atom_bond_batch(r=1.5) + features = {"bonds": torch.zeros(1, 1, dtype=torch.float64)} + ir = composer.parameterize(features, batch) + pos = batch["atoms", "pos"].clone() + + def energy_fn(p: torch.Tensor) -> torch.Tensor: + return composer.energy(ir, batch, pos=p) + + forces = ForceDerivation(method="autograd")(energy_fn, pos) + assert forces.shape == (2, 3) + # Along bond axis: atom0 pulls toward equilibrium, atom1 opposite (Newton 3) + assert forces[0, 0] > 0 # r > r0 → force on atom0 toward +x? F = -dE/dx + # E = 0.5*k*(r-r0)^2, r = x1-x0, dE/dx0 = k*(r-r0)*(-1) → F0 = +k*(r-r0) + assert math.isclose(float(forces[0, 0]), 2.0 * (1.5 - 1.0), abs_tol=1e-5) + assert math.isclose(float(forces[1, 0]), -2.0 * (1.5 - 1.0), abs_tol=1e-5) + + def test_source_has_no_hand_rolled_force_formula(self): + src = Path(__file__).resolve().parents[3] / "src/molpot/composition/classical_mm.py" + text = src.read_text() + # No analytic force kernels beyond ForceDerivation usage. + forbidden = [ + "force = -k *", + "forces = -", + "def calc_forces", + "def forces(", + ] + for needle in forbidden: + assert needle not in text, f"hand-rolled force pattern found: {needle!r}" + + +# --------------------------------------------------------------------------- +# Import boundary (ac-009) +# --------------------------------------------------------------------------- + + +class TestNoEncoderImports: + def test_classical_mm_and_mm_heads_have_no_molzoo_or_molrep_chem(self): + root = Path(__file__).resolve().parents[3] / "src/molpot/composition" + for name in ("classical_mm.py", "mm_heads.py"): + path = root / name + tree = ast.parse(path.read_text()) + for node in ast.walk(tree): + if isinstance(node, ast.Import): + for alias in node.names: + assert not alias.name.startswith("molzoo") + assert not alias.name.startswith("molrep.chem") + elif isinstance(node, ast.ImportFrom) and node.module: + assert not node.module.startswith("molzoo") + assert not node.module.startswith("molrep.chem") + + +# --------------------------------------------------------------------------- +# Docstrings (ac-010 partial) +# --------------------------------------------------------------------------- + + +class TestDocstrings: + def test_public_symbols_have_google_docstrings(self): + for cls in ( + BondParamHead, + AngleParamHead, + ProperTorsionParamHead, + ImproperParamHead, + ClassicalMMComposer, + ): + doc = cls.__doc__ or "" + assert len(doc.strip()) > 20, f"{cls.__name__} missing docstring" diff --git a/tests/test_molpot/test_composition/test_mm_heads.py b/tests/test_molpot/test_composition/test_mm_heads.py new file mode 100644 index 0000000..f5bf751 --- /dev/null +++ b/tests/test_molpot/test_composition/test_mm_heads.py @@ -0,0 +1,211 @@ +"""Tests for continuous Class-I MM parameter heads (learnable-classical-ff-03).""" + +from __future__ import annotations + +import math + +import torch + +from molpot.composition.heads import ChargeHead, LJParameterHead +from molpot.composition.mm_heads import ( + AngleParamHead, + BondParamHead, + ImproperParamHead, + ProperTorsionParamHead, +) +from molpot.composition.multihead import MultiHead + +# --------------------------------------------------------------------------- +# BondParamHead (ac-001) +# --------------------------------------------------------------------------- + + +class TestBondParamHead: + def test_shapes_and_positivity(self): + head = BondParamHead(feature_dim=8, hidden_dim=16) + features = torch.randn(5, 8) + out = head(features) + assert out["k"].shape == (5,) + assert out["r0"].shape == (5,) + assert torch.all(out["k"] > 0) + assert torch.all(out["r0"] > 0) + + def test_softplus_on_large_negative_inputs(self): + head = BondParamHead(feature_dim=4, hidden_dim=8, min_k=1e-3, min_r0=1e-3) + # Force MLP weights so pre-activations are large-negative if features are huge negative. + with torch.no_grad(): + for p in head.parameters(): + p.zero_() + if p.ndim == 1: + p.fill_(-50.0) + features = torch.full((3, 4), -100.0) + out = head(features) + assert torch.all(out["k"] >= head.min_k - 1e-12) + assert torch.all(out["r0"] >= head.min_r0 - 1e-12) + assert torch.all(out["k"] > 0) + assert torch.all(out["r0"] > 0) + + def test_endpoint_symmetry_same_features(self): + """Heads are pure MLPs: identical features → identical params (call-site symmetry). + + Endpoint symmetry for bonds (i,j)↔(j,i) is enforced by supplying + order-invariant pooled features at the call site; the head itself does + not reorder atoms. + """ + head = BondParamHead(feature_dim=6, hidden_dim=12) + feats = torch.randn(4, 6) + a = head(feats) + b = head(feats.clone()) + assert torch.allclose(a["k"], b["k"]) + assert torch.allclose(a["r0"], b["r0"]) + + +# --------------------------------------------------------------------------- +# AngleParamHead (ac-002) +# --------------------------------------------------------------------------- + + +class TestAngleParamHead: + def test_shapes_k_positive_theta0_in_open_interval(self): + head = AngleParamHead(feature_dim=8, hidden_dim=16) + features = torch.randn(7, 8) + out = head(features) + assert out["k"].shape == (7,) + assert out["theta0"].shape == (7,) + assert torch.all(out["k"] > 0) + # theta0 ∈ (0, π) + assert torch.all(out["theta0"] > 0) + assert torch.all(out["theta0"] < math.pi) + + def test_softplus_k_on_large_negative(self): + head = AngleParamHead(feature_dim=4, hidden_dim=8, min_k=1e-3) + with torch.no_grad(): + for p in head.parameters(): + p.zero_() + if p.ndim == 1: + p.fill_(-40.0) + out = head(torch.full((2, 4), -80.0)) + assert torch.all(out["k"] >= head.min_k - 1e-12) + assert torch.all(out["theta0"] > 0) + assert torch.all(out["theta0"] < math.pi) + + def test_endpoint_symmetry_same_features(self): + """(i,j,k)↔(k,j,i) symmetry is call-site feature pooling; head is pure MLP.""" + head = AngleParamHead(feature_dim=5, hidden_dim=10) + feats = torch.randn(3, 5) + a, b = head(feats), head(feats.clone()) + assert torch.allclose(a["k"], b["k"]) + assert torch.allclose(a["theta0"], b["theta0"]) + + +# --------------------------------------------------------------------------- +# ProperTorsionParamHead (ac-003) +# --------------------------------------------------------------------------- + + +class TestProperTorsionParamHead: + def test_multi_term_k_phase_shapes(self): + n_terms = 3 + head = ProperTorsionParamHead( + feature_dim=8, + hidden_dim=16, + n_terms=n_terms, + periodicity=(1, 2, 3), + ) + features = torch.randn(4, 8) + out = head(features) + assert out["k"].shape == (4, n_terms) + assert out["phase"].shape == (4, n_terms) + assert torch.all(out["k"] >= 0) + # Fixed periodicity buffer + assert out["periodicity"].shape == (n_terms,) + assert out["periodicity"].tolist() == [1, 2, 3] + assert out["idivf"].shape == (4,) + assert torch.all(out["idivf"] > 0) + + def test_k_nonneg_large_negative(self): + head = ProperTorsionParamHead(feature_dim=4, n_terms=2, periodicity=(1, 2)) + with torch.no_grad(): + for p in head.parameters(): + p.zero_() + if p.ndim == 1: + p.fill_(-30.0) + out = head(torch.full((2, 4), -50.0)) + assert torch.all(out["k"] >= 0) + + +# --------------------------------------------------------------------------- +# ImproperParamHead (ac-004) +# --------------------------------------------------------------------------- + + +class TestImproperParamHead: + def test_harmonic_mode(self): + head = ImproperParamHead( + feature_dim=8, + hidden_dim=16, + include_harmonic=True, + include_periodic=False, + ) + out = head(torch.randn(3, 8)) + assert "k" in out and "chi0" in out + assert out["k"].shape == (3,) + assert out["chi0"].shape == (3,) + assert torch.all(out["k"] > 0) + assert "phase" not in out + + def test_periodic_mode(self): + head = ImproperParamHead( + feature_dim=8, + n_terms=2, + periodicity=(2, 2), + include_harmonic=False, + include_periodic=True, + ) + out = head(torch.randn(2, 8)) + assert out["k"].shape == (2, 2) + assert out["phase"].shape == (2, 2) + assert torch.all(out["k"] >= 0) + assert out["periodicity"].tolist() == [2, 2] + assert "chi0" not in out + + def test_both_modes(self): + head = ImproperParamHead( + feature_dim=6, + n_terms=1, + periodicity=(2,), + include_harmonic=True, + include_periodic=True, + ) + out = head(torch.randn(2, 6)) + assert "k_harmonic" in out and "chi0" in out + assert "k_periodic" in out and "phase" in out + assert torch.all(out["k_harmonic"] > 0) + assert torch.all(out["k_periodic"] >= 0) + + +# --------------------------------------------------------------------------- +# MultiHead + ChargeHead + LJParameterHead reuse (ac-005) +# --------------------------------------------------------------------------- + + +class TestMultiHeadChargeLJReuse: + def test_merge_epsilon_sigma_charge_neutrality(self): + heads = MultiHead( + { + "lj": LJParameterHead(feature_dim=8, hidden_dim=16), + "q": ChargeHead(feature_dim=8, hidden_dim=16, total_charge=0.0), + } + ) + node_features = torch.randn(5, 8) + batch = torch.tensor([0, 0, 0, 1, 1], dtype=torch.long) + out = heads(node_features, batch=batch) + assert set(out) == {"epsilon", "sigma", "charge"} + assert out["epsilon"].shape == (5,) + assert out["sigma"].shape == (5,) + assert out["charge"].shape == (5,) + assert torch.all(out["epsilon"] > 0) + assert torch.all(out["sigma"] > 0) + # Per-molecule neutrality + assert out["charge"][:3].sum().abs() < 1e-5 + assert out["charge"][3:].sum().abs() < 1e-5 diff --git a/tests/test_molpot/test_composition/test_parameterizer.py b/tests/test_molpot/test_composition/test_parameterizer.py new file mode 100644 index 0000000..3e9d705 --- /dev/null +++ b/tests/test_molpot/test_composition/test_parameterizer.py @@ -0,0 +1,316 @@ +"""Tests for ClassicalMMParameterizer (learnable-classical-ff-05).""" + +from __future__ import annotations + +import ast +import math +from pathlib import Path +from typing import Mapping + +import torch +import torch.nn as nn +from tensordict import TensorDict + +from molpot.composition.classical_mm import ClassicalMMComposer +from molpot.composition.parameterizer import ( + KCAL_MOL_TO_EV, + ChemEmbeddingsLike, + ChemEncoderProtocol, + ClassicalMMParameterizer, + energy_kcal_to_ev, +) +from molpot.derivation import ForceDerivation +from molpot.ir import CLASS_I_CANONICAL, PotentialIR + +# --------------------------------------------------------------------------- +# Fixtures / fakes (no molzoo) +# --------------------------------------------------------------------------- + + +class _FakeEmbeddings: + """Structural ChemEmbeddingsLike: interaction_dict only.""" + + def __init__(self, features: Mapping[str, torch.Tensor]) -> None: + self._features = dict(features) + + def interaction_dict(self) -> dict[str, torch.Tensor]: + return dict(self._features) + + +class FakeEncoder(nn.Module): + """Minimal encoder satisfying ChemEncoderProtocol without molzoo/molrep.chem. + + Returns fixed feature tensors (constant heads ignore feature content). + """ + + def __init__(self, features: Mapping[str, torch.Tensor] | None = None) -> None: + super().__init__() + self._features = dict(features) if features is not None else {} + # Ensure at least one trainable param so encoder registers as a submodule. + self._dummy = nn.Parameter(torch.zeros(1)) + + def set_features(self, features: Mapping[str, torch.Tensor]) -> None: + self._features = dict(features) + + def forward(self, td: TensorDict) -> TensorDict: + return td + + def embeddings(self, td: TensorDict) -> _FakeEmbeddings: + return _FakeEmbeddings(self._features) + + +class _ConstantBondHead(nn.Module): + """Mock head returning fixed k, r0 (ignores features).""" + + def __init__(self, k: float, r0: float) -> None: + super().__init__() + self.k = k + self.r0 = r0 + + def forward(self, features: torch.Tensor) -> dict[str, torch.Tensor]: + n = features.shape[0] + return { + "k": torch.full((n,), self.k, dtype=features.dtype, device=features.device), + "r0": torch.full((n,), self.r0, dtype=features.dtype, device=features.device), + } + + +def _two_atom_bond_batch( + *, + r: float = 1.5, + dtype: torch.dtype = torch.float64, +) -> TensorDict: + pos = torch.tensor([[0.0, 0.0, 0.0], [r, 0.0, 0.0]], dtype=dtype) + return TensorDict( + { + "atoms": TensorDict( + { + "pos": pos, + "Z": torch.tensor([1, 1], dtype=torch.long), + "batch": torch.zeros(2, dtype=torch.long), + }, + batch_size=[2], + ), + "bonds": TensorDict( + { + "atomi": torch.tensor([0], dtype=torch.long), + "atomj": torch.tensor([1], dtype=torch.long), + }, + batch_size=[1], + ), + }, + batch_size=[], + ) + + +def _bond_features(dtype: torch.dtype = torch.float64) -> dict[str, torch.Tensor]: + return {"bonds": torch.zeros(1, 4, dtype=dtype)} + + +def _make_parameterizer( + *, + k: float = 2.0, + r0: float = 1.0, + features: Mapping[str, torch.Tensor] | None = None, +) -> ClassicalMMParameterizer: + feat = dict(features) if features is not None else _bond_features() + encoder = FakeEncoder(feat) + composer = ClassicalMMComposer(bond_head=_ConstantBondHead(k=k, r0=r0)) + return ClassicalMMParameterizer(encoder=encoder, composer=composer) + + +# --------------------------------------------------------------------------- +# Protocol structural typing +# --------------------------------------------------------------------------- + + +class TestChemEncoderProtocol: + def test_fake_encoder_satisfies_protocol(self): + enc = FakeEncoder(_bond_features()) + assert isinstance(enc, ChemEncoderProtocol) + + def test_fake_embeddings_satisfies_chem_embeddings_like(self): + emb = _FakeEmbeddings(_bond_features()) + assert isinstance(emb, ChemEmbeddingsLike) + + def test_encode_returns_interaction_feature_dict(self): + param = _make_parameterizer() + batch = _two_atom_bond_batch() + features = param.encode(batch) + assert "bonds" in features + assert features["bonds"].shape == (1, 4) + + +# --------------------------------------------------------------------------- +# parameterize → IR units kcal/mol + energy smoke +# --------------------------------------------------------------------------- + + +class TestParameterizeAndEnergy: + def test_parameterize_class_i_ir_kcal_mol(self): + param = _make_parameterizer(k=2.0, r0=1.0) + batch = _two_atom_bond_batch(r=1.5) + ir = param.parameterize(batch) + assert isinstance(ir, PotentialIR) + assert ir.unit_system == "class_i_canonical" + assert CLASS_I_CANONICAL["energy"] == "kcal/mol" + assert ir.bonds is not None + assert ir.bonds.k.shape == (1,) + assert math.isclose(float(ir.bonds.k[0]), 2.0, abs_tol=1e-8) + assert math.isclose(float(ir.bonds.r0[0]), 1.0, abs_tol=1e-8) + + def test_energy_matches_hand_built_bond_harmonic(self): + k, r0, r = 2.0, 1.0, 1.5 + expected = 0.5 * k * (r - r0) ** 2 # 0.25 kcal/mol + param = _make_parameterizer(k=k, r0=r0) + batch = _two_atom_bond_batch(r=r) + energy = param.energy(batch) + assert energy.shape == () + assert math.isclose(float(energy), expected, rel_tol=1e-5, abs_tol=1e-5) + + def test_energy_with_precomputed_ir(self): + param = _make_parameterizer(k=4.0, r0=1.2) + batch = _two_atom_bond_batch(r=1.4) + ir = param.parameterize(batch) + e1 = param.energy(batch, ir=ir) + e2 = param.energy(batch) + assert math.isclose(float(e1), float(e2), abs_tol=1e-10) + + def test_parameterize_accepts_explicit_features(self): + """Skip encoder when features= is provided.""" + encoder = FakeEncoder({"bonds": torch.ones(1, 4, dtype=torch.float64)}) + composer = ClassicalMMComposer(bond_head=_ConstantBondHead(k=2.0, r0=1.0)) + param = ClassicalMMParameterizer(encoder=encoder, composer=composer) + batch = _two_atom_bond_batch(r=1.5) + override = {"bonds": torch.zeros(1, 8, dtype=torch.float64)} + ir = param.parameterize(batch, features=override) + assert ir.bonds is not None + # Constant head ignores feature width; energy still golden. + e = param.energy(batch, ir=ir) + assert math.isclose(float(e), 0.25, abs_tol=1e-5) + + +# --------------------------------------------------------------------------- +# Forces via ForceDerivation only +# --------------------------------------------------------------------------- + + +class TestForces: + def test_forward_compute_forces_shape(self): + param = _make_parameterizer(k=2.0, r0=1.0) + batch = _two_atom_bond_batch(r=1.5) + out = param.forward(batch, compute_forces=True) + assert "energy" in out + assert "forces" in out + forces = out["forces"] + assert forces.shape == (2, 3) + # Analytic: F0_x = +k*(r-r0), F1_x = -k*(r-r0) + assert math.isclose(float(forces[0, 0]), 2.0 * (1.5 - 1.0), abs_tol=1e-5) + assert math.isclose(float(forces[1, 0]), -2.0 * (1.5 - 1.0), abs_tol=1e-5) + + def test_forward_writes_batch_keys(self): + param = _make_parameterizer(k=2.0, r0=1.0) + batch = _two_atom_bond_batch(r=1.5) + out = param.forward(batch, compute_forces=True) + # graphs.energy / atoms.forces convention + assert "graphs" in batch + assert "energy" in batch["graphs"] + assert "forces" in batch["atoms"] + assert torch.allclose(out["energy"], batch["graphs", "energy"]) + assert torch.allclose(out["forces"], batch["atoms", "forces"]) + + def test_forces_via_injected_force_derivation(self): + enc = FakeEncoder(_bond_features()) + composer = ClassicalMMComposer(bond_head=_ConstantBondHead(k=2.0, r0=1.0)) + fd = ForceDerivation(method="autograd") + param = ClassicalMMParameterizer(encoder=enc, composer=composer, force_derivation=fd) + batch = _two_atom_bond_batch(r=1.5) + out = param(batch, compute_forces=True) + assert out["forces"].shape == (2, 3) + + def test_source_has_no_hand_rolled_force_formula(self): + src = Path(__file__).resolve().parents[3] / "src/molpot/composition/parameterizer.py" + text = src.read_text() + forbidden = [ + "force = -k *", + "forces = -grad", + "def calc_forces", + ] + for needle in forbidden: + assert needle not in text, f"hand-rolled force pattern found: {needle!r}" + assert "ForceDerivation" in text + + +# --------------------------------------------------------------------------- +# Units boundary +# --------------------------------------------------------------------------- + + +class TestUnitsBoundary: + def test_ir_stays_kcal_not_mutated_to_ev(self): + param = _make_parameterizer() + batch = _two_atom_bond_batch() + ir = param.parameterize(batch) + assert ir.unit_system == "class_i_canonical" + # Conversion helper does not mutate IR + e_kcal = param.energy(batch, ir=ir) + e_ev = energy_kcal_to_ev(e_kcal) + assert ir.unit_system == "class_i_canonical" + assert not math.isclose(float(e_kcal), float(e_ev), rel_tol=0.1) + assert math.isclose( + float(e_ev), + float(e_kcal) * KCAL_MOL_TO_EV, + abs_tol=1e-12, + ) + + def test_kcal_to_ev_factor_matches_documented_constant(self): + # 1 eV ≈ 23.060547830619026 kcal/mol → KCAL_MOL_TO_EV ≈ 0.0433641 + assert math.isclose(KCAL_MOL_TO_EV, 1.0 / 23.060547830619026, rel_tol=1e-12) + x = torch.tensor(23.060547830619026) + assert math.isclose(float(energy_kcal_to_ev(x)), 1.0, abs_tol=1e-10) + + +# --------------------------------------------------------------------------- +# Import boundary + submodule registration +# --------------------------------------------------------------------------- + + +class TestImportBoundaryAndRegistration: + def test_parameterizer_module_has_no_molzoo_import(self): + path = Path(__file__).resolve().parents[3] / "src/molpot/composition/parameterizer.py" + tree = ast.parse(path.read_text()) + for node in ast.walk(tree): + if isinstance(node, ast.Import): + for alias in node.names: + assert not alias.name.startswith("molzoo") + assert not alias.name.startswith("molrep.chem") + elif isinstance(node, ast.ImportFrom) and node.module: + assert not node.module.startswith("molzoo") + assert not node.module.startswith("molrep.chem") + + def test_encoder_registered_as_submodule(self): + param = _make_parameterizer() + names = {name for name, _ in param.named_modules()} + assert "encoder" in names + assert "composer" in names + # Dummy param from FakeEncoder is in state_dict + assert any(k.startswith("encoder.") for k in param.state_dict()) + + +# --------------------------------------------------------------------------- +# Docstrings +# --------------------------------------------------------------------------- + + +class TestDocstrings: + def test_public_symbols_have_google_docstrings(self): + for obj in ( + ClassicalMMParameterizer, + ClassicalMMParameterizer.encode, + ClassicalMMParameterizer.parameterize, + ClassicalMMParameterizer.energy, + ClassicalMMParameterizer.forward, + energy_kcal_to_ev, + ): + doc = obj.__doc__ or "" + assert len(doc.strip()) > 20, f"{obj} missing docstring" diff --git a/tests/test_molpot/test_composition/test_potential_parity.py b/tests/test_molpot/test_composition/test_potential_parity.py new file mode 100644 index 0000000..929308e --- /dev/null +++ b/tests/test_molpot/test_composition/test_potential_parity.py @@ -0,0 +1,211 @@ +"""Validation B0 — Class-I PotentialIR + kernel parity goldens (kcal/mol, Å). + +Hard-coded analytical references only. No live OpenMM / molpy ForceField. +Units: CLASS_I_CANONICAL (kcal/mol, Å, kcal/mol/Å). Improper center-first +(molrs row 0 = center). Forces via ForceDerivation(method="autograd"). +""" + +from __future__ import annotations + +import math + +import torch +from tensordict import TensorDict + +from molpot.composition.classical_mm import ClassicalMMComposer +from molpot.derivation import ForceDerivation +from molpot.ir import NonbondedScaling, PotentialIR +from molpot.ir.bags import ( + BondBag, +) +from molpot.potentials.angles import AngleHarmonic +from molpot.potentials.bonds import BondHarmonic +from molpot.potentials.dihedrals.periodic import ProperTorsionPeriodic +from molpot.potentials.elec.potentials.coulomb import CoulombPotential +from molpot.potentials.elec.prefactors import kcalmol_A +from molpot.potentials.impropers.harmonic import ImproperHarmonic +from molpot.potentials.vdw.lj126 import lj126_pair_energy + +_TOL_E = 1e-10 +_TOL_F = 1e-6 + + +class TestPotentialParityTerms: + def test_bond_harmonic_golden(self): + pot = BondHarmonic( + k=torch.tensor([2.0], dtype=torch.float64), r0=torch.tensor([1.0], dtype=torch.float64) + ) + pos = torch.tensor([[0.0, 0.0, 0.0], [1.5, 0.0, 0.0]], dtype=torch.float64) + e = pot( + pos=pos, + bond_index=torch.tensor([[0], [1]], dtype=torch.long), + bond_types=torch.tensor([0], dtype=torch.long), + ) + assert abs(float(e) - 0.25) <= _TOL_E + + def test_angle_harmonic_golden(self): + pot = AngleHarmonic( + k=torch.tensor([2.0], dtype=torch.float64), + theta0=torch.tensor([math.pi / 3], dtype=torch.float64), + ) + # right angle at atom 1: (1,0,0)-(0,0,0)-(0,1,0) + pos = torch.tensor( + [[1.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 1.0, 0.0]], + dtype=torch.float64, + ) + e = pot( + pos=pos, + angle_index=torch.tensor([[0], [1], [2]], dtype=torch.long), + angle_types=torch.tensor([0], dtype=torch.long), + ) + expected = (math.pi / 6) ** 2 + assert abs(float(e) - expected) <= _TOL_E + + def test_proper_cis_golden(self): + # E = k/s * [1 + cos(n*φ - γ)]; cis φ≈0, n=1, γ=0, s=1 → 2k with k=1 + pot = ProperTorsionPeriodic( + k=torch.tensor([[1.0]], dtype=torch.float64), + periodicity=torch.tensor([1.0], dtype=torch.float64), + phase=torch.tensor([[0.0]], dtype=torch.float64), + idivf=torch.tensor([1.0], dtype=torch.float64), + ) + # planar cis: i-j-k-l with φ=0 + pos = torch.tensor( + [ + [1.0, 0.0, 0.0], + [0.0, 0.0, 0.0], + [0.0, 1.0, 0.0], + [1.0, 1.0, 0.0], + ], + dtype=torch.float64, + ) + e = pot( + pos=pos, + proper_index=torch.tensor([[0], [1], [2], [3]], dtype=torch.long), + proper_types=torch.tensor([0], dtype=torch.long), + ) + assert abs(float(e) - 2.0) <= _TOL_E + + def test_improper_harmonic_golden(self): + """E = ½ k (χ − χ₀)² with k=2, χ₀=0, χ=π/6 → (π/6)².""" + pot = ImproperHarmonic( + k=torch.tensor([2.0], dtype=torch.float64), + chi0=torch.tensor([0.0], dtype=torch.float64), + ) + chi = math.pi / 6 + # molrs center-first improper_index [c,i,j,k]; dihedral uses (i,c,j,k). + pos = torch.tensor( + [ + [0.0, 0.0, 0.0], # center + [1.0, 0.0, 0.0], # i + [0.0, 1.0, 0.0], # j + [math.cos(chi), 1.0, math.sin(chi)], # k → |χ|=π/6 + ], + dtype=torch.float64, + ) + e = pot( + pos=pos, + improper_index=torch.tensor([[0], [1], [2], [3]], dtype=torch.long), + improper_types=torch.tensor([0], dtype=torch.long), + ) + expected = chi**2 # 0.5 * k=2 → E = χ² + assert abs(float(e) - expected) <= 1e-8 + + def test_coulomb_golden(self): + pot = CoulombPotential(prefactor=kcalmol_A) + r = torch.tensor([2.0], dtype=torch.float64) + # pair energy q_i q_j / r * prefactor with q=±1 + e_pair = pot.from_dist(r) * (-1.0) # charges product -1 + expected = -kcalmol_A / 2.0 + assert abs(float(e_pair) - expected) <= _TOL_E + + def test_lj_golden(self): + r = torch.tensor([2.0], dtype=torch.float64) + e = lj126_pair_energy( + r, torch.tensor([1.0], dtype=torch.float64), torch.tensor([1.0], dtype=torch.float64) + ) + assert abs(float(e) - (-0.0615234375)) <= _TOL_E + + +class TestPotentialParityTotal: + def test_composer_bonded_bond_only(self): + ir = PotentialIR( + bonds=BondBag( + k=torch.tensor([2.0], dtype=torch.float64), + r0=torch.tensor([1.0], dtype=torch.float64), + ) + ) + batch = TensorDict( + { + "atoms": TensorDict( + {"pos": torch.tensor([[0.0, 0.0, 0.0], [1.5, 0.0, 0.0]], dtype=torch.float64)}, + batch_size=[2], + ), + "bonds": TensorDict( + { + "bond_index": torch.tensor([[0], [1]], dtype=torch.long), + }, + batch_size=[], + ), + }, + batch_size=[], + ) + # ClassicalMMComposer may read bond_index differently - check helpers + composer = ClassicalMMComposer() + # Ensure bond_index layout matches composer expectations + e = composer.energy(ir, batch) + assert abs(float(e) - 0.25) <= _TOL_E + terms = composer.term_energies(ir, batch) + assert abs(float(terms["bonds"]) - 0.25) <= _TOL_E + + +class TestPotentialParityForces: + def test_bond_forces_golden(self): + pot = BondHarmonic( + k=torch.tensor([2.0], dtype=torch.float64), + r0=torch.tensor([1.0], dtype=torch.float64), + ) + pos = torch.tensor( + [[0.0, 0.0, 0.0], [1.5, 0.0, 0.0]], dtype=torch.float64, requires_grad=True + ) + bond_index = torch.tensor([[0], [1]], dtype=torch.long) + bond_types = torch.tensor([0], dtype=torch.long) + + def energy_fn(p: torch.Tensor) -> torch.Tensor: + return pot(pos=p, bond_index=bond_index, bond_types=bond_types) + + forces = ForceDerivation(method="autograd")(energy_fn, pos) + assert abs(float(forces[0, 0].detach()) - 1.0) <= _TOL_F + assert abs(float(forces[1, 0].detach()) - (-1.0)) <= _TOL_F + + def test_force_is_neg_grad(self): + pot = BondHarmonic( + k=torch.tensor([2.0], dtype=torch.float64), + r0=torch.tensor([1.0], dtype=torch.float64), + ) + pos = torch.tensor( + [[0.0, 0.0, 0.0], [1.5, 0.0, 0.0]], dtype=torch.float64, requires_grad=True + ) + bond_index = torch.tensor([[0], [1]], dtype=torch.long) + bond_types = torch.tensor([0], dtype=torch.long) + e = pot(pos=pos, bond_index=bond_index, bond_types=bond_types) + (g,) = torch.autograd.grad(e, pos, create_graph=False) + forces = ForceDerivation(method="autograd")( + lambda p: pot(pos=p, bond_index=bond_index, bond_types=bond_types), + pos.detach().requires_grad_(True), + ) + assert torch.allclose(forces, -g.detach(), atol=_TOL_F) + + +class TestPotentialParityUnits: + def test_nonbonded_scaling_defaults(self): + s = NonbondedScaling() + assert s.scale_q_12 == 0.0 + assert s.scale_lj_12 == 0.0 + assert abs(s.scale_q_14 - 5.0 / 6.0) < 1e-12 + assert abs(s.scale_lj_14 - 0.5) < 1e-12 + ir = PotentialIR() + assert ir.unit_system == "class_i_canonical" + + def test_kcalmol_prefactor(self): + assert abs(kcalmol_A - 332.0637132991921) < 1e-9 diff --git a/tests/test_molpot/test_composition/test_sonata.py b/tests/test_molpot/test_composition/test_sonata.py index 6d36921..efcb6ed 100644 --- a/tests/test_molpot/test_composition/test_sonata.py +++ b/tests/test_molpot/test_composition/test_sonata.py @@ -1,8 +1,8 @@ -"""RED tests for `molpot.composition.sonata` — Sonata composer. +"""Unit tests for `molpot.composition.sonata` — Sonata composer. Sub-spec 01 of the Sonata model line. Tests cover: -* ac-001 — `build_sonata` returns a wired `nn.Module` with the right +* ac-001 — `Sonata.from_encoder` returns a wired `nn.Module` with the right sub-module types. * ac-002 — `Sonata.__init__` refuses `kappa_head=`, `alpha_head=`, `induced_*` kwargs (future `LesPolarizable` composer territory). @@ -15,10 +15,10 @@ * ac-007 — `compute_stress=True` adds a symmetric `(B, 3, 3)` stress. * ac-008 — `Sonata.from_spec(sonata.config, encoder)` round-trips when state_dict is transferred. -* ac-009 — `build_sonata` validates `encoder.expose_tensor_track` and - `encoder.l_max`. -* ac-010 — `Sonata`, `SonataSpec`, `build_sonata` are exported from - `molpot` and `molpot.composition`. +* ac-009 — `Sonata.from_encoder` validates `encoder.expose_tensor_track` + and `encoder.l_max`. +* ac-010 — `Sonata` and `SonataSpec` are exported from `molpot` and + `molpot.composition`. """ from __future__ import annotations @@ -29,7 +29,7 @@ from tensordict import TensorDict from molpot import Polarization -from molpot.composition import Sonata, SonataSpec, build_sonata +from molpot.composition import Sonata, SonataSpec from molpot.heads import EdgeEnergyHead, PermMultipoleHead from molpot.potentials import EwaldMultipoleEnergy from molzoo import Allegro @@ -149,8 +149,8 @@ def _short_range_head(encoder: Allegro) -> EdgeEnergyHead: # --------------------------------------------------------------------------- -def test_build_sonata_returns_wired_model(encoder): - model = build_sonata( +def test_from_encoder_returns_wired_model(encoder): + model = Sonata.from_encoder( encoder, charge=True, dipole=True, @@ -243,7 +243,7 @@ def test_refuse_polarization_in_short_range_list(encoder): def test_forward_output_schema(encoder, batch): - sonata = build_sonata( + sonata = Sonata.from_encoder( encoder, charge=True, dipole=True, @@ -283,7 +283,7 @@ def test_forward_output_schema(encoder, batch): def test_energy_decomposition(encoder, batch): - sonata = build_sonata( + sonata = Sonata.from_encoder( encoder, charge=True, dipole=True, @@ -303,7 +303,7 @@ def test_energy_decomposition(encoder, batch): def test_compute_forces(encoder, batch): - sonata = build_sonata( + sonata = Sonata.from_encoder( encoder, charge=True, dipole=True, @@ -324,7 +324,7 @@ def test_compute_forces(encoder, batch): def test_compute_stress(encoder, batch_with_cell): - sonata = build_sonata( + sonata = Sonata.from_encoder( encoder, charge=True, dipole=True, @@ -343,7 +343,7 @@ def test_compute_stress(encoder, batch_with_cell): def test_spec_round_trip(encoder, batch): - sonata1 = build_sonata( + sonata1 = Sonata.from_encoder( encoder, charge=True, dipole=True, @@ -369,13 +369,13 @@ def test_spec_round_trip(encoder, batch): # --------------------------------------------------------------------------- -# ac-009 — build_sonata validates encoder +# ac-009 — from_encoder validates encoder # --------------------------------------------------------------------------- -def test_build_sonata_requires_expose_tensor_track(encoder_no_tensor_track): +def test_from_encoder_requires_expose_tensor_track(encoder_no_tensor_track): with pytest.raises(ValueError, match="expose_tensor_track"): - build_sonata( + Sonata.from_encoder( encoder_no_tensor_track, charge=True, dipole=True, @@ -384,9 +384,9 @@ def test_build_sonata_requires_expose_tensor_track(encoder_no_tensor_track): ) -def test_build_sonata_requires_lmax_for_dipole(encoder_lmax1): +def test_from_encoder_requires_lmax_for_dipole(encoder_lmax1): with pytest.raises(ValueError, match="l_max"): - build_sonata( + Sonata.from_encoder( encoder_lmax1, charge=True, dipole=True, @@ -406,10 +406,11 @@ def test_public_surface_reexports(): assert "Sonata" in molpot.__all__ assert "SonataSpec" in molpot.__all__ - assert "build_sonata" in molpot.__all__ + assert "build_sonata" not in molpot.__all__ assert "Sonata" in molpot.composition.__all__ assert "SonataSpec" in molpot.composition.__all__ - assert "build_sonata" in molpot.composition.__all__ + assert "build_sonata" not in molpot.composition.__all__ assert molpot.Sonata is Sonata assert molpot.SonataSpec is SonataSpec - assert molpot.build_sonata is build_sonata + assert not hasattr(molpot, "build_sonata") + assert not hasattr(molpot.composition, "build_sonata") diff --git a/tests/test_molpot/test_composition/test_sonata_batch.py b/tests/test_molpot/test_composition/test_sonata_batch.py index 5c387c8..f98e338 100644 --- a/tests/test_molpot/test_composition/test_sonata_batch.py +++ b/tests/test_molpot/test_composition/test_sonata_batch.py @@ -21,7 +21,7 @@ from tensordict import TensorDict from molix.config import config -from molpot.composition import Sonata, build_sonata +from molpot.composition import Sonata from molzoo import Allegro # --------------------------------------------------------------------------- @@ -49,7 +49,7 @@ def sonata_fp64() -> Sonata: avg_num_neighbors=12.0, expose_tensor_track=True, ) - sonata = build_sonata( + sonata = Sonata.from_encoder( encoder, sigma=1.0, dl=2.0, diff --git a/tests/test_molpot/test_composition/test_sonata_boundary.py b/tests/test_molpot/test_composition/test_sonata_boundary.py index 02b8d07..62681f5 100644 --- a/tests/test_molpot/test_composition/test_sonata_boundary.py +++ b/tests/test_molpot/test_composition/test_sonata_boundary.py @@ -36,7 +36,7 @@ from molix.config import config from molpot import Polarization -from molpot.composition import Sonata, build_sonata +from molpot.composition import Sonata from molzoo import Allegro # --------------------------------------------------------------------------- @@ -64,7 +64,7 @@ def sonata_and_batch() -> tuple[Sonata, TensorDict]: avg_num_neighbors=12.0, expose_tensor_track=True, ) - sonata = build_sonata( + sonata = Sonata.from_encoder( encoder, sigma=1.0, dl=2.0, diff --git a/tests/test_molpot/test_composition/test_sonata_invariants.py b/tests/test_molpot/test_composition/test_sonata_invariants.py index f4206fe..2a13df7 100644 --- a/tests/test_molpot/test_composition/test_sonata_invariants.py +++ b/tests/test_molpot/test_composition/test_sonata_invariants.py @@ -27,14 +27,14 @@ import cuequivariance_torch as cuet import torch from tensordict import TensorDict + +from molpot.composition import Sonata from tests.conftest import ( permute_graph, rotate_graph, translate_graph, ) -from molpot.composition import Sonata - # --------------------------------------------------------------------------- # Local helpers # --------------------------------------------------------------------------- diff --git a/tests/test_molpot/test_composition/test_sonata_periodic.py b/tests/test_molpot/test_composition/test_sonata_periodic.py index 24611c6..d5b2ed6 100644 --- a/tests/test_molpot/test_composition/test_sonata_periodic.py +++ b/tests/test_molpot/test_composition/test_sonata_periodic.py @@ -50,7 +50,7 @@ from tensordict import TensorDict from molix.config import config -from molpot.composition import Sonata, build_sonata +from molpot.composition import Sonata from molzoo import Allegro # ``cuequivariance_ops_torch`` is the optimized CUDA-only backend for @@ -158,7 +158,7 @@ def single_graph_periodic() -> tuple[Sonata, TensorDict]: avg_num_neighbors=12.0, expose_tensor_track=True, ) - sonata = build_sonata( + sonata = Sonata.from_encoder( encoder, sigma=1.0, dl=2.0, diff --git a/tests/test_molpot/test_derivation/__init__.py b/tests/test_molpot/test_derivation/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/test_molpot/test_derivation/conftest.py b/tests/test_molpot/test_derivation/conftest.py new file mode 100644 index 0000000..ff79c06 --- /dev/null +++ b/tests/test_molpot/test_derivation/conftest.py @@ -0,0 +1,125 @@ +"""Shared toy potential / batch fixtures for the derivation unit tests. + +The toy is the analytic quadratic potential used throughout this directory:: + + E_g = s · Σ_{i ∈ g} ‖r_i‖² (eV, positions in Å) + F_i = -∂E/∂r_i = -2 s · r_i (eV/Å) + +Both derivation backends must reproduce ``-2 s r`` exactly, so it pins the +kernels without any reference implementation. ``n_forward`` counts energy-core +invocations — that is how the "one forward per force pass" contract is checked. + +``test_readouts.py`` carries its own float32 copy of the same toy (predating +this conftest); it is intentionally left untouched, so the names here are +distinct (``KernelToyPotential`` / ``kernel_toy_batch``) and default to float64 +for the 1e-12 kernel assertions. +""" + +from __future__ import annotations + +from collections.abc import Callable + +import torch +import torch.nn as nn +from tensordict import TensorDict + +from molpot.derivation.protocol import write_energy + + +class KernelToyPotential(nn.Module): + """``E = s · Σ‖pos‖²`` written onto the batch; counts energy-core calls. + + Args: + scale: Initial value of the single parameter ``s`` (eV/Ų). + write_atomic: Also write per-atom energies at ``("atoms", "energy")``. + ``False`` exercises the "no atomic energy" branch of the kernels. + write_energy_key: Write ``("graphs", "energy")`` at all. ``False`` + makes the core violate the contract, which the kernels must + reject with a ``RuntimeError`` naming ``graphs.energy``. + """ + + def __init__( + self, + *, + scale: float = 1.0, + write_atomic: bool = True, + write_energy_key: bool = True, + ) -> None: + super().__init__() + self.scale = nn.Parameter(torch.tensor(scale, dtype=torch.float64)) + self.write_atomic = write_atomic + self.write_energy_key = write_energy_key + self.n_forward = 0 + + def forward(self, batch: TensorDict) -> TensorDict: + """Write ``graphs.energy`` (eV) for ``batch`` and return it in place.""" + self.n_forward += 1 + if not self.write_energy_key: + return batch + pos = batch["atoms", "pos"] + atom_batch = batch["atoms", "batch"] + atom_energy = self.scale.to(pos.dtype) * (pos * pos).sum(dim=-1) + n_graphs = int(batch["graphs"].batch_size[0]) + energy = torch.zeros(n_graphs, dtype=pos.dtype, device=pos.device) + energy = energy.index_add(0, atom_batch, atom_energy) + write_energy( + batch, + energy, + atomic_energy=atom_energy if self.write_atomic else None, + ) + return batch + + def total_energy(self, pos: torch.Tensor) -> torch.Tensor: + """Scalar total energy (eV) for raw positions ``(N, 3)`` — finite-difference oracle.""" + return self.scale.to(pos.dtype) * (pos * pos).sum() + + +def kernel_toy_batch( + n_atoms: int = 4, + n_graphs: int = 2, + *, + dtype: torch.dtype = torch.float64, +) -> TensorDict: + """Post-collate batch with ``pos = arange(3 N).reshape(N, 3) * 0.1`` Å.""" + per_graph = n_atoms // n_graphs + pos = torch.arange(n_atoms * 3, dtype=dtype).reshape(n_atoms, 3) * 0.1 + return TensorDict( + atoms=TensorDict( + pos=pos, + batch=torch.arange(n_graphs).repeat_interleave(per_graph), + batch_size=[n_atoms], + ), + graphs=TensorDict(batch_size=[n_graphs]), + batch_size=[], + ) + + +def central_difference_forces( + energy_fn: Callable[[torch.Tensor], torch.Tensor], + pos: torch.Tensor, + *, + h: float = 1e-5, +) -> torch.Tensor: + """``F = -dE/dr`` by central differences — the domain oracle for the kernels. + + Args: + energy_fn: Total energy (eV) of a position configuration ``(N, 3)`` Å. + Must be free of autograd side effects; called ``6 N`` times. + pos: Positions ``(N, 3)`` in Å (values only; gradients are not used). + h: Displacement in Å. ``1e-5`` keeps float64 round-off near 1e-10 eV/Å + while the ``O(h²)`` truncation term stays far below 1e-6 eV/Å. + + Returns: + Numerical forces ``(N, 3)`` in eV/Å. + """ + base = pos.detach().clone() + forces = torch.zeros_like(base) + with torch.no_grad(): + for atom in range(base.shape[0]): + for axis in range(base.shape[1]): + plus = base.clone() + plus[atom, axis] += h + minus = base.clone() + minus[atom, axis] -= h + forces[atom, axis] = -(energy_fn(plus) - energy_fn(minus)) / (2.0 * h) + return forces diff --git a/tests/test_molpot/test_derivation/test_force.py b/tests/test_molpot/test_derivation/test_force.py index d9872a0..4d93e96 100644 --- a/tests/test_molpot/test_derivation/test_force.py +++ b/tests/test_molpot/test_derivation/test_force.py @@ -1,13 +1,20 @@ -"""Unit tests for molpot.derivation.force.ForceDerivation backends.""" +"""Unit tests for molpot.derivation.force.ForceDerivation backends. + +Each backend is pure: functorch helpers never run under method=autograd and +autograd helpers never run under method=functorch. +""" from __future__ import annotations import torch -from molpot.derivation.force import ForceDerivation +from molpot.derivation.force import ( + ForceDerivation, + autograd_forces_from_energy, +) -def test_functorch_and_autograd_agree(): +def test_functorch_and_autograd_agree_on_energy_fn(): pos = torch.randn(5, 3) def energy_fn(p: torch.Tensor) -> torch.Tensor: @@ -18,9 +25,44 @@ def energy_fn(p: torch.Tensor) -> torch.Tensor: torch.testing.assert_close(f_ft, f_ag, atol=1e-5, rtol=1e-5) +def test_autograd_forces_from_energy_matches_autograd_forward(): + pos = torch.randn(4, 3, requires_grad=True) + energy = (pos.pow(2).sum(dim=-1) * torch.tensor([1.0, 2.0, 0.5, 1.5])).sum() + f_from_e = autograd_forces_from_energy(energy, pos, create_graph=False) + + def energy_fn(p: torch.Tensor) -> torch.Tensor: + return (p.pow(2).sum(dim=-1) * torch.tensor([1.0, 2.0, 0.5, 1.5])).sum() + + f_fwd = ForceDerivation(method="autograd")(energy_fn, pos.detach()) + torch.testing.assert_close(f_from_e, f_fwd, atol=1e-5, rtol=1e-5) + + +def test_forces_from_energy_rejected_on_functorch_instance(): + deriv = ForceDerivation(method="functorch") + pos = torch.zeros(2, 3, requires_grad=True) + try: + deriv.forces_from_energy(pos.pow(2).sum(), pos) + raise AssertionError("expected ValueError") + except ValueError as e: + assert "autograd-only" in str(e) + + +def test_has_aux_rejected_on_autograd_instance(): + deriv = ForceDerivation(method="autograd") + + def energy_fn_aux(p: torch.Tensor): + return p.pow(2).sum(), p + + try: + deriv(energy_fn_aux, torch.zeros(2, 3), has_aux=True) + raise AssertionError("expected ValueError") + except ValueError as e: + assert "functorch-only" in str(e) + + def test_invalid_method_raises(): try: - ForceDerivation(method="magic") # type: ignore[arg-type] + ForceDerivation(method="magic") raise AssertionError("expected ValueError") except ValueError: pass diff --git a/tests/test_molpot/test_derivation/test_kernels.py b/tests/test_molpot/test_derivation/test_kernels.py new file mode 100644 index 0000000..59f7477 --- /dev/null +++ b/tests/test_molpot/test_derivation/test_kernels.py @@ -0,0 +1,199 @@ +"""Unit tests for the shared batch-level force passes in molpot.derivation.kernels. + +Mirrors ``src/molpot/derivation/kernels.py``. Each test targets exactly one of +the two kernels on the analytic toy potential ``E = s · Σ‖r‖²`` (conftest), for +which ``F = -2 s r`` is known in closed form — no reference implementation and +no third-party oracle is involved. + +Units: positions Å, energy eV, forces eV/Å. All assertions run in float64 at +``atol=1e-12, rtol=0`` (the "exact" energy/force band for a closed-form +gradient), except the finite-difference domain check, which uses the numerical +band ``1e-6 eV/Å``. +""" + +from __future__ import annotations + +import pytest +import torch +from tensordict import TensorDict + +from molpot.derivation.kernels import func_force_pass, grad_force_pass +from molpot.derivation.protocol import ENERGY_KEY, FORCES_KEY, POS_KEY +from tests.test_molpot.test_derivation.conftest import ( + KernelToyPotential, + central_difference_forces, + kernel_toy_batch, +) + +#: ``F = -2 s r`` for ``s = 1.0`` and ``pos = arange(12).reshape(4, 3) * 0.1`` Å, +#: i.e. ``-0.2 * arange(12)`` in eV/Å. Hard-coded, not recomputed from ``pos``. +EXPECTED_FORCES: list[list[float]] = [ + [-0.0, -0.2, -0.4], + [-0.6, -0.8, -1.0], + [-1.2, -1.4, -1.6], + [-1.8, -2.0, -2.2], +] + +EXACT_FORCE_ATOL: float = 1e-12 +FD_FORCE_ATOL: float = 1e-6 + + +def _expected(dtype: torch.dtype = torch.float64) -> torch.Tensor: + return torch.tensor(EXPECTED_FORCES, dtype=dtype) + + +def _with_position_leaf(batch: TensorDict) -> torch.Tensor: + """Replace ``atoms.pos`` with a live ``requires_grad`` leaf; return it.""" + leaf = batch[POS_KEY].detach().clone().requires_grad_(True) + batch[POS_KEY] = leaf + return leaf + + +class TestGradForcePass: + """``grad_force_pass`` — one energy forward + ``torch.autograd.grad``.""" + + def test_forces_match_hardcoded_analytic_gradient(self) -> None: + toy = KernelToyPotential() + batch = grad_force_pass(toy, kernel_toy_batch()) + torch.testing.assert_close(batch[FORCES_KEY], _expected(), atol=EXACT_FORCE_ATOL, rtol=0.0) + + def test_energy_core_none_reuses_materialised_energy(self) -> None: + toy = KernelToyPotential() + batch = kernel_toy_batch() + _with_position_leaf(batch) + with torch.enable_grad(): + batch = toy(batch) + toy.n_forward = 0 # only count what the kernel does + + batch = grad_force_pass(None, batch, detach_energy=False) + + assert toy.n_forward == 0 + torch.testing.assert_close(batch[FORCES_KEY], _expected(), atol=EXACT_FORCE_ATOL, rtol=0.0) + + def test_energy_core_none_without_position_leaf_raises(self) -> None: + batch = kernel_toy_batch() # atoms.pos does not require grad + with pytest.raises(RuntimeError): + grad_force_pass(None, batch) + + def test_energy_core_none_with_non_leaf_position_raises(self) -> None: + toy = KernelToyPotential() + batch = kernel_toy_batch() + leaf = batch[POS_KEY].detach().clone().requires_grad_(True) + batch[POS_KEY] = leaf * 1.0 # requires_grad, but not a leaf + with torch.enable_grad(): + batch = toy(batch) + with pytest.raises(RuntimeError): + grad_force_pass(None, batch) + + def test_energy_core_without_graphs_energy_raises(self) -> None: + toy = KernelToyPotential(write_energy_key=False) + with pytest.raises(RuntimeError, match=r"graphs\.energy"): + grad_force_pass(toy, kernel_toy_batch()) + + def test_detach_energy_false_keeps_energy_attached(self) -> None: + toy = KernelToyPotential() + batch = grad_force_pass(toy, kernel_toy_batch(), detach_energy=False) + assert batch[ENERGY_KEY].requires_grad is True + + def test_detach_energy_true_detaches_energy(self) -> None: + toy = KernelToyPotential() + batch = grad_force_pass(toy, kernel_toy_batch(), detach_energy=True) + assert batch[ENERGY_KEY].requires_grad is False + + def test_detach_energy_none_detaches_when_kernel_creates_leaf(self) -> None: + toy = KernelToyPotential() + batch = grad_force_pass(toy, kernel_toy_batch(), detach_energy=None) + assert batch[ENERGY_KEY].requires_grad is False + + def test_detach_energy_none_keeps_energy_when_caller_supplies_leaf(self) -> None: + toy = KernelToyPotential() + batch = kernel_toy_batch() + leaf = _with_position_leaf(batch) + + batch = grad_force_pass(toy, batch, detach_energy=None) + + assert batch[POS_KEY] is leaf # kernel must not create a second leaf + assert batch[ENERGY_KEY].requires_grad is True + + def test_create_graph_true_force_loss_reaches_parameters(self) -> None: + toy = KernelToyPotential() + batch = grad_force_pass(toy, kernel_toy_batch(), create_graph=True, detach_energy=False) + batch[FORCES_KEY].pow(2).mean().backward() + assert toy.scale.grad is not None + assert toy.scale.grad.abs().item() > 0.0 + + def test_create_graph_false_disconnects_forces(self) -> None: + toy = KernelToyPotential() + batch = grad_force_pass(toy, kernel_toy_batch(), create_graph=False, detach_energy=False) + assert batch[FORCES_KEY].requires_grad is False + + def test_no_grad_context_still_returns_forces(self) -> None: + toy = KernelToyPotential() + with torch.no_grad(): + batch = grad_force_pass(toy, kernel_toy_batch()) + torch.testing.assert_close(batch[FORCES_KEY], _expected(), atol=EXACT_FORCE_ATOL, rtol=0.0) + + def test_forces_match_central_finite_differences(self) -> None: + toy = KernelToyPotential() + batch = kernel_toy_batch() + pos = batch[POS_KEY].detach().clone() + + batch = grad_force_pass(toy, batch) + numerical = central_difference_forces(toy.total_energy, pos, h=1e-5) + + deviation = (batch[FORCES_KEY].detach() - numerical).abs().max().item() + assert deviation < FD_FORCE_ATOL + + +class TestFuncForcePass: + """``func_force_pass`` — single ``torch.func.grad(..., has_aux=True)`` pass.""" + + def test_forces_match_hardcoded_analytic_gradient(self) -> None: + toy = KernelToyPotential() + batch = func_force_pass(toy, kernel_toy_batch()) + torch.testing.assert_close(batch[FORCES_KEY], _expected(), atol=EXACT_FORCE_ATOL, rtol=0.0) + + def test_agrees_with_grad_force_pass(self) -> None: + func_batch = func_force_pass(KernelToyPotential(), kernel_toy_batch()) + grad_batch = grad_force_pass(KernelToyPotential(), kernel_toy_batch()) + torch.testing.assert_close( + func_batch[FORCES_KEY].detach(), + grad_batch[FORCES_KEY].detach(), + atol=EXACT_FORCE_ATOL, + rtol=0.0, + ) + + def test_runs_exactly_one_energy_forward(self) -> None: + toy = KernelToyPotential() + func_force_pass(toy, kernel_toy_batch()) + assert toy.n_forward == 1 + + def test_atomic_energy_written_back_when_core_provides_it(self) -> None: + toy = KernelToyPotential(write_atomic=True) + batch = func_force_pass(toy, kernel_toy_batch()) + assert "energy" in batch["atoms"].keys() + assert batch["atoms", "energy"].shape == (4,) + + def test_core_without_atomic_energy_is_accepted(self) -> None: + toy = KernelToyPotential(write_atomic=False) + batch = func_force_pass(toy, kernel_toy_batch()) + assert "energy" not in batch["atoms"].keys() + torch.testing.assert_close(batch[FORCES_KEY], _expected(), atol=EXACT_FORCE_ATOL, rtol=0.0) + + def test_force_loss_reaches_parameters_in_one_backward(self) -> None: + toy = KernelToyPotential() + batch = func_force_pass(toy, kernel_toy_batch()) + batch[FORCES_KEY].pow(2).mean().backward() + assert toy.scale.grad is not None + assert toy.scale.grad.abs().item() > 0.0 + + def test_forces_match_central_finite_differences(self) -> None: + toy = KernelToyPotential() + batch = kernel_toy_batch() + pos = batch[POS_KEY].detach().clone() + + batch = func_force_pass(toy, batch) + numerical = central_difference_forces(toy.total_energy, pos, h=1e-5) + + deviation = (batch[FORCES_KEY].detach() - numerical).abs().max().item() + assert deviation < FD_FORCE_ATOL diff --git a/tests/test_molpot/test_derivation/test_protocol.py b/tests/test_molpot/test_derivation/test_protocol.py new file mode 100644 index 0000000..49ede17 --- /dev/null +++ b/tests/test_molpot/test_derivation/test_protocol.py @@ -0,0 +1,118 @@ +"""Unit tests for the batch-schema helpers in molpot.derivation.protocol. + +Mirrors ``src/molpot/derivation/protocol.py``. Only the two helpers that own +the ``graphs`` sub-TensorDict are covered here — :func:`ensure_graphs` (its +shape) and :func:`write_energy` (the shape it asks for). The session +side-channel and :func:`absorb_model_output` are exercised elsewhere. + +The contract under test is CLAUDE.md's post-collate schema: ``"graphs"`` is a +``TensorDict(batch_size=[B])`` for ``B`` graphs, the same shape the canonical +collate writes (``src/molix/data/collate.py``). Consumers already read +``batch["graphs"].batch_size[0]`` (``src/molzoo/pinet/potential.py``), so a +``batch_size=[]`` produced here is a latent ``IndexError``. + +Units: energy eV. Every expected shape is a hard-coded ``torch.Size``; no +value is recomputed from the input by the assertion side. +""" + +from __future__ import annotations + +import torch +from tensordict import TensorDict + +from molpot.derivation.protocol import ( + ATOMIC_ENERGY_KEY, + ENERGY_KEY, + ensure_graphs, + write_energy, +) + + +def _empty_batch() -> TensorDict: + """Root batch with no ``graphs`` namespace yet.""" + return TensorDict(batch_size=[]) + + +def _batch_with_graphs(num_graphs: int = 2) -> TensorDict: + """Root batch already carrying a schema-conforming ``graphs`` namespace. + + Args: + num_graphs: ``B`` — the ``graphs`` batch size, and the length of the + ``num_atoms`` entry that makes the namespace non-empty. + + Returns: + ``TensorDict`` with ``graphs.batch_size == [num_graphs]`` and + ``graphs.num_atoms == arange(num_graphs) + 1``. + """ + return TensorDict( + graphs=TensorDict( + num_atoms=torch.arange(num_graphs, dtype=torch.long) + 1, + batch_size=[num_graphs], + ), + batch_size=[], + ) + + +def _batch_with_atoms(n_atoms: int = 3) -> TensorDict: + """Root batch with an ``atoms`` namespace, so ``atoms.energy`` is writable.""" + return TensorDict( + atoms=TensorDict( + pos=torch.zeros(n_atoms, 3, dtype=torch.float64), + batch_size=[n_atoms], + ), + batch_size=[], + ) + + +class TestEnsureGraphs: + """Target: :func:`molpot.derivation.protocol.ensure_graphs`.""" + + def test_num_graphs_sets_graphs_batch_size(self): + batch = ensure_graphs(_empty_batch(), num_graphs=3) + assert batch["graphs"].batch_size == torch.Size([3]) + + def test_omitted_num_graphs_keeps_scalar_batch_size(self): + batch = ensure_graphs(_empty_batch()) + assert batch["graphs"].batch_size == torch.Size([]) + + def test_existing_graphs_keeps_its_contents(self): + batch = ensure_graphs(_batch_with_graphs(2), num_graphs=2) + assert torch.equal(batch["graphs", "num_atoms"], torch.tensor([1, 2])) + + def test_existing_graphs_batch_size_is_not_rewritten(self): + batch = ensure_graphs(_batch_with_graphs(2), num_graphs=5) + assert batch["graphs"].batch_size == torch.Size([2]) + + def test_zero_graphs_gives_empty_first_dim(self): + batch = ensure_graphs(_empty_batch(), num_graphs=0) + assert batch["graphs"].batch_size == torch.Size([0]) + + def test_single_graph_gives_unit_first_dim(self): + batch = ensure_graphs(_empty_batch(), num_graphs=1) + assert batch["graphs"].batch_size == torch.Size([1]) + + +class TestWriteEnergyGraphShape: + """Target: :func:`molpot.derivation.protocol.write_energy`.""" + + def test_vector_energy_sets_graphs_batch_size(self): + batch = write_energy(_empty_batch(), torch.zeros(4, dtype=torch.float64)) + assert batch["graphs"].batch_size == torch.Size([4]) + + def test_vector_energy_is_written_elementwise(self): + energy = torch.tensor([-1.5, 0.0, 2.25, 7.0], dtype=torch.float64) + batch = write_energy(_empty_batch(), energy) + assert torch.equal(batch[ENERGY_KEY], energy) + + def test_zero_dim_energy_keeps_scalar_batch_size(self): + batch = write_energy(_empty_batch(), torch.tensor(-3.5, dtype=torch.float64)) + assert batch["graphs"].batch_size == torch.Size([]) + + def test_atomic_energy_is_written_under_atoms(self): + atomic = torch.tensor([-1.0, -2.0, -3.0], dtype=torch.float64) + batch = write_energy( + _batch_with_atoms(3), + torch.tensor([-6.0], dtype=torch.float64), + atomic_energy=atomic, + ) + assert torch.equal(batch[ATOMIC_ENERGY_KEY], atomic) diff --git a/tests/test_molpot/test_derivation/test_readouts.py b/tests/test_molpot/test_derivation/test_readouts.py new file mode 100644 index 0000000..eae3015 --- /dev/null +++ b/tests/test_molpot/test_derivation/test_readouts.py @@ -0,0 +1,126 @@ +"""In-place EnergyReadout / ForceReadout (no user-facing Derivative).""" + +from __future__ import annotations + +import torch +import torch.nn as nn +from tensordict import TensorDict + +from molpot.derivation import EnergyReadout, ForceReadout +from molpot.derivation.protocol import ENERGY_KEY, FORCES_KEY, write_energy + + +class _ToyPotential(nn.Module): + """E = s · Σ ‖pos‖² written onto the batch; counts forward calls.""" + + def __init__(self) -> None: + super().__init__() + self.scale = nn.Parameter(torch.tensor(1.0)) + self.n_forward = 0 + + def forward(self, batch: TensorDict) -> TensorDict: + self.n_forward += 1 + pos = batch["atoms", "pos"] + atom_batch = batch["atoms", "batch"] + atom_e = self.scale * (pos * pos).sum(dim=-1) + n_graphs = int(batch["graphs"].batch_size[0]) + energy = torch.zeros(n_graphs, dtype=pos.dtype, device=pos.device) + energy.scatter_add_(0, atom_batch, atom_e) + write_energy(batch, energy, atomic_energy=atom_e) + return batch + + +def _batch(n_atoms: int = 4, n_graphs: int = 2) -> TensorDict: + per = n_atoms // n_graphs + pos = torch.arange(n_atoms * 3, dtype=torch.float32).reshape(n_atoms, 3) * 0.1 + return TensorDict( + atoms=TensorDict( + pos=pos.clone(), + batch=torch.arange(n_graphs).repeat_interleave(per), + batch_size=[n_atoms], + ), + graphs=TensorDict(batch_size=[n_graphs]), + batch_size=[], + ) + + +def _model(*, train: bool = True) -> _ToyPotential: + m = _ToyPotential() + m.train(train) + return m + + +class TestGradMode: + def test_energy_then_force_is_single_forward(self): + model = _model() + batch = _batch() + batch = EnergyReadout(model, method="grad", backward=True)(batch) + batch = ForceReadout(model, method="grad")(batch) + assert model.n_forward == 1 + assert batch[FORCES_KEY].shape == (4, 3) + assert batch[ENERGY_KEY].shape == (2,) + + def test_forces_match_analytic(self): + model = _model() + batch = _batch() + pos0 = batch["atoms", "pos"].clone() + batch = EnergyReadout(model, method="grad", backward=True)(batch) + batch = ForceReadout(model, method="grad")(batch) + expected = -2.0 * model.scale.detach() * pos0 + assert torch.allclose(batch[FORCES_KEY], expected, atol=1e-5) + + def test_force_loss_reaches_parameters(self): + model = _model(train=True) + batch = _batch() + batch = EnergyReadout(model, method="grad", backward=True)(batch) + batch = ForceReadout(model, method="grad")(batch) + batch[FORCES_KEY].pow(2).mean().backward() + assert model.scale.grad is not None + assert model.scale.grad.abs().item() > 0 + + def test_energy_only(self): + model = _model() + batch = EnergyReadout(model, method="grad", backward=False)(_batch()) + assert model.n_forward == 1 + assert ENERGY_KEY[1] in batch["graphs"].keys() + + +class TestFuncMode: + def test_energy_then_force_is_single_forward(self): + model = _model() + batch = _batch() + batch = EnergyReadout(model, method="func", backward=True)(batch) + assert model.n_forward == 0 # lazy + batch = ForceReadout(model, method="func")(batch) + assert model.n_forward == 1 + assert batch[FORCES_KEY].shape == (4, 3) + + def test_forces_match_analytic(self): + model = _model() + batch = _batch() + pos0 = batch["atoms", "pos"].clone() + batch = EnergyReadout(model, method="func", backward=True)(batch) + batch = ForceReadout(model, method="func")(batch) + expected = -2.0 * model.scale.detach() * pos0 + assert torch.allclose(batch[FORCES_KEY], expected, atol=1e-5) + + def test_force_loss_reaches_parameters(self): + model = _model(train=True) + batch = _batch() + batch = EnergyReadout(model, method="func", backward=True)(batch) + batch = ForceReadout(model, method="func")(batch) + batch[FORCES_KEY].pow(2).mean().backward() + assert model.scale.grad is not None + assert model.scale.grad.abs().item() > 0 + + def test_energy_only(self): + model = _model() + batch = EnergyReadout(model, method="func", backward=False)(_batch()) + assert model.n_forward == 1 + assert batch[ENERGY_KEY].shape == (2,) + + def test_force_alone_is_single_forward(self): + model = _model() + batch = ForceReadout(model, method="func")(_batch()) + assert model.n_forward == 1 + assert FORCES_KEY[1] in batch["atoms"].keys() diff --git a/tests/test_molpot/test_elec/calculators/test_padding.py b/tests/test_molpot/test_elec/calculators/test_padding.py index 8909042..2998b5b 100644 --- a/tests/test_molpot/test_elec/calculators/test_padding.py +++ b/tests/test_molpot/test_elec/calculators/test_padding.py @@ -1,11 +1,10 @@ import os import time -import numpy import torch -from ase.io import read from torch.nn.utils.rnn import pad_sequence +from molix.datasets._extxyz import parse_extxyz_frames from molpot.potentials.elec import CoulombPotential, EwaldCalculator from molpot.potentials.elec.lib import compute_batched_kvectors from tests.test_molpot.test_elec.conftest import periodic_neighbor_list @@ -18,7 +17,7 @@ xyz_path = os.path.join(os.path.dirname(__file__), "pbc_structures.xyz") -systems = read(xyz_path, index=":") +systems = parse_extxyz_frames(xyz_path) i_list, j_list, d_list, pos_list, cell_list, charges_list, periodic_list = ( [], @@ -30,20 +29,21 @@ [], ) -for atoms in systems: +for frame in systems: # Full periodic neighbour list via our own pure-torch op. - pos_t = torch.tensor(atoms.get_positions(), dtype=torch.float64) - cell_t = torch.tensor(numpy.array(atoms.get_cell()), dtype=torch.float64) + pos_t = torch.tensor(frame.pos, dtype=torch.float64) + cell_t = torch.tensor(frame.cell, dtype=torch.float64) pairs_, _, dist_ = periodic_neighbor_list( - pos_t, cell_t, cutoff=5.0, full_list=True, periodic=bool(atoms.get_pbc().all()) + pos_t, cell_t, cutoff=5.0, full_list=True, periodic=all(frame.pbc) ) i_list.append(pairs_[:, 0].to(torch.long)) j_list.append(pairs_[:, 1].to(torch.long)) d_list.append(dist_.to(torch.float32)) - pos_list.append(torch.tensor(atoms.get_positions(), dtype=torch.float32)) - cell_list.append(torch.tensor(numpy.array(atoms.get_cell()), dtype=torch.float32)) - charges_list.append(torch.tensor(atoms.get_initial_charges(), dtype=torch.float32) + 1) - periodic_list.append(torch.tensor(atoms.get_pbc(), dtype=torch.bool)) + pos_list.append(torch.tensor(frame.pos, dtype=torch.float32)) + cell_list.append(torch.tensor(frame.cell, dtype=torch.float32)) + # Fixture has no charge column; use unit charges (was ASE zeros + 1). + charges_list.append(torch.ones(frame.n_atoms, dtype=torch.float32)) + periodic_list.append(torch.tensor(frame.pbc, dtype=torch.bool)) # Pad neighbor indices/distances to the same length for batching i_batch = pad_sequence(i_list, batch_first=True, padding_value=0) diff --git a/tests/test_molpot/test_elec/calculators/test_workflow.py b/tests/test_molpot/test_elec/calculators/test_workflow.py index 288a7cc..76937fd 100644 --- a/tests/test_molpot/test_elec/calculators/test_workflow.py +++ b/tests/test_molpot/test_elec/calculators/test_workflow.py @@ -232,6 +232,9 @@ def test_kspace_filter_error_catch(): full_neighbor_list=True, mesh_spacing=0.5, ) + # The NaN scan is a full-mesh reduction (device sync barrier), so it is + # opt-in per instance rather than always on — enable it to exercise the guard. + calculator.kspace_filter.check_nan = True charges = torch.ones([4, 1]) positions = torch.arange(4 * 3).reshape(4, 3).to(torch.float32) @@ -249,7 +252,7 @@ def test_kspace_filter_error_catch(): match = ( "NaNs detected in the k-space filter result. This are probably caused " "by an unsuitable `mesh_spacing`, resulting in a problematic grid of " - r"shape: \[1, 16, 16, 32\]. Try adjsuting the grid by using a " + r"shape: \[1, 16, 16, 32\]. Try adjusting the grid by using a " "different `mesh_spacing` value." ) with pytest.raises(ValueError, match=match): diff --git a/tests/test_molpot/test_heads/conftest.py b/tests/test_molpot/test_heads/conftest.py new file mode 100644 index 0000000..e5d8a7d --- /dev/null +++ b/tests/test_molpot/test_heads/conftest.py @@ -0,0 +1,25 @@ +"""Shared fixtures for the :mod:`molpot.heads` unit suite.""" + +from __future__ import annotations + +from collections.abc import Iterator + +import pytest +import torch + +from molix import config + + +@pytest.fixture +def fp64() -> Iterator[None]: + """Run the case under the global fp64 precision, restoring the previous one. + + Every head in :mod:`molpot.heads` bakes ``config["ftype"]`` into its + parameters at construction time (the contract documented in + :mod:`molix.config`), so the precision has to be switched *before* the + head is built and handed back afterwards. + """ + previous = config["ftype"] + config.set_precision("fp64") + yield + config.set_precision("fp64" if previous == torch.float64 else "fp32") diff --git a/tests/test_molpot/test_heads/test_charge_bond.py b/tests/test_molpot/test_heads/test_charge_bond.py index d78542e..c88353b 100644 --- a/tests/test_molpot/test_heads/test_charge_bond.py +++ b/tests/test_molpot/test_heads/test_charge_bond.py @@ -23,10 +23,15 @@ directly — the kernel is a pure tensor function, and the head *calls* it. The shape/dtype contract holds and the resulting energy is finite and differentiable through the head MLP. +7. **Construction-time precision.** Built under + ``config.set_precision("fp64")`` the head's parameters come out fp64 + on their own — the post-hoc ``.to(dtype=torch.float64)`` inside + ``_build_head`` is a convenience for the cases above, not the contract. """ from __future__ import annotations +import pytest import torch from molpot.heads import BondChargeHead @@ -279,3 +284,35 @@ def test_head_calls_kernel_directly(): assert torch.isfinite(grad_mu).all() assert torch.isfinite(grad_feat).all() assert grad_feat.abs().max() > 0.0 + + +@pytest.mark.parametrize("edge_dim", [None, 4], ids=["no-edge-features", "edge-features"]) +def test_all_parameters_honour_the_fp64_precision(fp64, edge_dim): + """Every parameter is fp64 when the head is built under fp64. + + ``config["ftype"]`` is the single source of truth for the working + precision; a layer that ignores it leaves the head mixed-precision and + its first forward dies on a dtype-mismatched matmul. Built directly — + no ``_build_head`` cast — because the cast is exactly what would mask + the defect. + """ + head = BondChargeHead(node_dim=8, edge_dim=edge_dim, hidden_dim=32) + + assert {p.dtype for p in head.parameters()} == {torch.float64} + + +def test_forward_runs_under_the_fp64_precision(fp64): + """A head built at fp64 emits fp64 charges from the fp64 sample inputs.""" + head = BondChargeHead(node_dim=8, hidden_dim=32, charge_projection=False) + inp = _make_inputs() + + out = head( + node_features=inp["node_features"], + edge_index=inp["edge_index"], + edge_dist=inp["edge_dist"], + atom_batch=inp["atom_batch"], + num_graphs=1, + ) + + assert out["atomic_charges"].dtype == torch.float64 + assert out["bond_charges"].dtype == torch.float64 diff --git a/tests/test_molpot/test_heads/test_charge_response.py b/tests/test_molpot/test_heads/test_charge_response.py new file mode 100644 index 0000000..26785b9 --- /dev/null +++ b/tests/test_molpot/test_heads/test_charge_response.py @@ -0,0 +1,66 @@ +"""Tests for :class:`molpot.heads.charge_response.ChargeResponseHead`. + +Scope: the construction-time dtype contract. ``sigma_raw`` and +``edge_vector_mlp`` already read ``config["ftype"]``; the scalar MLPs +(``atom_diag_mlp`` / ``edge_scalar_mlp`` / ``iso_mlp``) must do the same or +the head comes out mixed-precision under ``config.set_precision("fp64")``. +""" + +from __future__ import annotations + +import pytest +import torch + +from molpot.heads.charge_response import ChargeResponseHead + + +def _make_head(*, variant: str = "localchi", iso: bool = False) -> ChargeResponseHead: + """A minimal head — the dtype contract is size-independent.""" + return ChargeResponseHead( + node_scalar_dim=4, + edge_scalar_dim=3, + edge_vector_dim=2, + hidden_dim=8, + variant=variant, + iso=iso, + ) + + +class TestChargeResponseHead: + @pytest.mark.parametrize( + "variant, iso", + [("localchi", False), ("localchi", True), ("eem", True)], + ids=["localchi", "localchi-iso", "eem-iso"], + ) + def test_all_parameters_honour_the_fp64_precision(self, fp64, variant, iso): + """Every parameter is fp64 when the head is built under fp64. + + ``config["ftype"]`` is the single source of truth for the working + precision; a layer that ignores it leaves the head mixed-precision + and its first forward dies on a dtype-mismatched matmul. The + ``iso`` cases additionally cover the optional ``iso_mlp`` branch. + """ + head = _make_head(variant=variant, iso=iso) + + assert {p.dtype for p in head.parameters()} == {torch.float64} + + def test_forward_runs_under_the_fp64_precision(self, fp64): + """A head built at fp64 turns fp64 features into an fp64 polarizability.""" + head = _make_head() + + pos = torch.tensor([[0.0, 0.0, 0.0], [0.0, 0.0, 1.0]], dtype=torch.float64) + edge_index = torch.tensor([[0, 1], [1, 0]], dtype=torch.long) + out = head( + pos=pos, + Z=torch.tensor([1, 8], dtype=torch.long), + atom_batch=torch.zeros(2, dtype=torch.long), + num_graphs=1, + edge_index=edge_index, + edge_diff=pos[edge_index[:, 1]] - pos[edge_index[:, 0]], + node_scalars=torch.ones(2, 4, dtype=torch.float64), + edge_scalars=torch.ones(2, 3, dtype=torch.float64), + ) + + assert out["alpha"].shape == (1, 3, 3) + assert out["alpha"].dtype == torch.float64 + assert out["atom_diag"].dtype == torch.float64 diff --git a/tests/test_molpot/test_heads/test_dipole.py b/tests/test_molpot/test_heads/test_dipole.py new file mode 100644 index 0000000..775166d --- /dev/null +++ b/tests/test_molpot/test_heads/test_dipole.py @@ -0,0 +1,68 @@ +"""Tests for :class:`molpot.heads.dipole.DipoleHead`. + +Scope: the construction-time dtype contract. ``atomic_dipole_gate`` and +``bond_vector_mlp`` already read ``config["ftype"]``; the scalar MLPs +(``charge_mlp`` / ``bond_scalar_mlp``) must do the same or the head comes +out mixed-precision under ``config.set_precision("fp64")``. +""" + +from __future__ import annotations + +import pytest +import torch + +from molpot.heads.dipole import DipoleHead + + +def _make_head(variant: str) -> DipoleHead: + """A minimal head with every optional dim supplied, so ``variant`` alone + decides which submodules exist.""" + return DipoleHead( + node_scalar_dim=4, + node_vector_dim=2, + edge_scalar_dim=3, + edge_vector_dim=2, + hidden_dim=8, + variant=variant, + ) + + +class TestDipoleHead: + @pytest.mark.parametrize( + "variant", + ["ac", "ad", "bc", "ac_ad", "ac_bc"], + ids=["ac", "ad", "bc", "ac_ad", "ac_bc"], + ) + def test_all_parameters_honour_the_fp64_precision(self, fp64, variant): + """Every parameter is fp64 when the head is built under fp64. + + ``config["ftype"]`` is the single source of truth for the working + precision; a layer that ignores it leaves the head mixed-precision + and its first forward dies on a dtype-mismatched matmul. The + variants select which term MLPs are constructed, so the sweep + covers each submodule the head can own. + """ + head = _make_head(variant) + + assert {p.dtype for p in head.parameters()} == {torch.float64} + + def test_forward_runs_under_the_fp64_precision(self, fp64): + """A head built at fp64 assembles an fp64 dipole from fp64 features.""" + head = _make_head("ac_bc") + + pos = torch.tensor([[0.0, 0.0, 0.0], [0.0, 0.0, 1.0]], dtype=torch.float64) + edge_index = torch.tensor([[0, 1], [1, 0]], dtype=torch.long) + out = head( + pos=pos, + atom_batch=torch.zeros(2, dtype=torch.long), + num_graphs=1, + node_scalars=torch.ones(2, 4, dtype=torch.float64), + edge_scalars=torch.ones(2, 3, dtype=torch.float64), + edge_vectors=torch.ones(2, 3, 2, dtype=torch.float64), + edge_index=edge_index, + edge_diff=pos[edge_index[:, 1]] - pos[edge_index[:, 0]], + ) + + assert out["dipole"].shape == (1, 3) + assert out["dipole"].dtype == torch.float64 + assert out["atomic_charges"].dtype == torch.float64 diff --git a/tests/test_molpot/test_heads/test_energy.py b/tests/test_molpot/test_heads/test_energy.py new file mode 100644 index 0000000..ffde34c --- /dev/null +++ b/tests/test_molpot/test_heads/test_energy.py @@ -0,0 +1,66 @@ +"""Tests for the per-atom energy heads in :mod:`molpot.heads.energy`. + +``TestEnergyHead.test_forward_energy`` was moved from the stale +``tests/test_molpot/test_readout/`` mirror — the production module is +``src/molpot/heads/energy.py``, so the mirror is +``tests/test_molpot/test_heads/test_energy.py``. +""" + +from __future__ import annotations + +import torch + +from molpot.heads.energy import AtomicEnergyMLP, AtomicReferenceEnergy, EnergyHead + + +class TestAtomicEnergyMLP: + def test_all_parameters_honour_the_fp64_precision(self, fp64): + """Every parameter is fp64 when the head is built under fp64. + + ``config["ftype"]`` is the single source of truth for the working + precision; a layer that ignores it leaves the head mixed-precision + and its first forward dies on a dtype-mismatched matmul. + """ + head = AtomicEnergyMLP(hidden_dim=4) + + assert {p.dtype for p in head.parameters()} == {torch.float64} + + +class TestEnergyHead: + def test_forward_energy(self): + head = EnergyHead(hidden_dim=4) + h = torch.ones(3, 4) + batch = torch.tensor([0, 0, 1]) + out = head(h, batch) + assert out.shape == torch.Size([2]) + + def test_all_parameters_honour_the_fp64_precision(self, fp64): + """Every parameter is fp64 when the head is built under fp64.""" + head = EnergyHead(hidden_dim=4) + + assert {p.dtype for p in head.parameters()} == {torch.float64} + + def test_forward_runs_under_the_fp64_precision(self, fp64): + """A head built at fp64 pools fp64 atomic energies into fp64 totals.""" + head = EnergyHead(hidden_dim=4) + + out = head(torch.ones(3, 4, dtype=torch.float64), torch.tensor([0, 0, 1])) + + assert out.shape == torch.Size([2]) + assert out.dtype == torch.float64 + + +class TestAtomicReferenceEnergy: + def test_lookup_buffer_honours_the_fp64_precision(self, fp64): + """The Z-indexed ``E0`` buffer follows ``config["ftype"]``. + + Already honoured by the production code (the buffer is built with + ``dtype=config.ftype``); pinned here so the dtype contract of the + module is covered end to end. + """ + head = AtomicReferenceEnergy( + atomic_energies=[-13.6, -1000.0], + atomic_numbers=[1, 8], + ) + + assert head.atomic_energies.dtype == torch.float64 diff --git a/tests/test_molpot/test_heads/test_multipole.py b/tests/test_molpot/test_heads/test_multipole.py new file mode 100644 index 0000000..1e7aca6 --- /dev/null +++ b/tests/test_molpot/test_heads/test_multipole.py @@ -0,0 +1,112 @@ +"""Tests for :class:`molpot.heads.multipole.PermMultipoleHead`. + +Scope: the construction-time dtype contract. ``q_head`` / ``mu_proj`` / +``theta_proj`` already read ``config["ftype"]``; the two ``cuet.Linear`` +moment-collapse layers (``mu_collapse`` / ``theta_collapse``) must do the +same or the head comes out mixed-precision under +``config.set_precision("fp64")`` and its first forward dies on the l=1 / l=2 +collapse with ``expected scalar type Float but found Double``. +""" + +from __future__ import annotations + +import cuequivariance as cue +import pytest +import torch +from tensordict import TensorDict + +from molpot.heads import PermMultipoleHead + +#: Encoder tensor-track irreps with the uniform multiplicity the head +#: requires, carrying both the ``1o`` block the dipole readout slices and +#: the ``2e`` block the quadrupole readout slices. +TENSOR_IRREPS = cue.Irreps(cue.O3, [(4, "0e"), (4, "1o"), (4, "2e")]) + +#: Which ``(charge, dipole, quadrupole)`` combinations to sweep. ``dipole`` +#: and ``quadrupole`` are the flags that construct a ``cuet.Linear``. +MOMENT_CASES: list[tuple[str, ...]] = [ + ("charge",), + ("charge", "dipole"), + ("charge", "quadrupole"), + ("charge", "dipole", "quadrupole"), + ("dipole", "quadrupole"), +] + + +def _make_head(moments: tuple[str, ...]) -> PermMultipoleHead: + """A minimal head with only ``moments`` enabled.""" + return PermMultipoleHead( + input_dim=8, + hidden_dim=8, + avg_num_neighbors=4.0, + charge="charge" in moments, + dipole="dipole" in moments, + quadrupole="quadrupole" in moments, + tensor_irreps=TENSOR_IRREPS, + ) + + +def _stub_batch(dtype: torch.dtype) -> TensorDict: + """Single-graph batch wired through every key the head reads. + + Three atoms, four (bidirectional) edges, with both the scalar and the + tensor edge tracks present so the q / μ / Θ readouts all run. + """ + n_atoms, n_edges = 3, 4 + atoms = TensorDict( + Z=torch.tensor([1, 6, 8]), + pos=torch.tensor( + [[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0]], + dtype=dtype, + ), + batch=torch.zeros(n_atoms, dtype=torch.long), + batch_size=[n_atoms], + ) + edges = TensorDict( + edge_index=torch.tensor([[0, 1], [1, 0], [1, 2], [2, 1]], dtype=torch.long), + batch_size=[n_edges], + ) + edges["edge_features"] = torch.ones(n_edges, 8, dtype=dtype) + edges["edge_tensor_features"] = torch.ones(n_edges, TENSOR_IRREPS.dim, dtype=dtype) + graphs = TensorDict( + num_atoms=torch.tensor([n_atoms]), + total_charge=torch.zeros(1, dtype=dtype), + batch_size=[1], + ) + return TensorDict(atoms=atoms, edges=edges, graphs=graphs, batch_size=[]) + + +class TestPermMultipoleHead: + @pytest.mark.parametrize( + "moments", + MOMENT_CASES, + ids=["q", "q_mu", "q_theta", "q_mu_theta", "mu_theta"], + ) + def test_all_parameters_honour_the_fp64_precision( + self, fp64: None, moments: tuple[str, ...] + ) -> None: + """Every parameter is fp64 when the head is built under fp64. + + ``config["ftype"]`` is the single source of truth for the working + precision; a layer that ignores it leaves the head mixed-precision + and its first forward dies on a dtype-mismatched matmul. The sweep + covers each moment flag, so both ``cuet.Linear`` collapse weights + (``mu_collapse`` for l=1, ``theta_collapse`` for l=2) are exercised. + """ + head = _make_head(moments) + + assert {p.dtype for p in head.parameters()} == {torch.float64} + + def test_forward_runs_under_the_fp64_precision(self, fp64: None) -> None: + """A head built at fp64 reads fp64 features and emits fp64 moments.""" + head = _make_head(("charge", "dipole", "quadrupole")) + + out = head(_stub_batch(torch.float64)) + + assert out["atomic_charges"].shape == (3,) + assert out["atomic_charges"].dtype == torch.float64 + assert out["atomic_dipoles"].shape == (3, 3) + assert out["atomic_dipoles"].dtype == torch.float64 + assert out["atomic_quadrupoles"].shape == (3, 5) + assert out["atomic_quadrupoles"].dtype == torch.float64 + assert out["molecular_dipole"].dtype == torch.float64 diff --git a/tests/test_molpot/test_heads/test_pinet.py b/tests/test_molpot/test_heads/test_pinet.py index 078ee69..c757f9b 100644 --- a/tests/test_molpot/test_heads/test_pinet.py +++ b/tests/test_molpot/test_heads/test_pinet.py @@ -5,6 +5,7 @@ import pytest import torch +from molpot.derivation.protocol import ENERGY_KEY, FORCES_KEY from molrep.utils.equivariance import random_rotation_matrix, rotate_vectors from molzoo import PiNet from molzoo.pinet import PiNetDipole, PiNetPolarizability, PiNetPotential @@ -62,13 +63,16 @@ def test_energy_and_force_shapes(self): depth=2, rank=3, hidden_dim=8, + # Monomorphic since b85d12f: force derivation is fixed here, not + # requested per call. + compute_forces=True, ) g = _graph() g["atoms", "pos"] = g["atoms", "pos"].clone().requires_grad_(True) - out = model(g, compute_forces=True) - assert out["energy"].shape == (1,) - assert out["forces"].shape == (4, 3) - (out["energy"].sum() + out["forces"].square().sum()).backward() + out = model(g) + assert out[ENERGY_KEY].shape == (1,) + assert out[FORCES_KEY].shape == (4, 3) + (out[ENERGY_KEY].sum() + out[FORCES_KEY].square().sum()).backward() class TestPiNetDipole: diff --git a/tests/test_molpot/test_heads/test_provenance.py b/tests/test_molpot/test_heads/test_provenance.py new file mode 100644 index 0000000..1c3812c --- /dev/null +++ b/tests/test_molpot/test_heads/test_provenance.py @@ -0,0 +1,268 @@ +"""Tests for molpot.heads.provenance (coverage + ParameterProvenance). + +Spec: learnable-classical-ff-09-provenance. +""" + +from __future__ import annotations + +import ast +from dataclasses import FrozenInstanceError +from pathlib import Path + +import pytest +import torch + +from molpot.heads.provenance import ( + CoverageRegime, + ParameterProvenance, + SupportClassifier, + attach_provenance, +) +from molpot.heads.type import TypeHead +from molrep.condensation.classes import InteractionClass +from molrep.embedding.support import ChemicalSupportIndex + + +class TestCoverageRegime: + def test_members(self): + names = {m.name for m in CoverageRegime} + assert names == { + "IN_SUPPORT", + "NEAR_SUPPORT", + "EXTRAPOLATING", + "UNKNOWN", + } + + +class TestSupportClassifier: + """Confidence bands + support membership → CoverageRegime.""" + + def test_in_support_high_confidence(self): + support = ChemicalSupportIndex.from_type_ids({0, 1, 2}, radius=0.0) + clf = SupportClassifier(support, conf_in=0.8, conf_near=0.5) + type_ids = torch.tensor([0, 1], dtype=torch.long) + conf = torch.tensor([0.95, 0.90]) + regimes = clf.classify(type_ids, conf) + assert regimes == [ + CoverageRegime.IN_SUPPORT, + CoverageRegime.IN_SUPPORT, + ] + + def test_near_support_mid_confidence(self): + support = ChemicalSupportIndex.from_type_ids({0, 1}, radius=0.0) + clf = SupportClassifier(support, conf_in=0.8, conf_near=0.5) + type_ids = torch.tensor([0], dtype=torch.long) + conf = torch.tensor([0.65]) + assert clf.classify(type_ids, conf) == [CoverageRegime.NEAR_SUPPORT] + + def test_low_confidence_is_unknown(self): + """Documented policy: conf < conf_near → UNKNOWN (usable floor).""" + support = ChemicalSupportIndex.from_type_ids({0}, radius=0.0) + clf = SupportClassifier(support, conf_in=0.8, conf_near=0.5) + type_ids = torch.tensor([0], dtype=torch.long) + conf = torch.tensor([0.2]) + assert clf.classify(type_ids, conf) == [CoverageRegime.UNKNOWN] + + def test_out_of_support_never_in_support(self): + support = ChemicalSupportIndex.from_type_ids({0}, radius=0.0) + clf = SupportClassifier(support, conf_in=0.8, conf_near=0.5) + type_ids = torch.tensor([99, 99], dtype=torch.long) + conf = torch.tensor([0.99, 0.60]) + regimes = clf.classify(type_ids, conf) + assert CoverageRegime.IN_SUPPORT not in regimes + assert regimes[0] == CoverageRegime.EXTRAPOLATING + assert regimes[1] == CoverageRegime.EXTRAPOLATING + + def test_out_of_support_low_conf_unknown(self): + support = ChemicalSupportIndex.from_type_ids({0}, radius=0.0) + clf = SupportClassifier(support, conf_in=0.8, conf_near=0.5) + type_ids = torch.tensor([99], dtype=torch.long) + conf = torch.tensor([0.1]) + assert clf.classify(type_ids, conf) == [CoverageRegime.UNKNOWN] + + def test_missing_support_confidence_only(self): + """No support index: confidence bands alone (membership assumed True).""" + clf = SupportClassifier(None, conf_in=0.8, conf_near=0.5) + type_ids = torch.tensor([3, 4, 5], dtype=torch.long) + conf = torch.tensor([0.9, 0.6, 0.1]) + assert clf.classify(type_ids, conf) == [ + CoverageRegime.IN_SUPPORT, + CoverageRegime.NEAR_SUPPORT, + CoverageRegime.UNKNOWN, + ] + + def test_embedding_query_path(self): + """Continuous embedding bank: classify via L2 membership.""" + bank = torch.tensor([[0.0, 0.0], [1.0, 0.0]], dtype=torch.float64) + support = ChemicalSupportIndex(bank, radius=0.25, k=1) + clf = SupportClassifier(support, conf_in=0.8, conf_near=0.5) + # type_ids unused for membership when query= is given + type_ids = torch.tensor([0, 0], dtype=torch.long) + conf = torch.tensor([0.95, 0.95]) + query = torch.tensor([[0.0, 0.0], [5.0, 5.0]], dtype=torch.float64) + regimes = clf.classify(type_ids, conf, query=query) + assert regimes == [ + CoverageRegime.IN_SUPPORT, + CoverageRegime.EXTRAPOLATING, + ] + + def test_reuses_typehead_decode_with_confidence(self): + """(indices, confidence) from TypeHead — no local softmax helper.""" + head = TypeHead(hidden_dim=4, num_types=3) + # Hand-crafted logits so softmax max is deterministic. + logits = torch.tensor( + [ + [10.0, 0.0, 0.0], # high conf → class 0 + [0.0, 1.0, 0.5], # mid conf → class 1 + [0.1, 0.0, 0.0], # low conf → class 0 + ], + dtype=torch.float64, + ) + # Bypass nn weights: call decode on synthetic logits directly. + indices, confidence = head.decode_with_confidence(logits) + + support = ChemicalSupportIndex.from_type_ids({0, 1}, radius=0.0) + clf = SupportClassifier(support, conf_in=0.8, conf_near=0.5) + regimes = clf.classify(indices, confidence) + + assert int(indices[0]) == 0 + assert float(confidence[0]) > 0.8 + assert regimes[0] == CoverageRegime.IN_SUPPORT + # mid-band softmax max on [0,1,0.5] ≈ 0.51 → NEAR_SUPPORT + assert regimes[1] == CoverageRegime.NEAR_SUPPORT + # low-band on [0.1,0,0] ≈ 0.36 → UNKNOWN + assert regimes[2] == CoverageRegime.UNKNOWN + + # Prove provenance package does not reimplement softmax-max. + prov_root = ( + Path(__file__).resolve().parents[3] + / "src" + / "molpot" + / "heads" + / "provenance" + ) + for py in prov_root.glob("*.py"): + src = py.read_text() + # No parallel softmax implementation (docs may mention the word). + assert "torch.softmax" not in src + assert "F.softmax" not in src + assert "nn.functional.softmax" not in src + + +class TestParameterProvenance: + def test_construction_and_fields(self): + rec = ParameterProvenance( + interaction="bond", + type_id=3, + confidence=0.91, + regime=CoverageRegime.IN_SUPPORT, + source="condensed_type", + pattern="[#6]-[#8]", + ) + assert rec.interaction == "bond" + assert rec.type_id == 3 + assert rec.confidence == 0.91 + assert rec.regime is CoverageRegime.IN_SUPPORT + assert rec.source == "condensed_type" + assert rec.pattern == "[#6]-[#8]" + assert rec.ir_units == "class_i_canonical" + assert rec.notes == "" + + def test_frozen(self): + rec = ParameterProvenance( + interaction="angle", + type_id=None, + confidence=None, + regime=CoverageRegime.UNKNOWN, + source="neural_continuous", + ) + with pytest.raises((FrozenInstanceError, AttributeError)): + rec.type_id = 1 # type: ignore[misc] + + def test_as_dict_json_friendly(self): + rec = ParameterProvenance( + interaction="lj", + type_id=1, + confidence=0.5, + regime=CoverageRegime.NEAR_SUPPORT, + source="symbolic", + pattern=None, + notes="unit", + ) + d = rec.as_dict() + assert d["interaction"] == "lj" + assert d["regime"] == "NEAR_SUPPORT" + assert d["source"] == "symbolic" + assert d["type_id"] == 1 + assert d["confidence"] == 0.5 + assert d["ir_units"] == "class_i_canonical" + + def test_interaction_class_serializes(self): + rec = ParameterProvenance( + interaction=InteractionClass.BOND, + type_id=0, + confidence=0.9, + regime=CoverageRegime.IN_SUPPORT, + source="condensed_type", + ) + assert rec.as_dict()["interaction"] == "bond" + + def test_attach_provenance_on_metadata_object(self): + class _Spec: + def __init__(self) -> None: + self.metadata: dict = {} + + rec = ParameterProvenance( + interaction="bond", + type_id=1, + confidence=0.9, + regime=CoverageRegime.IN_SUPPORT, + source="neural_continuous", + ) + spec = _Spec() + out = attach_provenance(spec, [rec]) + assert out is spec + assert spec.metadata["provenance"][0]["type_id"] == 1 + assert spec.metadata["provenance"][0]["regime"] == "IN_SUPPORT" + + def test_attach_provenance_on_dict(self): + rec = ParameterProvenance( + interaction="angle", + type_id=None, + confidence=None, + regime=CoverageRegime.UNKNOWN, + source="symbolic", + ) + meta: dict = {} + attach_provenance(meta, [rec]) + assert meta["provenance"][0]["interaction"] == "angle" + + +class TestNoActiveLearningApis: + def test_public_modules_have_no_acquisition(self): + prov_root = ( + Path(__file__).resolve().parents[3] + / "src" + / "molpot" + / "heads" + / "provenance" + ) + forbidden = ( + "acquisition", + "active_learning", + "query_selector", + "retrain_loop", + "select_batch", + ) + for py in prov_root.glob("*.py"): + tree = ast.parse(py.read_text()) + names: set[str] = set() + for node in ast.walk(tree): + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)): + names.add(node.name.lower()) + if isinstance(node, ast.Assign): + for t in node.targets: + if isinstance(t, ast.Name): + names.add(t.id.lower()) + for bad in forbidden: + assert not any(bad in n for n in names), (py.name, names) diff --git a/tests/test_molpot/test_heads/test_type.py b/tests/test_molpot/test_heads/test_type.py new file mode 100644 index 0000000..91b2eb5 --- /dev/null +++ b/tests/test_molpot/test_heads/test_type.py @@ -0,0 +1,47 @@ +"""Tests for :class:`molpot.heads.type.TypeHead`. + +Moved from the stale ``tests/test_molpot/test_readout/`` mirror — the +production module is ``src/molpot/heads/type.py``, so the mirror is +``tests/test_molpot/test_heads/test_type.py``. +""" + +from __future__ import annotations + +import torch + +from molpot.heads.type import TypeHead + + +class TestTypeHead: + def test_forward_logits(self): + head = TypeHead(hidden_dim=4, num_types=5) + h = torch.ones(3, 4) + out = head(h) + assert out.shape == torch.Size([3, 5]) + + def test_all_parameters_honour_the_fp64_precision(self, fp64): + """Every parameter is fp64 when the head is built under fp64. + + ``config["ftype"]`` is the single source of truth for the working + precision; a layer that ignores it leaves the head mixed-precision + and its first forward dies on a dtype-mismatched matmul. + """ + head = TypeHead(hidden_dim=4, num_types=5) + + assert {p.dtype for p in head.parameters()} == {torch.float64} + + def test_forward_runs_under_the_fp64_precision(self, fp64): + """A head built at fp64 consumes fp64 features and emits fp64 logits.""" + head = TypeHead(hidden_dim=4, num_types=5) + + out = head(torch.ones(3, 4, dtype=torch.float64)) + + assert out.shape == torch.Size([3, 5]) + assert out.dtype == torch.float64 + + def test_decode_with_confidence(self): + head = TypeHead(hidden_dim=4, num_types=3) + logits = torch.tensor([[5.0, 0.0, 0.0], [0.0, 2.0, 0.0]]) + indices, confidence = head.decode_with_confidence(logits) + assert indices.tolist() == [0, 1] + assert float(confidence[0]) > float(confidence[1]) diff --git a/tests/test_molpot/test_ir/__init__.py b/tests/test_molpot/test_ir/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/test_molpot/test_ir/test_bags.py b/tests/test_molpot/test_ir/test_bags.py new file mode 100644 index 0000000..d097ded --- /dev/null +++ b/tests/test_molpot/test_ir/test_bags.py @@ -0,0 +1,117 @@ +"""IR parameter bags — field shapes and validation (learnable-classical-ff-01).""" + +import pytest +import torch + +from molpot.ir import ( + AngleBag, + BondBag, + ChargeBag, + ImproperHarmonicBag, + ImproperPeriodicBag, + LJBag, + ProperTorsionBag, +) + + +class TestBondBag: + def test_construct_with_equal_length_tensors(self): + bag = BondBag(k=torch.tensor([100.0, 200.0]), r0=torch.tensor([1.0, 1.5])) + assert bag.k.shape == bag.r0.shape + assert bag.k.shape[0] == 2 + + def test_mismatch_lengths_raise(self): + with pytest.raises((ValueError, TypeError)): + BondBag(k=torch.tensor([100.0]), r0=torch.tensor([1.0, 1.5])) + + +class TestAngleBag: + def test_construct_with_equal_length_tensors(self): + bag = AngleBag( + k=torch.tensor([50.0, 60.0]), + theta0=torch.tensor([1.9, 2.0]), + ) + assert bag.k.shape == bag.theta0.shape + + def test_mismatch_lengths_raise(self): + with pytest.raises((ValueError, TypeError)): + AngleBag(k=torch.tensor([50.0, 60.0]), theta0=torch.tensor([1.9])) + + +class TestProperTorsionBag: + def test_construct_with_equal_term_tables(self): + # k / phase: [n_types, n_terms]; periodicity: [n_terms]; idivf/s: [n_types] + bag = ProperTorsionBag( + k=torch.tensor([[1.0, 0.5]]), + periodicity=torch.tensor([1, 2], dtype=torch.long), + phase=torch.tensor([[0.0, 0.0]]), + idivf=torch.tensor([1.0]), + ) + assert bag.k.shape[-1] == bag.periodicity.shape[0] + assert bag.phase.shape == bag.k.shape + + def test_mismatch_term_counts_raise(self): + with pytest.raises((ValueError, TypeError)): + ProperTorsionBag( + k=torch.tensor([[1.0, 0.5]]), + periodicity=torch.tensor([1], dtype=torch.long), + phase=torch.tensor([[0.0, 0.0]]), + idivf=torch.tensor([1.0]), + ) + + def test_non_positive_periodicity_raises(self): + with pytest.raises((ValueError, TypeError)): + ProperTorsionBag( + k=torch.tensor([[1.0]]), + periodicity=torch.tensor([0], dtype=torch.long), + phase=torch.tensor([[0.0]]), + idivf=torch.tensor([1.0]), + ) + + def test_negative_periodicity_raises(self): + with pytest.raises((ValueError, TypeError)): + ProperTorsionBag( + k=torch.tensor([[1.0]]), + periodicity=torch.tensor([-1], dtype=torch.long), + phase=torch.tensor([[0.0]]), + idivf=torch.tensor([1.0]), + ) + + +class TestImproperPeriodicBag: + def test_construct_with_equal_term_tables(self): + bag = ImproperPeriodicBag( + k=torch.tensor([[1.0]]), + periodicity=torch.tensor([2], dtype=torch.long), + phase=torch.tensor([[3.141592653589793]]), + idivf=torch.tensor([1.0]), + ) + assert bag.k.shape[-1] == 1 + + +class TestImproperHarmonicBag: + def test_construct_with_equal_length_tensors(self): + bag = ImproperHarmonicBag( + k=torch.tensor([10.0]), + chi0=torch.tensor([0.0]), + ) + assert bag.k.shape == bag.chi0.shape + + def test_mismatch_lengths_raise(self): + with pytest.raises((ValueError, TypeError)): + ImproperHarmonicBag(k=torch.tensor([10.0, 20.0]), chi0=torch.tensor([0.0])) + + +class TestLJBag: + def test_construct_with_equal_length_tensors(self): + bag = LJBag( + epsilon=torch.tensor([0.1, 0.2]), + sigma=torch.tensor([3.0, 3.5]), + ) + assert bag.epsilon.shape == bag.sigma.shape + + +class TestChargeBag: + def test_construct(self): + bag = ChargeBag(q=torch.tensor([0.1, -0.1, 0.0])) + assert bag.q.shape[0] == 3 diff --git a/tests/test_molpot/test_ir/test_potential_ir.py b/tests/test_molpot/test_ir/test_potential_ir.py new file mode 100644 index 0000000..46cf625 --- /dev/null +++ b/tests/test_molpot/test_ir/test_potential_ir.py @@ -0,0 +1,43 @@ +"""PotentialIR aggregate — optional bags + unit_system validation.""" + +import pytest +import torch + +from molpot.ir import ( + BondBag, + NonbondedScaling, + PotentialIR, +) + + +class TestPotentialIR: + def test_empty_ir_is_valid(self): + ir = PotentialIR() + assert ir is not None + # Empty IR means zero contribution from every term bag. + assert getattr(ir, "bonds", None) is None or ir.bonds is None + assert ( + getattr(ir, "unit_system", "class_i_canonical") + in ( + "class_i_canonical", + None, + ) + or ir.unit_system == "class_i_canonical" + ) + + def test_empty_ir_defaults_to_class_i_canonical_units(self): + ir = PotentialIR() + unit_system = getattr(ir, "unit_system", "class_i_canonical") + assert unit_system == "class_i_canonical" + + def test_unknown_unit_system_raises(self): + with pytest.raises(ValueError): + PotentialIR(unit_system="not_a_real_unit_system") + + def test_subset_of_bags_is_valid(self): + ir = PotentialIR( + bonds=BondBag(k=torch.tensor([100.0]), r0=torch.tensor([1.5])), + scaling=NonbondedScaling(), + ) + assert ir.bonds is not None + assert ir.scaling is not None diff --git a/tests/test_molpot/test_ir/test_scaling.py b/tests/test_molpot/test_ir/test_scaling.py new file mode 100644 index 0000000..4587314 --- /dev/null +++ b/tests/test_molpot/test_ir/test_scaling.py @@ -0,0 +1,36 @@ +"""NonbondedScaling Class-I defaults (AMBER/GAFF / SMIRNOFF). + +Provenance: OpenMM §19 / SMIRNOFF nonbonded section / Cornell 1995 AMBER. +1–2 and 1–3 interactions are fully excluded (scale 0); 1–4 electrostatics +are scaled by 5/6 and 1–4 LJ by 1/2. +""" + +import math + +from molpot.ir import NonbondedScaling + + +class TestNonbondedScaling: + def test_class_i_default_scale_q_12_is_zero(self): + scaling = NonbondedScaling() + assert float(scaling.scale_q_12) == 0.0 + + def test_class_i_default_scale_q_13_is_zero(self): + scaling = NonbondedScaling() + assert float(scaling.scale_q_13) == 0.0 + + def test_class_i_default_scale_q_14_is_five_sixths(self): + scaling = NonbondedScaling() + assert math.isclose(float(scaling.scale_q_14), 5.0 / 6.0, rel_tol=1e-12, abs_tol=1e-12) + + def test_class_i_default_scale_lj_12_is_zero(self): + scaling = NonbondedScaling() + assert float(scaling.scale_lj_12) == 0.0 + + def test_class_i_default_scale_lj_13_is_zero(self): + scaling = NonbondedScaling() + assert float(scaling.scale_lj_13) == 0.0 + + def test_class_i_default_scale_lj_14_is_half(self): + scaling = NonbondedScaling() + assert math.isclose(float(scaling.scale_lj_14), 0.5, rel_tol=1e-12, abs_tol=1e-12) diff --git a/tests/test_molpot/test_ir/test_units.py b/tests/test_molpot/test_ir/test_units.py new file mode 100644 index 0000000..a54812f --- /dev/null +++ b/tests/test_molpot/test_ir/test_units.py @@ -0,0 +1,65 @@ +"""UnitTag + CLASS_I_CANONICAL contract (learnable-classical-ff-01-ir-kernels).""" + +from collections.abc import Mapping + +from molpot.ir import CLASS_I_CANONICAL, UnitTag + + +class TestUnitTag: + """UnitTag enumerates the physical dimensions named by the Potential IR.""" + + def test_required_dimension_members_exist(self): + names = {m.name.lower() if hasattr(m, "name") else str(m).lower() for m in UnitTag} + # Accept either Enum members or string values for the dimension tags. + values = set() + for m in UnitTag: + if hasattr(m, "value"): + values.add(str(m.value).lower()) + values.add(str(m).lower().split(".")[-1]) + combined = names | values + for required in ( + "energy", + "length", + "charge", + "angle", + ): + assert any(required in c for c in combined), f"missing UnitTag for {required}" + + +class TestClassICanonical: + """CLASS_I_CANONICAL freezes Class-I SI-free internal units.""" + + def test_energy_is_kcal_per_mol(self): + assert CLASS_I_CANONICAL["energy"] == "kcal/mol" + + def test_length_is_angstrom(self): + length = CLASS_I_CANONICAL["length"] + assert length in ("angstrom", "Å", "A") + + def test_charge_is_elementary_e(self): + assert CLASS_I_CANONICAL["charge"] == "e" + + def test_angle_is_radian(self): + angle = CLASS_I_CANONICAL["angle"] + assert angle in ("radian", "rad") + + def test_mapping_is_frozen_not_silently_mutable(self): + # MappingProxyType / frozendict / Final: mutation must raise or be a no-op + # that leaves the public mapping unchanged. + original = dict(CLASS_I_CANONICAL) + raised = False + try: + CLASS_I_CANONICAL["energy"] = "kJ/mol" # type: ignore[index] + except (TypeError, AttributeError, ValueError): + raised = True + if not raised: + # If assignment did not raise, the public view must still be unchanged. + assert dict(CLASS_I_CANONICAL) == original + assert CLASS_I_CANONICAL["energy"] == "kcal/mol" + else: + assert CLASS_I_CANONICAL["energy"] == "kcal/mol" + + def test_is_mapping_with_required_keys(self): + assert isinstance(CLASS_I_CANONICAL, Mapping) + for key in ("energy", "length", "charge", "angle"): + assert key in CLASS_I_CANONICAL diff --git a/tests/test_molpot/test_les_parity/conftest.py b/tests/test_molpot/test_les_parity/conftest.py index 8bd851d..60a1375 100644 --- a/tests/test_molpot/test_les_parity/conftest.py +++ b/tests/test_molpot/test_les_parity/conftest.py @@ -67,11 +67,11 @@ # scipy is not a hard dep of MolNex tests; fall back to vectorised math.erf # if scipy isn't on the path. Both produce identical numerical results. try: - from scipy.special import erf as _scipy_erf # type: ignore[import-not-found] + from scipy.special import erf as _scipy_erf except ImportError: # pragma: no cover - depends on local install _erf_vec = np.vectorize(math.erf, otypes=[np.float64]) - def _scipy_erf(x: np.ndarray) -> np.ndarray: # type: ignore[no-redef] + def _scipy_erf(x: np.ndarray) -> np.ndarray: return _erf_vec(x) diff --git a/tests/test_molpot/test_potentials/test_dihedrals/__init__.py b/tests/test_molpot/test_potentials/test_dihedrals/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/test_molpot/test_potentials/test_dihedrals/test_periodic.py b/tests/test_molpot/test_potentials/test_dihedrals/test_periodic.py new file mode 100644 index 0000000..e9ea610 --- /dev/null +++ b/tests/test_molpot/test_potentials/test_dihedrals/test_periodic.py @@ -0,0 +1,170 @@ +"""ProperTorsionPeriodic — Class-I multi-term cosine proper torsion. + +Energy: E = sum_n (k_n / s) * [1 + cos(n * phi - gamma_n)] + +Hard-coded goldens (no external oracle). Geometry fixtures yield known +dihedral angles under the standard atan2(n1,n2) convention used by +DihedralHarmonic (i-j-k-l, b1=j-i, b2=k-j, b3=l-k). +""" + +import math + +import pytest +import torch + +from molpot.potentials import ProperTorsionPeriodic + +# --------------------------------------------------------------------------- +# Geometry fixtures — known dihedral angles +# --------------------------------------------------------------------------- + + +def _cis_pos() -> torch.Tensor: + """Four atoms in a plane with proper torsion phi = 0 (cis). + + i=(0,1,0), j=(0,0,0), k=(1,0,0), l=(1,1,0) + → n1 = n2 = (0,0,1) → phi = 0. + """ + return torch.tensor( + [ + [0.0, 1.0, 0.0], # i + [0.0, 0.0, 0.0], # j + [1.0, 0.0, 0.0], # k + [1.0, 1.0, 0.0], # l + ], + dtype=torch.float64, + ) + + +def _trans_pos() -> torch.Tensor: + """phi = pi (trans): l flipped below the jk axis.""" + return torch.tensor( + [ + [0.0, 1.0, 0.0], + [0.0, 0.0, 0.0], + [1.0, 0.0, 0.0], + [1.0, -1.0, 0.0], + ], + dtype=torch.float64, + ) + + +def _perp_pos() -> torch.Tensor: + """phi = pi/2: l out of plane along +z.""" + return torch.tensor( + [ + [0.0, 1.0, 0.0], + [0.0, 0.0, 0.0], + [1.0, 0.0, 0.0], + [1.0, 0.0, 1.0], + ], + dtype=torch.float64, + ) + + +def _proper_index() -> torch.Tensor: + """COO-style [4, 1] proper torsion over atoms 0-1-2-3.""" + return torch.tensor([[0], [1], [2], [3]], dtype=torch.long) + + +def _single_term_potential(*, k: float = 1.0, n: int = 1, gamma: float = 0.0, s: float = 1.0): + """Build a one-type, one-term ProperTorsionPeriodic. + + Formula-correct goldens use k=1 so cis (phi=0) yields E=2.0: + E = (k/s)[1 + cos(n*phi - gamma)] = 1*[1+1] = 2.0 + """ + return ProperTorsionPeriodic( + k=torch.tensor([[k]], dtype=torch.float64), + periodicity=torch.tensor([n], dtype=torch.long), + phase=torch.tensor([[gamma]], dtype=torch.float64), + idivf=torch.tensor([s], dtype=torch.float64), + ) + + +class TestProperTorsionPeriodic: + def test_cis_phi0_k1_gives_energy_2(self): + # E = (1/1)[1 + cos(0)] = 2.0 (hard-coded golden; Class-I identity) + pot = _single_term_potential(k=1.0, n=1, gamma=0.0, s=1.0) + e = pot( + pos=_cis_pos(), + proper_index=_proper_index(), + proper_types=torch.tensor([0], dtype=torch.long), + ) + assert math.isclose(float(e), 2.0, rel_tol=1e-10, abs_tol=1e-10) + + def test_trans_phi_pi_gives_energy_0(self): + # E = (1/1)[1 + cos(pi)] = 0.0 + pot = _single_term_potential(k=1.0, n=1, gamma=0.0, s=1.0) + e = pot( + pos=_trans_pos(), + proper_index=_proper_index(), + proper_types=torch.tensor([0], dtype=torch.long), + ) + assert math.isclose(float(e), 0.0, rel_tol=1e-10, abs_tol=1e-10) + + def test_perp_phi_half_pi_gives_energy_1(self): + # E = (1/1)[1 + cos(pi/2)] = 1.0 + pot = _single_term_potential(k=1.0, n=1, gamma=0.0, s=1.0) + e = pot( + pos=_perp_pos(), + proper_index=_proper_index(), + proper_types=torch.tensor([0], dtype=torch.long), + ) + assert math.isclose(float(e), 1.0, rel_tol=1e-10, abs_tol=1e-10) + + def test_empty_proper_index_returns_zero(self): + pot = _single_term_potential() + e = pot( + pos=_cis_pos(), + proper_index=torch.zeros(4, 0, dtype=torch.long), + proper_types=torch.zeros(0, dtype=torch.long), + ) + assert float(e) == 0.0 + + def test_rejects_row_major_n_by_4_shape(self): + pot = _single_term_potential() + wrong = torch.tensor([[0, 1, 2, 3]], dtype=torch.long) # [1, 4] == [N, 4] + with pytest.raises(ValueError, match=r"\[4"): + pot( + pos=_cis_pos(), + proper_index=wrong, + proper_types=torch.tensor([0], dtype=torch.long), + ) + + def test_rejects_edge_index_shape(self): + pot = _single_term_potential() + edge_like = torch.tensor([[0, 1], [1, 2]], dtype=torch.long) # [E, 2] + with pytest.raises(ValueError, match=r"\[4"): + pot( + pos=_cis_pos(), + proper_index=edge_like, + proper_types=torch.tensor([0, 0], dtype=torch.long), + ) + + def test_rejects_bond_like_2_by_n_shape(self): + pot = _single_term_potential() + bond_like = torch.tensor([[0], [1]], dtype=torch.long) # [2, 1] COO bond shape + with pytest.raises(ValueError, match=r"\[4"): + pot( + pos=_cis_pos(), + proper_index=bond_like, + proper_types=torch.tensor([0], dtype=torch.long), + ) + + def test_multi_term_sums_component_energies(self): + # Two terms at phi=0, gamma=0: + # term1: k=1, n=1 → (1/1)[1+cos(0)] = 2.0 + # term2: k=0.5, n=2 → (0.5/1)[1+cos(0)] = 1.0 + # total = 3.0 + pot = ProperTorsionPeriodic( + k=torch.tensor([[1.0, 0.5]], dtype=torch.float64), + periodicity=torch.tensor([1, 2], dtype=torch.long), + phase=torch.tensor([[0.0, 0.0]], dtype=torch.float64), + idivf=torch.tensor([1.0], dtype=torch.float64), + ) + e = pot( + pos=_cis_pos(), + proper_index=_proper_index(), + proper_types=torch.tensor([0], dtype=torch.long), + ) + assert math.isclose(float(e), 3.0, rel_tol=1e-10, abs_tol=1e-10) diff --git a/tests/test_molpot/test_potentials/test_impropers/__init__.py b/tests/test_molpot/test_potentials/test_impropers/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/test_molpot/test_potentials/test_impropers/test_harmonic.py b/tests/test_molpot/test_potentials/test_impropers/test_harmonic.py new file mode 100644 index 0000000..7c8c0cc --- /dev/null +++ b/tests/test_molpot/test_potentials/test_impropers/test_harmonic.py @@ -0,0 +1,116 @@ +"""ImproperHarmonic — E = 1/2 k (chi - chi0)^2. + +Hard-coded geometry goldens. improper_index is molrs layout +[center, i, j, k] (center at row 0); chi is dihedral(i, center, j, k). +""" + +import math + +import pytest +import torch + +from molpot.potentials import ImproperHarmonic + + +def _planar_pos() -> torch.Tensor: + """chi = 0 (planar).""" + return torch.tensor( + [ + [0.0, 1.0, 0.0], + [0.0, 0.0, 0.0], # center + [1.0, 0.0, 0.0], + [1.0, 1.0, 0.0], + ], + dtype=torch.float64, + ) + + +def _pi_over_6_pos() -> torch.Tensor: + """chi = pi/6: k rotated out of the plane. + + i=(0,1,0), center=(0,0,0), j=(1,0,0), k=(1, cos(pi/6), sin(pi/6)) + → atan2 path yields phi = pi/6 on (i,center,j,k). + """ + return torch.tensor( + [ + [0.0, 1.0, 0.0], + [0.0, 0.0, 0.0], # center + [1.0, 0.0, 0.0], + [1.0, math.cos(math.pi / 6.0), math.sin(math.pi / 6.0)], + ], + dtype=torch.float64, + ) + + +def _improper_index() -> torch.Tensor: + """molrs [center, i, j, k] with center = atom 1.""" + return torch.tensor([[1], [0], [2], [3]], dtype=torch.long) + + +class TestImproperHarmonic: + def test_planar_equilibrium_gives_zero_energy(self): + # chi = chi0 = 0 → E = 0 + pot = ImproperHarmonic( + k=torch.tensor([2.0], dtype=torch.float64), + chi0=torch.tensor([0.0], dtype=torch.float64), + ) + e = pot( + pos=_planar_pos(), + improper_index=_improper_index(), + improper_types=torch.tensor([0], dtype=torch.long), + ) + assert math.isclose(float(e), 0.0, rel_tol=1e-10, abs_tol=1e-10) + + def test_pi_over_6_matches_half_k_delta_squared(self): + # k=2, chi=pi/6, chi0=0 → E = 0.5 * 2 * (pi/6)^2 = (pi/6)^2 + pot = ImproperHarmonic( + k=torch.tensor([2.0], dtype=torch.float64), + chi0=torch.tensor([0.0], dtype=torch.float64), + ) + e = pot( + pos=_pi_over_6_pos(), + improper_index=_improper_index(), + improper_types=torch.tensor([0], dtype=torch.long), + ) + expected = (math.pi / 6.0) ** 2 + # float64 geometry + atan2 leaves ~1e-9 residual on (pi/6)**2 + assert math.isclose(float(e), expected, rel_tol=1e-8, abs_tol=1e-8) + + def test_empty_improper_index_returns_zero(self): + pot = ImproperHarmonic( + k=torch.tensor([2.0], dtype=torch.float64), + chi0=torch.tensor([0.0], dtype=torch.float64), + ) + e = pot( + pos=_planar_pos(), + improper_index=torch.zeros(4, 0, dtype=torch.long), + improper_types=torch.zeros(0, dtype=torch.long), + ) + assert float(e) == 0.0 + + def test_rejects_row_major_n_by_4_shape(self): + pot = ImproperHarmonic( + k=torch.tensor([2.0], dtype=torch.float64), + chi0=torch.tensor([0.0], dtype=torch.float64), + ) + wrong = torch.tensor([[0, 1, 2, 3]], dtype=torch.long) + with pytest.raises(ValueError, match=r"\[4"): + pot( + pos=_planar_pos(), + improper_index=wrong, + improper_types=torch.tensor([0], dtype=torch.long), + ) + + def test_forces_are_finite_off_equilibrium(self): + pot = ImproperHarmonic( + k=torch.tensor([2.0], dtype=torch.float64), + chi0=torch.tensor([0.0], dtype=torch.float64), + ) + pos = _pi_over_6_pos().clone().requires_grad_(True) + forces = pot.calc_forces( + pos=pos.detach(), + improper_index=_improper_index(), + improper_types=torch.tensor([0], dtype=torch.long), + as_numpy=False, + ) + assert torch.isfinite(forces).all() diff --git a/tests/test_molpot/test_potentials/test_impropers/test_periodic.py b/tests/test_molpot/test_potentials/test_impropers/test_periodic.py new file mode 100644 index 0000000..d35c039 --- /dev/null +++ b/tests/test_molpot/test_potentials/test_impropers/test_periodic.py @@ -0,0 +1,142 @@ +"""ImproperPeriodic — Class-I multi-term cosine improper torsion. + +Same cosine form as ProperTorsionPeriodic, on improper_index [4, N]. +Central atom is at row index 0 (molrs Topology layout [center, i, j, k]). +""" + +import math + +import pytest +import torch + +from molpot.potentials import ImproperPeriodic + + +def _cis_pos() -> torch.Tensor: + """phi = 0 planar fixture (same geometry as proper cis).""" + return torch.tensor( + [ + [0.0, 1.0, 0.0], # atom 0 = i (peripheral) + [0.0, 0.0, 0.0], # atom 1 = center + [1.0, 0.0, 0.0], # atom 2 = j + [1.0, 1.0, 0.0], # atom 3 = k + ], + dtype=torch.float64, + ) + + +def _trans_pos() -> torch.Tensor: + return torch.tensor( + [ + [0.0, 1.0, 0.0], + [0.0, 0.0, 0.0], # center + [1.0, 0.0, 0.0], + [1.0, -1.0, 0.0], + ], + dtype=torch.float64, + ) + + +def _perp_pos() -> torch.Tensor: + return torch.tensor( + [ + [0.0, 1.0, 0.0], + [0.0, 0.0, 0.0], # center + [1.0, 0.0, 0.0], + [1.0, 0.0, 1.0], + ], + dtype=torch.float64, + ) + + +def _improper_index_center_first() -> torch.Tensor: + """COO [4, 1]; molrs layout [center, i, j, k] with center = atom 1.""" + return torch.tensor([[1], [0], [2], [3]], dtype=torch.long) + + +def _single_term(*, k: float = 1.0, n: int = 1, gamma: float = 0.0, s: float = 1.0): + return ImproperPeriodic( + k=torch.tensor([[k]], dtype=torch.float64), + periodicity=torch.tensor([n], dtype=torch.long), + phase=torch.tensor([[gamma]], dtype=torch.float64), + idivf=torch.tensor([s], dtype=torch.float64), + ) + + +class TestImproperPeriodic: + def test_cis_phi0_k1_gives_energy_2(self): + pot = _single_term(k=1.0, n=1, gamma=0.0, s=1.0) + e = pot( + pos=_cis_pos(), + improper_index=_improper_index_center_first(), + improper_types=torch.tensor([0], dtype=torch.long), + ) + assert math.isclose(float(e), 2.0, rel_tol=1e-10, abs_tol=1e-10) + + def test_trans_phi_pi_gives_energy_0(self): + pot = _single_term(k=1.0, n=1, gamma=0.0, s=1.0) + e = pot( + pos=_trans_pos(), + improper_index=_improper_index_center_first(), + improper_types=torch.tensor([0], dtype=torch.long), + ) + assert math.isclose(float(e), 0.0, rel_tol=1e-10, abs_tol=1e-10) + + def test_perp_phi_half_pi_gives_energy_1(self): + pot = _single_term(k=1.0, n=1, gamma=0.0, s=1.0) + e = pot( + pos=_perp_pos(), + improper_index=_improper_index_center_first(), + improper_types=torch.tensor([0], dtype=torch.long), + ) + assert math.isclose(float(e), 1.0, rel_tol=1e-10, abs_tol=1e-10) + + def test_central_atom_is_row_index_0(self): + """Lock molrs layout: central lives at row 0 of improper_index. + + Center-first [center, i, j, k] with center=atom 1 yields the cis + cosine golden; a silent trefoil reorder would break the energy. + """ + pot = _single_term(k=1.0, n=1, gamma=0.0, s=1.0) + index = _improper_index_center_first() + assert int(index[0, 0]) == 1 # central atom id at row 0 + e = pot( + pos=_cis_pos(), + improper_index=index, + improper_types=torch.tensor([0], dtype=torch.long), + ) + assert math.isclose(float(e), 2.0, rel_tol=1e-10, abs_tol=1e-10) + + def test_empty_improper_index_returns_zero(self): + pot = _single_term() + e = pot( + pos=_cis_pos(), + improper_index=torch.zeros(4, 0, dtype=torch.long), + improper_types=torch.zeros(0, dtype=torch.long), + ) + assert float(e) == 0.0 + + def test_rejects_row_major_n_by_4_shape(self): + pot = _single_term() + wrong = torch.tensor([[0, 1, 2, 3]], dtype=torch.long) # [N, 4] + with pytest.raises(ValueError, match=r"\[4"): + pot( + pos=_cis_pos(), + improper_index=wrong, + improper_types=torch.tensor([0], dtype=torch.long), + ) + + def test_multi_term_sums_component_energies(self): + pot = ImproperPeriodic( + k=torch.tensor([[1.0, 0.5]], dtype=torch.float64), + periodicity=torch.tensor([1, 2], dtype=torch.long), + phase=torch.tensor([[0.0, 0.0]], dtype=torch.float64), + idivf=torch.tensor([1.0], dtype=torch.float64), + ) + e = pot( + pos=_cis_pos(), + improper_index=_improper_index_center_first(), + improper_types=torch.tensor([0], dtype=torch.long), + ) + # 2.0 + 1.0 = 3.0 at phi=0 + assert math.isclose(float(e), 3.0, rel_tol=1e-10, abs_tol=1e-10) diff --git a/tests/test_molpot/test_potentials/test_repulsion.py b/tests/test_molpot/test_potentials/test_repulsion.py new file mode 100644 index 0000000..15e0e7c --- /dev/null +++ b/tests/test_molpot/test_potentials/test_repulsion.py @@ -0,0 +1,85 @@ +"""Tests for molpot.potentials.repulsion module.""" + +import pytest +import torch + +from molpot.potentials.repulsion import ZBLRepulsion + + +@pytest.fixture +def zbl(): + """ZBL term at MACE's foundation-model settings (p = 5).""" + return ZBLRepulsion(exponent=5).double() + + +class TestZBLRepulsion: + """Test the ZBL screened-nuclear-repulsion pair term.""" + + def test_output_is_per_atom(self, zbl): + """One energy per atom, not per edge.""" + Z = torch.tensor([8, 1, 1]) + edge_index = torch.tensor([[0, 1, 0, 2], [1, 0, 2, 0]]).t() # (E, 2) + r = torch.tensor([1.0, 1.0, 1.0, 1.0], dtype=torch.float64) + assert zbl(r, Z, edge_index).shape == (3,) + + def test_repulsion_is_positive(self, zbl): + """Nuclear repulsion never lowers the energy.""" + Z = torch.tensor([8, 8]) + edge_index = torch.tensor([[0, 1], [1, 0]]).t() # (E, 2) + r = torch.tensor([0.8, 0.8], dtype=torch.float64) + assert bool((zbl(r, Z, edge_index) > 0).all()) + + def test_decays_with_distance(self, zbl): + """Closer nuclei repel harder.""" + Z = torch.tensor([8, 8]) + edge_index = torch.tensor([[0, 1], [1, 0]]).t() # (E, 2) + close = zbl(torch.tensor([0.6, 0.6], dtype=torch.float64), Z, edge_index) + far = zbl(torch.tensor([1.1, 1.1], dtype=torch.float64), Z, edge_index) + assert float(close[0]) > float(far[0]) + + def test_vanishes_beyond_summed_covalent_radii(self, zbl): + """The envelope cuts the pair off at ``R_i + R_j``, not at the model cutoff.""" + Z = torch.tensor([1, 1]) # H-H: covalent radii sum to 0.62 A + edge_index = torch.tensor([[0, 1], [1, 0]]).t() # (E, 2) + beyond = zbl(torch.tensor([0.7, 0.7], dtype=torch.float64), Z, edge_index) + assert float(beyond.abs().max()) == 0.0 + + def test_heavier_nuclei_repel_more(self, zbl): + """The ``Z_i Z_j`` prefactor makes heavy pairs stiffer at equal distance.""" + edge_index = torch.tensor([[0, 1], [1, 0]]).t() # (E, 2) + r = torch.tensor([0.5, 0.5], dtype=torch.float64) + light = zbl(r, torch.tensor([1, 1]), edge_index) + heavy = zbl(r, torch.tensor([8, 8]), edge_index) + assert float(heavy[0]) > float(light[0]) + + def test_pair_energy_is_split_between_both_atoms(self, zbl): + """Each direction carries half, so a bidirectional list sums to the pair.""" + Z = torch.tensor([8, 1]) + both = torch.tensor([[0, 1], [1, 0]]) # (E, 2): 0->1 and 1->0 + one = torch.tensor([[0, 1]]) # (E, 2): the single direction 0->1 + r2 = torch.tensor([0.7, 0.7], dtype=torch.float64) + r1 = torch.tensor([0.7], dtype=torch.float64) + assert float(zbl(r2, Z, both).sum()) == pytest.approx( + 2.0 * float(zbl(r1, Z, one).sum()), rel=1e-12 + ) + + def test_is_differentiable_wrt_distance(self, zbl): + """Forces come from ``-dE/dr``, so the term must carry gradient.""" + Z = torch.tensor([8, 8]) + edge_index = torch.tensor([[0, 1], [1, 0]]).t() # (E, 2) + r = torch.tensor([0.8, 0.8], dtype=torch.float64, requires_grad=True) + zbl(r, Z, edge_index).sum().backward() + assert r.grad is not None + # Repulsive: energy decreases as the pair separates. + assert bool((r.grad < 0).all()) + + def test_frozen_parameters_by_default(self, zbl): + """Screening constants are buffers unless the caller asks for training.""" + assert not any(p.requires_grad for p in zbl.parameters()) + assert any(n == "a_exp" for n, _ in zbl.named_buffers()) + + def test_trainable_promotes_screening_to_parameters(self): + """``trainable=True`` exposes the screening length for fine-tuning.""" + trainable = ZBLRepulsion(trainable=True) + names = {n for n, _ in trainable.named_parameters()} + assert {"a_exp", "a_prefactor"} <= names diff --git a/tests/test_molpot/test_readout/test_energy.py b/tests/test_molpot/test_readout/test_energy.py deleted file mode 100644 index 1ee1866..0000000 --- a/tests/test_molpot/test_readout/test_energy.py +++ /dev/null @@ -1,12 +0,0 @@ -import torch - -from molpot.heads.energy import EnergyHead - - -class TestEnergyHead: - def test_forward_energy(self): - head = EnergyHead(hidden_dim=4) - h = torch.ones(3, 4) - batch = torch.tensor([0, 0, 1]) - out = head(h, batch) - assert out.shape == torch.Size([2]) diff --git a/tests/test_molpot/test_readout/test_max.py b/tests/test_molpot/test_readout/test_max.py index 5a7195a..cb2df23 100644 --- a/tests/test_molpot/test_readout/test_max.py +++ b/tests/test_molpot/test_readout/test_max.py @@ -10,3 +10,27 @@ def test_max_pooling(self): batch = torch.tensor([0, 0, 1]) out = pooling(x, batch) assert torch.allclose(out, torch.tensor([[3.0, 2.0], [10.0, 20.0]])) + + def test_max_pooling_1d(self): + out = MaxPooling()(torch.tensor([1.0, 5.0, 2.0, -3.0]), torch.tensor([0, 0, 1, 1])) + assert torch.equal(out, torch.tensor([5.0, 2.0])) + + def test_max_pooling_unsorted_batch(self): + """Rows need not be grouped by graph — the scatter handles any order.""" + x = torch.tensor([[7.0], [1.0], [9.0], [2.0]]) + out = MaxPooling()(x, torch.tensor([1, 0, 1, 0]), 2) + assert torch.equal(out, torch.tensor([[2.0], [9.0]])) + + def test_max_pooling_empty_graph_is_neg_inf(self): + """A graph with no atoms keeps the -inf identity (documented behaviour).""" + x = torch.tensor([[1.0, 2.0], [3.0, 4.0]]) + out = MaxPooling()(x, torch.tensor([0, 2]), 3) + assert torch.equal(out[0], torch.tensor([1.0, 2.0])) + assert torch.isinf(out[1]).all() and (out[1] < 0).all() + assert torch.equal(out[2], torch.tensor([3.0, 4.0])) + + def test_max_pooling_is_differentiable(self): + x = torch.tensor([[1.0, 5.0], [3.0, 2.0]], requires_grad=True) + MaxPooling()(x, torch.tensor([0, 0]), 1).sum().backward() + # Gradient reaches exactly the arg-max element of each feature column. + assert torch.equal(x.grad, torch.tensor([[0.0, 1.0], [1.0, 0.0]])) diff --git a/tests/test_molpot/test_readout/test_type.py b/tests/test_molpot/test_readout/test_type.py deleted file mode 100644 index d6d1267..0000000 --- a/tests/test_molpot/test_readout/test_type.py +++ /dev/null @@ -1,11 +0,0 @@ -import torch - -from molpot.heads.type import TypeHead - - -class TestTypeHead: - def test_forward_logits(self): - head = TypeHead(hidden_dim=4, num_types=5) - h = torch.ones(3, 4) - out = head(h) - assert out.shape == torch.Size([3, 5]) diff --git a/tests/test_molrep/test_analysis/__init__.py b/tests/test_molrep/test_analysis/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/test_molrep/test_analysis/test_latent_store.py b/tests/test_molrep/test_analysis/test_latent_store.py new file mode 100644 index 0000000..97c940b --- /dev/null +++ b/tests/test_molrep/test_analysis/test_latent_store.py @@ -0,0 +1,14 @@ +import torch + +from molrep.analysis import AtomLatentTable + + +class TestAtomLatentTable: + def test_basic(self): + t = AtomLatentTable( + features=torch.zeros(3, 4), + molecule_id=["a", "a", "b"], + ref_atom_type=torch.tensor([0, 0, 1]), + ) + assert t.n_atoms == 3 + assert t.dim == 4 diff --git a/tests/test_molrep/test_analysis/test_projection.py b/tests/test_molrep/test_analysis/test_projection.py new file mode 100644 index 0000000..e997f35 --- /dev/null +++ b/tests/test_molrep/test_analysis/test_projection.py @@ -0,0 +1,30 @@ +import torch + +from molrep.analysis import AtomLatentTable, LatentAnalysisArtifacts, LatentPCA2D + + +class TestLatentPCA2D: + def test_shape(self): + t = AtomLatentTable(torch.randn(5, 4), ["m"] * 5) + c = LatentPCA2D().project(t) + assert c.shape == (5, 2) + + +class TestLatentAnalysisArtifacts: + def test_keys(self, tmp_path): + t = AtomLatentTable( + torch.tensor([[0.0, 0.0], [0.1, 0.0], [10.0, 0.0], [10.1, 0.0]]), + ["a", "a", "b", "b"], + torch.tensor([0, 0, 1, 1]), + ) + m = LatentAnalysisArtifacts(tmp_path).write(t) + assert set(m) >= { + "nn_type_purity_mean", + "nn_type_purity_k", + "n_atoms", + "n_labeled", + "n_scored", + } + assert (tmp_path / "latent_points.jsonl").is_file() + assert (tmp_path / "type_purity.txt").is_file() + assert m["nn_type_purity_mean"] == 1.0 diff --git a/tests/test_molrep/test_analysis/test_type_purity.py b/tests/test_molrep/test_analysis/test_type_purity.py new file mode 100644 index 0000000..b01c8a6 --- /dev/null +++ b/tests/test_molrep/test_analysis/test_type_purity.py @@ -0,0 +1,26 @@ +import torch + +from molrep.analysis import AtomLatentTable, NearestNeighbourTypePurity + + +class TestNearestNeighbourTypePurity: + def test_separated_clusters(self): + feats = torch.tensor([[0.0, 0.0], [0.1, 0.0], [10.0, 0.0], [10.1, 0.0]]) + y = torch.tensor([0, 0, 1, 1]) + t = AtomLatentTable(feats, ["m"] * 4, y) + r = NearestNeighbourTypePurity(k=1).score(t) + assert r.mean_purity == 1.0 + + def test_alternating_line(self): + feats = torch.tensor([[0.0], [1.0], [2.0], [3.0]]) + y = torch.tensor([0, 1, 0, 1]) + t = AtomLatentTable(feats, ["m"] * 4, y) + r = NearestNeighbourTypePurity(k=1).score(t) + assert r.mean_purity == 0.0 + + def test_all_unlabeled(self): + feats = torch.zeros(3, 2) + y = torch.tensor([-1, -1, -1]) + t = AtomLatentTable(feats, ["m"] * 3, y) + r = NearestNeighbourTypePurity(k=1).score(t) + assert r.mean_purity is None diff --git a/tests/test_molrep/test_chem/__init__.py b/tests/test_molrep/test_chem/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/test_molrep/test_chem/conftest.py b/tests/test_molrep/test_chem/conftest.py new file mode 100644 index 0000000..e78527d --- /dev/null +++ b/tests/test_molrep/test_chem/conftest.py @@ -0,0 +1,79 @@ +"""Shared fixtures for molrep.chem unit tests.""" + +from __future__ import annotations + +import pytest +import torch +from tensordict import TensorDict + + +@pytest.fixture +def atom_dim() -> int: + return 8 + + +@pytest.fixture +def bond_dim() -> int: + return 8 + + +@pytest.fixture +def mini_batch() -> TensorDict: + """Synthetic water-like topology with full valence namespaces. + + Atoms: O(0), H(1), H(2), C(3) for a second bond leg. + Bonds: 0-1, 0-2, 0-3 + Angles: 1-0-2, 1-0-3 + Propers: 1-0-3-2 (one) + Impropers: center=0, outer={1,2,3} + """ + z = torch.tensor([8, 1, 1, 6], dtype=torch.long) + n = z.shape[0] + batch = TensorDict( + { + "atoms": TensorDict( + { + "Z": z, + "pos": torch.zeros(n, 3), + "batch": torch.zeros(n, dtype=torch.long), + }, + batch_size=[n], + ), + "bonds": TensorDict( + { + "atomi": torch.tensor([0, 0, 0], dtype=torch.long), + "atomj": torch.tensor([1, 2, 3], dtype=torch.long), + "bond_types": torch.tensor([1, 1, 1], dtype=torch.long), + }, + batch_size=[3], + ), + "angles": TensorDict( + { + "atomi": torch.tensor([1, 1], dtype=torch.long), + "atomj": torch.tensor([0, 0], dtype=torch.long), + "atomk": torch.tensor([2, 3], dtype=torch.long), + }, + batch_size=[2], + ), + "propers": TensorDict( + { + "atomi": torch.tensor([1], dtype=torch.long), + "atomj": torch.tensor([0], dtype=torch.long), + "atomk": torch.tensor([3], dtype=torch.long), + "atoml": torch.tensor([2], dtype=torch.long), + }, + batch_size=[1], + ), + "impropers": TensorDict( + { + "atomi": torch.tensor([0], dtype=torch.long), # center + "atomj": torch.tensor([1], dtype=torch.long), + "atomk": torch.tensor([2], dtype=torch.long), + "atoml": torch.tensor([3], dtype=torch.long), + }, + batch_size=[1], + ), + }, + batch_size=[], + ) + return batch diff --git a/tests/test_molrep/test_chem/test_context.py b/tests/test_molrep/test_chem/test_context.py new file mode 100644 index 0000000..c3919b6 --- /dev/null +++ b/tests/test_molrep/test_chem/test_context.py @@ -0,0 +1,144 @@ +"""Symmetry tests for valence interaction context builders.""" + +from __future__ import annotations + +import torch + +from molrep.chem.context import ( + AngleContext, + BondContext, + ImproperContext, + ProperContext, +) +from molrep.chem.embed import AtomChemEmbedding + + +class TestBondContext: + def test_delegates_to_bond_embedding_shape(self): + atom_dim, bond_dim = 8, 6 + atom_emb = AtomChemEmbedding(atom_dim=atom_dim, num_elements=20) + ctx = BondContext(atom_dim=atom_dim, bond_dim=bond_dim) + h = atom_emb(torch.tensor([1, 6, 8], dtype=torch.long)) + atomi = torch.tensor([0, 1], dtype=torch.long) + atomj = torch.tensor([1, 2], dtype=torch.long) + out = ctx(h, atomi, atomj) + assert out.shape == (2, bond_dim) + + def test_reverse_symmetry(self): + torch.manual_seed(1) + atom_dim, bond_dim = 8, 6 + atom_emb = AtomChemEmbedding(atom_dim=atom_dim, num_elements=20) + ctx = BondContext(atom_dim=atom_dim, bond_dim=bond_dim) + h = atom_emb(torch.tensor([1, 6, 8, 7], dtype=torch.long)) + atomi = torch.tensor([0, 1], dtype=torch.long) + atomj = torch.tensor([2, 3], dtype=torch.long) + assert torch.allclose( + ctx(h, atomi, atomj), + ctx(h, atomj, atomi), + rtol=1e-5, + atol=1e-6, + ) + + +class TestAngleContext: + def test_shape(self): + atom_dim, angle_dim = 8, 6 + atom_emb = AtomChemEmbedding(atom_dim=atom_dim, num_elements=20) + ctx = AngleContext(atom_dim=atom_dim, angle_dim=angle_dim) + h = atom_emb(torch.tensor([1, 6, 8], dtype=torch.long)) + atomi = torch.tensor([0], dtype=torch.long) + atomj = torch.tensor([1], dtype=torch.long) + atomk = torch.tensor([2], dtype=torch.long) + out = ctx(h, atomi, atomj, atomk) + assert out.shape == (1, angle_dim) + + def test_reverse_symmetry_ijk_to_kji(self): + """Angle context invariant under (i,j,k) → (k,j,i).""" + torch.manual_seed(2) + atom_dim, angle_dim = 8, 6 + atom_emb = AtomChemEmbedding(atom_dim=atom_dim, num_elements=20) + ctx = AngleContext(atom_dim=atom_dim, angle_dim=angle_dim) + h = atom_emb(torch.tensor([1, 6, 8, 7, 16], dtype=torch.long)) + atomi = torch.tensor([0, 1], dtype=torch.long) + atomj = torch.tensor([2, 2], dtype=torch.long) + atomk = torch.tensor([3, 4], dtype=torch.long) + forward = ctx(h, atomi, atomj, atomk) + reverse = ctx(h, atomk, atomj, atomi) + assert torch.allclose(forward, reverse, rtol=1e-5, atol=1e-6) + + +class TestProperContext: + def test_shape(self): + atom_dim, proper_dim = 8, 6 + atom_emb = AtomChemEmbedding(atom_dim=atom_dim, num_elements=20) + ctx = ProperContext(atom_dim=atom_dim, proper_dim=proper_dim) + h = atom_emb(torch.tensor([1, 6, 8, 7], dtype=torch.long)) + out = ctx( + h, + torch.tensor([0], dtype=torch.long), + torch.tensor([1], dtype=torch.long), + torch.tensor([2], dtype=torch.long), + torch.tensor([3], dtype=torch.long), + ) + assert out.shape == (1, proper_dim) + + def test_reverse_symmetry_ijkl_to_lkji(self): + """Proper context invariant under (i,j,k,l) → (l,k,j,i).""" + torch.manual_seed(3) + atom_dim, proper_dim = 8, 6 + atom_emb = AtomChemEmbedding(atom_dim=atom_dim, num_elements=20) + ctx = ProperContext(atom_dim=atom_dim, proper_dim=proper_dim) + h = atom_emb(torch.tensor([1, 6, 8, 7, 16, 15], dtype=torch.long)) + atomi = torch.tensor([0, 1], dtype=torch.long) + atomj = torch.tensor([1, 2], dtype=torch.long) + atomk = torch.tensor([2, 3], dtype=torch.long) + atoml = torch.tensor([3, 4], dtype=torch.long) + forward = ctx(h, atomi, atomj, atomk, atoml) + reverse = ctx(h, atoml, atomk, atomj, atomi) + assert torch.allclose(forward, reverse, rtol=1e-5, atol=1e-6) + + +class TestImproperContext: + def test_shape(self): + atom_dim, improper_dim = 8, 6 + atom_emb = AtomChemEmbedding(atom_dim=atom_dim, num_elements=20) + ctx = ImproperContext(atom_dim=atom_dim, improper_dim=improper_dim) + h = atom_emb(torch.tensor([6, 1, 1, 8], dtype=torch.long)) + # center = atomi = 0 + out = ctx( + h, + torch.tensor([0], dtype=torch.long), + torch.tensor([1], dtype=torch.long), + torch.tensor([2], dtype=torch.long), + torch.tensor([3], dtype=torch.long), + ) + assert out.shape == (1, improper_dim) + + def test_outer_swap_center_fixed(self): + """Improper invariant under outer-leg swap; center (atomi) fixed.""" + torch.manual_seed(4) + atom_dim, improper_dim = 8, 6 + atom_emb = AtomChemEmbedding(atom_dim=atom_dim, num_elements=20) + ctx = ImproperContext(atom_dim=atom_dim, improper_dim=improper_dim) + h = atom_emb(torch.tensor([6, 1, 7, 8, 16], dtype=torch.long)) + center = torch.tensor([0, 0], dtype=torch.long) + j = torch.tensor([1, 1], dtype=torch.long) + k = torch.tensor([2, 2], dtype=torch.long) + l = torch.tensor([3, 4], dtype=torch.long) + base = ctx(h, center, j, k, l) + # Swap outer j ↔ k + swapped_jk = ctx(h, center, k, j, l) + # Swap outer k ↔ l + swapped_kl = ctx(h, center, j, l, k) + # Swap outer j ↔ l + swapped_jl = ctx(h, center, l, k, j) + assert torch.allclose(base, swapped_jk, rtol=1e-5, atol=1e-6) + assert torch.allclose(base, swapped_kl, rtol=1e-5, atol=1e-6) + assert torch.allclose(base, swapped_jl, rtol=1e-5, atol=1e-6) + + def test_empty(self): + ctx = ImproperContext(atom_dim=8, improper_dim=6) + h = torch.zeros(3, 8) + empty = torch.zeros(0, dtype=torch.long) + out = ctx(h, empty, empty, empty, empty) + assert out.shape == (0, 6) diff --git a/tests/test_molrep/test_chem/test_embed.py b/tests/test_molrep/test_chem/test_embed.py new file mode 100644 index 0000000..2433cab --- /dev/null +++ b/tests/test_molrep/test_chem/test_embed.py @@ -0,0 +1,61 @@ +"""Tests for AtomChemEmbedding and BondChemEmbedding.""" + +from __future__ import annotations + +import torch + +from molrep.chem.embed import AtomChemEmbedding, BondChemEmbedding +from molrep.embedding.node import JointEmbedding + + +class TestAtomChemEmbedding: + """AtomChemEmbedding maps Z → (N, D_a).""" + + def test_shape_z_only(self, atom_dim: int): + emb = AtomChemEmbedding(atom_dim=atom_dim, num_elements=20) + z = torch.tensor([1, 6, 8, 1], dtype=torch.long) + out = emb(z) + assert out.shape == (4, atom_dim) + + def test_reuses_joint_embedding(self, atom_dim: int): + emb = AtomChemEmbedding(atom_dim=atom_dim, num_elements=20) + assert isinstance(emb.joint, JointEmbedding) + + def test_empty_atoms(self, atom_dim: int): + emb = AtomChemEmbedding(atom_dim=atom_dim, num_elements=20) + z = torch.zeros(0, dtype=torch.long) + out = emb(z) + assert out.shape == (0, atom_dim) + + +class TestBondChemEmbedding: + """BondChemEmbedding is endpoint-symmetric: h_ij == h_ji.""" + + def test_shape(self, atom_dim: int, bond_dim: int): + atom_emb = AtomChemEmbedding(atom_dim=atom_dim, num_elements=20) + bond_emb = BondChemEmbedding(atom_dim=atom_dim, bond_dim=bond_dim) + z = torch.tensor([6, 8, 1], dtype=torch.long) + h = atom_emb(z) + atomi = torch.tensor([0, 0], dtype=torch.long) + atomj = torch.tensor([1, 2], dtype=torch.long) + out = bond_emb(h, atomi, atomj) + assert out.shape == (2, bond_dim) + + def test_endpoint_symmetry(self, atom_dim: int, bond_dim: int): + torch.manual_seed(0) + atom_emb = AtomChemEmbedding(atom_dim=atom_dim, num_elements=20) + bond_emb = BondChemEmbedding(atom_dim=atom_dim, bond_dim=bond_dim) + z = torch.tensor([6, 8, 1, 7], dtype=torch.long) + h = atom_emb(z) + atomi = torch.tensor([0, 1, 2], dtype=torch.long) + atomj = torch.tensor([1, 2, 3], dtype=torch.long) + forward = bond_emb(h, atomi, atomj) + reverse = bond_emb(h, atomj, atomi) + assert torch.allclose(forward, reverse, rtol=1e-5, atol=1e-6) + + def test_empty_bonds(self, atom_dim: int, bond_dim: int): + bond_emb = BondChemEmbedding(atom_dim=atom_dim, bond_dim=bond_dim) + h = torch.zeros(3, atom_dim) + empty = torch.zeros(0, dtype=torch.long) + out = bond_emb(h, empty, empty) + assert out.shape == (0, bond_dim) diff --git a/tests/test_molrep/test_chem/test_encoder.py b/tests/test_molrep/test_chem/test_encoder.py new file mode 100644 index 0000000..a301896 --- /dev/null +++ b/tests/test_molrep/test_chem/test_encoder.py @@ -0,0 +1,131 @@ +"""Tests for ChemEncoder TensorDict I/O and isolation.""" + +from __future__ import annotations + +import ast +from pathlib import Path + +import torch +from tensordict import TensorDict + +from molrep.chem.encoder import ChemEncoder +from molrep.chem.features import ChemEmbeddings + + +class TestChemEncoder: + """ChemEncoder reads valence namespaces and writes chem features.""" + + def test_forward_writes_keys(self, mini_batch: TensorDict): + enc = ChemEncoder( + atom_dim=8, + bond_dim=8, + angle_dim=8, + proper_dim=8, + improper_dim=8, + num_elements=20, + ) + out = enc(mini_batch) + assert ("atoms", "chem_features") in out.keys(include_nested=True) + assert ("bonds", "chem_features") in out.keys(include_nested=True) + assert ("angles", "chem_features") in out.keys(include_nested=True) + assert ("propers", "chem_features") in out.keys(include_nested=True) + assert ("impropers", "chem_features") in out.keys(include_nested=True) + + assert out["atoms", "chem_features"].shape == (4, 8) + assert out["bonds", "chem_features"].shape == (3, 8) + assert out["angles", "chem_features"].shape == (2, 8) + assert out["propers", "chem_features"].shape == (1, 8) + assert out["impropers", "chem_features"].shape == (1, 8) + + def test_compose_returns_chem_embeddings(self, mini_batch: TensorDict): + enc = ChemEncoder(atom_dim=8, bond_dim=6, angle_dim=6, proper_dim=6, improper_dim=6) + emb = enc.compose(mini_batch) + assert isinstance(emb, ChemEmbeddings) + assert emb.atom.shape == (4, 8) + assert emb.bond.shape == (3, 6) + assert emb.angle.shape == (2, 6) + assert emb.proper.shape == (1, 6) + assert emb.improper.shape == (1, 6) + + def test_embeddings_view(self, mini_batch: TensorDict): + enc = ChemEncoder(atom_dim=8, bond_dim=8, angle_dim=8, proper_dim=8, improper_dim=8) + out = enc(mini_batch) + viewed = enc.embeddings(out) + assert torch.allclose(viewed.atom, out["atoms", "chem_features"]) + assert torch.allclose(viewed.bond, out["bonds", "chem_features"]) + + def test_write_batch(self, mini_batch: TensorDict): + enc = ChemEncoder(atom_dim=8, bond_dim=8, angle_dim=8, proper_dim=8, improper_dim=8) + emb = enc.compose(mini_batch) + written = enc.write_batch(mini_batch, emb) + assert written is mini_batch or isinstance(written, TensorDict) + assert written["atoms", "chem_features"].shape[0] == 4 + + def test_bond_index_coo_path(self): + """Bonds may carry COO bond_index (2, N) instead of atomi/atomj.""" + n = 3 + # bond_index contract: COO (2, N_bonds) + bond_index = torch.tensor([[0, 0], [1, 2]], dtype=torch.long) + batch = TensorDict( + { + "atoms": TensorDict( + {"Z": torch.tensor([6, 8, 1], dtype=torch.long)}, + batch_size=[n], + ), + "bonds": TensorDict( + {"bond_index": bond_index}, + batch_size=[], + ), + }, + batch_size=[], + ) + enc = ChemEncoder(atom_dim=4, bond_dim=4, angle_dim=4, proper_dim=4, improper_dim=4) + out = enc(batch) + assert out["bonds", "chem_features"].shape == (2, 4) + + def test_missing_optional_namespaces_empty(self): + """Missing angles/propers/impropers write empty (0, D) features.""" + batch = TensorDict( + { + "atoms": TensorDict( + {"Z": torch.tensor([1, 1], dtype=torch.long)}, + batch_size=[2], + ), + "bonds": TensorDict( + { + "atomi": torch.tensor([0], dtype=torch.long), + "atomj": torch.tensor([1], dtype=torch.long), + }, + batch_size=[1], + ), + }, + batch_size=[], + ) + enc = ChemEncoder(atom_dim=4, bond_dim=4, angle_dim=5, proper_dim=5, improper_dim=5) + emb = enc.compose(batch) + assert emb.angle.shape == (0, 5) + assert emb.proper.shape == (0, 5) + assert emb.improper.shape == (0, 5) + + def test_no_energy_keys(self, mini_batch: TensorDict): + enc = ChemEncoder(atom_dim=4, bond_dim=4, angle_dim=4, proper_dim=4, improper_dim=4) + out = enc(mini_batch) + flat_keys = {str(k) for k in out.keys(include_nested=True)} + assert not any("energy" in k for k in flat_keys) + assert not any("force" in k for k in flat_keys) + + +class TestNoMolpotImport: + """molrep.chem must not import molpot.""" + + def test_source_tree_has_no_molpot(self): + root = Path(__file__).resolve().parents[3] / "src" / "molrep" / "chem" + assert root.is_dir() + for path in root.rglob("*.py"): + tree = ast.parse(path.read_text()) + for node in ast.walk(tree): + if isinstance(node, ast.Import): + for alias in node.names: + assert not alias.name.startswith("molpot"), path + elif isinstance(node, ast.ImportFrom) and node.module: + assert not node.module.startswith("molpot"), path diff --git a/tests/test_molrep/test_chem/test_features.py b/tests/test_molrep/test_chem/test_features.py new file mode 100644 index 0000000..4f0a7d1 --- /dev/null +++ b/tests/test_molrep/test_chem/test_features.py @@ -0,0 +1,48 @@ +"""Tests for ChemEmbeddings first-class container.""" + +from __future__ import annotations + +import torch + +from molrep.chem.features import ChemEmbeddings + + +class TestChemEmbeddings: + """ChemEmbeddings holds the five feature tensors.""" + + def test_fields_and_shapes(self): + emb = ChemEmbeddings( + atom=torch.zeros(4, 8), + bond=torch.zeros(3, 8), + angle=torch.zeros(2, 8), + proper=torch.zeros(1, 8), + improper=torch.zeros(1, 8), + ) + assert emb.atom.shape == (4, 8) + assert emb.bond.shape == (3, 8) + assert emb.angle.shape == (2, 8) + assert emb.proper.shape == (1, 8) + assert emb.improper.shape == (1, 8) + + def test_empty_optional_impropers(self): + """Missing impropers use empty (0, D) rather than absent.""" + emb = ChemEmbeddings( + atom=torch.zeros(2, 4), + bond=torch.zeros(1, 4), + angle=torch.zeros(0, 4), + proper=torch.zeros(0, 4), + improper=torch.zeros(0, 4), + ) + assert emb.improper.shape == (0, 4) + assert emb.angle.shape == (0, 4) + + def test_as_dict_keys(self): + emb = ChemEmbeddings( + atom=torch.zeros(1, 2), + bond=torch.zeros(1, 2), + angle=torch.zeros(0, 2), + proper=torch.zeros(0, 2), + improper=torch.zeros(0, 2), + ) + d = emb.as_dict() + assert set(d) == {"atom", "bond", "angle", "proper", "improper"} diff --git a/tests/test_molrep/test_chem/test_typing_metrics.py b/tests/test_molrep/test_chem/test_typing_metrics.py new file mode 100644 index 0000000..82c837b --- /dev/null +++ b/tests/test_molrep/test_chem/test_typing_metrics.py @@ -0,0 +1,41 @@ +"""Tests for TypingRecoveryMetrics.""" + +from __future__ import annotations + +import torch + +from molrep.chem.typing_metrics import TypingRecoveryMetrics + + +class TestTypingRecoveryMetrics: + def test_perfect_accuracy(self): + m = TypingRecoveryMetrics(num_types=3) + pred = torch.tensor([0, 1, 2, 1]) + true = torch.tensor([0, 1, 2, 1]) + m.update(pred, true, Z=torch.tensor([6, 6, 1, 8]), molecule_ids=["a", "a", "b", "b"]) + r = m.compute() + assert r.overall_accuracy == 1.0 + assert r.n_atoms == 4 + assert r.molecule_error_counts == {} + assert r.per_element_accuracy[6] == 1.0 + + def test_errors_and_confusion(self): + m = TypingRecoveryMetrics(num_types=2, rare_max_count=1) + m.update( + torch.tensor([0, 1, 0]), + torch.tensor([0, 0, 1]), + molecule_ids=["m0", "m0", "m1"], + ) + r = m.compute() + assert abs(r.overall_accuracy - 1 / 3) < 1e-6 + assert r.confusion[0, 0] == 1 + assert r.confusion[0, 1] == 1 + assert r.confusion[1, 0] == 1 + assert r.molecule_error_counts["m0"] == 1 + assert r.molecule_error_counts["m1"] == 1 + + def test_reset(self): + m = TypingRecoveryMetrics(num_types=2) + m.update(torch.tensor([0]), torch.tensor([0])) + m.reset() + assert m.compute().n_atoms == 0 diff --git a/tests/test_molrep/test_chem/test_typing_probe.py b/tests/test_molrep/test_chem/test_typing_probe.py new file mode 100644 index 0000000..8c7d42a --- /dev/null +++ b/tests/test_molrep/test_chem/test_typing_probe.py @@ -0,0 +1,59 @@ +"""Tests for AtomTypeReadout isolation + overfit path.""" + +from __future__ import annotations + +import torch +from tensordict import TensorDict + +from molrep.chem import AtomTypeReadout, ChemEncoder + + +def _batch(Z: list[int], bonds: list[tuple[int, int]] | None = None) -> TensorDict: + n = len(Z) + atoms = TensorDict({"Z": torch.tensor(Z, dtype=torch.long)}, batch_size=[n]) + td = TensorDict({"atoms": atoms}, batch_size=[]) + if bonds: + atomi = torch.tensor([a for a, _ in bonds], dtype=torch.long) + atomj = torch.tensor([b for _, b in bonds], dtype=torch.long) + td["bonds"] = TensorDict( + {"atomi": atomi, "atomj": atomj}, + batch_size=[len(bonds)], + ) + return td + + +class TestAtomTypeReadout: + def test_forward_shapes(self): + enc = ChemEncoder(atom_dim=8, bond_dim=4) + probe = AtomTypeReadout(enc, num_types=3) + batch = _batch([6, 1], bonds=[(0, 1)]) + out = probe(batch) + assert out["logits"].shape == (2, 3) + assert out["pred_type_id"].shape == (2,) + + def test_labels_do_not_affect_encoder(self): + enc = ChemEncoder(atom_dim=8, bond_dim=4) + probe = AtomTypeReadout(enc, num_types=2) + b1 = _batch([6, 8], bonds=[(0, 1)]) + b2 = _batch([6, 8], bonds=[(0, 1)]) + b2["atoms", "atom_type_id"] = torch.tensor([0, 1], dtype=torch.long) + f1 = probe.encode(b1) + f2 = probe.encode(b2) + assert torch.allclose(f1, f2) + + def test_overfit_fixture_accuracy(self): + torch.manual_seed(0) + enc = ChemEncoder(atom_dim=16, bond_dim=8) + probe = AtomTypeReadout(enc, num_types=2) + batch = _batch([6, 6, 1, 1], bonds=[(0, 1), (2, 3)]) + y = torch.tensor([0, 0, 1, 1], dtype=torch.long) + opt = torch.optim.Adam(probe.parameters(), lr=0.05) + for _ in range(80): + opt.zero_grad() + logits = probe(batch)["logits"] + loss = torch.nn.functional.cross_entropy(logits, y) + loss.backward() + opt.step() + pred = probe(batch)["pred_type_id"] + acc = float((pred == y).float().mean().item()) + assert acc >= 0.95 diff --git a/tests/test_molrep/test_condensation/__init__.py b/tests/test_molrep/test_condensation/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/test_molrep/test_condensation/test_condenser.py b/tests/test_molrep/test_condensation/test_condenser.py new file mode 100644 index 0000000..18e75f9 --- /dev/null +++ b/tests/test_molrep/test_condensation/test_condenser.py @@ -0,0 +1,161 @@ +"""Tests for Condenser greedy merge (learnable-classical-ff-06).""" + +from __future__ import annotations + +import ast +from pathlib import Path + +import torch + +from molrep.condensation import ( + Condenser, + InteractionClass, + MergeCriterion, + PhysicalErrorMetrics, + bond_default_criterion, +) + + +class TestCondenser: + def test_identical_params_merge_to_one_type(self): + condenser = Condenser() + params = { + "k": torch.tensor([300.0, 300.0, 300.0]), + "r0": torch.tensor([1.09, 1.09, 1.09]), + } + result = condenser.merge( + [params], + interaction=InteractionClass.BOND, + criterion=bond_default_criterion(), + ) + assert result.type_system.n_types == 1 + assert result.type_system.get(0).member_count == 3 + assert result.assignment.type_ids.tolist() == [0, 0, 0] + + def test_far_params_remain_distinct(self): + condenser = Condenser() + params = { + "k": torch.tensor([300.0, 300.0]), + "r0": torch.tensor([1.09, 1.5]), # Δr0 >> 0.01 Å + } + result = condenser.merge( + [params], + interaction=InteractionClass.BOND, + criterion=bond_default_criterion(), + ) + assert result.type_system.n_types == 2 + assert result.assignment.type_ids.tolist() == [0, 1] + + def test_deterministic_sort_key(self): + condenser = Condenser() + params = { + "k": torch.tensor([300.0, 500.0, 300.0]), + "r0": torch.tensor([1.09, 1.5, 1.09]), + } + r1 = condenser.greedy_merge( + [params], + interaction=InteractionClass.BOND, + criterion=bond_default_criterion(), + ) + r2 = condenser.greedy_merge( + [params], + interaction=InteractionClass.BOND, + criterion=bond_default_criterion(), + ) + assert r1.type_system.n_types == r2.type_system.n_types + assert r1.assignment.type_ids.tolist() == r2.assignment.type_ids.tolist() + assert r1.type_system.prototypes_table() == r2.type_system.prototypes_table() + + def test_multi_system_global_ids(self): + condenser = Condenser() + sys_a = {"k": torch.tensor([300.0]), "r0": torch.tensor([1.09])} + sys_b = { + "k": torch.tensor([302.0, 500.0]), # near 300/1.09 and far + "r0": torch.tensor([1.091, 1.5]), + } + result = condenser.merge( + [sys_a, sys_b], + interaction=InteractionClass.BOND, + criterion=bond_default_criterion(), + ) + # Near-duplicate across systems → one global type; far → second type + assert result.type_system.n_types == 2 + assert result.assignment.for_system(0).tolist() == [0] + assert result.assignment.for_system(1).tolist() == [0, 1] + assert result.assignment.type_id_at(0, 0) == 0 + assert result.assignment.type_id_at(1, 1) == 1 + + def test_physics_gate_rejects_param_match(self): + """physical_eval can block a param-budget merge.""" + condenser = Condenser() + + def always_hot(proto, cand, interaction): + return {"energy_error": 10.0} + + params = { + "k": torch.tensor([300.0, 301.0]), + "r0": torch.tensor([1.09, 1.09]), + } + result = condenser.merge( + [params], + interaction=InteractionClass.BOND, + criterion=bond_default_criterion(), + physical_eval=always_hot, + energy_tol=0.1, + ) + # First row spawns type 0; second would match params but physics rejects + assert result.type_system.n_types == 2 + assert result.metrics.rejected_by_physics == 1 + assert result.metrics.n_compared >= 1 + + def test_physics_eval_records_metrics_without_gate(self): + condenser = Condenser() + + def mild(proto, cand, interaction): + return 0.01 + + params = { + "k": torch.tensor([300.0, 300.0]), + "r0": torch.tensor([1.09, 1.09]), + } + result = condenser.merge( + [params], + interaction=InteractionClass.BOND, + criterion=bond_default_criterion(), + physical_eval=mild, + ) + assert result.type_system.n_types == 1 + assert isinstance(result.metrics, PhysicalErrorMetrics) + assert result.metrics.n_compared == 1 + assert result.metrics.max_abs_energy_error == 0.01 + + def test_zero_abs_budget_one_type_per_unique_row(self): + crit = MergeCriterion( + interaction=InteractionClass.BOND, + abs_tol={"r0": 0.0, "k": 0.0}, + required_keys=("k", "r0"), + ) + condenser = Condenser() + params = { + "k": torch.tensor([100.0, 100.0, 200.0]), + "r0": torch.tensor([1.0, 1.0, 1.0]), + } + result = condenser.merge([params], interaction=InteractionClass.BOND, criterion=crit) + assert result.type_system.n_types == 2 + + def test_no_smarts_emitters_in_package(self): + root = Path(__file__).resolve().parents[3] / "src/molrep/condensation" + for path in root.glob("*.py"): + tree = ast.parse(path.read_text()) + for node in ast.walk(tree): + if isinstance(node, ast.FunctionDef): + name = node.name.lower() + assert "smarts" not in name and "smirks" not in name + if isinstance(node, ast.ClassDef): + name = node.name.lower() + assert "smarts" not in name and "smirks" not in name + text = path.read_text().lower() + # No function that returns pattern text as its purpose + assert "def to_smarts" not in text + assert "def emit_smarts" not in text + assert "def to_smirks" not in text diff --git a/tests/test_molrep/test_condensation/test_criterion.py b/tests/test_molrep/test_condensation/test_criterion.py new file mode 100644 index 0000000..a9ea91a --- /dev/null +++ b/tests/test_molrep/test_condensation/test_criterion.py @@ -0,0 +1,67 @@ +"""Tests for MergeCriterion budgets (learnable-classical-ff-06).""" + +from __future__ import annotations + +import math + +import pytest +import torch + +from molrep.condensation import ( + InteractionClass, + MergeCriterion, + bond_default_criterion, + default_criterion, +) + + +class TestMergeCriterion: + def test_bond_accepts_within_budget(self): + crit = bond_default_criterion() + proto = {"k": torch.tensor(300.0), "r0": torch.tensor(1.09)} + cand = {"k": torch.tensor(310.0), "r0": torch.tensor(1.095)} # ~3% k, 0.005 Å + assert crit.accepts(proto, cand) is True + + def test_bond_rejects_r0_outside_budget(self): + crit = bond_default_criterion() + proto = {"k": torch.tensor(300.0), "r0": torch.tensor(1.09)} + cand = {"k": torch.tensor(300.0), "r0": torch.tensor(1.12)} # Δr0 = 0.03 > 0.01 + assert crit.accepts(proto, cand) is False + + def test_bond_rejects_k_outside_relative_budget(self): + crit = bond_default_criterion() + proto = {"k": torch.tensor(300.0), "r0": torch.tensor(1.09)} + # 20% k change with r0 exact — relative k budget is 5% + cand = {"k": torch.tensor(360.0), "r0": torch.tensor(1.09)} + assert crit.accepts(proto, cand) is False + + def test_bond_units_are_class_i(self): + crit = bond_default_criterion() + assert crit.interaction is InteractionClass.BOND + assert crit.abs_tol["r0"] == 0.01 # Å + assert crit.rel_tol["k"] == 0.05 + assert crit.required_keys == ("k", "r0") + + def test_zero_budget_forces_exact_match_on_abs(self): + crit = MergeCriterion( + interaction=InteractionClass.BOND, + abs_tol={"r0": 0.0, "k": 0.0}, + required_keys=("k", "r0"), + ) + proto = {"k": 100.0, "r0": 1.0} + assert crit.accepts(proto, {"k": 100.0, "r0": 1.0}) is True + assert crit.accepts(proto, {"k": 100.0, "r0": 1.0001}) is False + + def test_angle_default_theta0_in_radians(self): + crit = default_criterion(InteractionClass.ANGLE) + assert math.isclose(crit.abs_tol["theta0"], math.radians(1.0)) + proto = {"k": torch.tensor(50.0), "theta0": torch.tensor(1.9)} + near = {"k": torch.tensor(50.0), "theta0": torch.tensor(1.9 + math.radians(0.5))} + far = {"k": torch.tensor(50.0), "theta0": torch.tensor(1.9 + math.radians(2.0))} + assert crit.accepts(proto, near) is True + assert crit.accepts(proto, far) is False + + def test_missing_required_key_raises(self): + crit = bond_default_criterion() + with pytest.raises(KeyError): + crit.accepts({"k": 1.0}, {"k": 1.0, "r0": 1.0}) diff --git a/tests/test_molrep/test_condensation/test_type_system.py b/tests/test_molrep/test_condensation/test_type_system.py new file mode 100644 index 0000000..3568ac5 --- /dev/null +++ b/tests/test_molrep/test_condensation/test_type_system.py @@ -0,0 +1,78 @@ +"""Tests for TypeSystem / TypeRecord (learnable-classical-ff-06).""" + +from __future__ import annotations + +import torch + +from molrep.condensation import ( + UNMATCHED_TYPE_ID, + InteractionClass, + TypeRecord, + TypeSystem, + bond_default_criterion, +) + + +class TestTypeSystem: + def test_from_prototypes_dense_ids(self): + ts = TypeSystem.from_prototypes( + InteractionClass.BOND, + [{"k": 300.0, "r0": 1.09}, {"k": 500.0, "r0": 1.5}], + ) + assert ts.n_types == 2 + assert ts.get(0).prototype["r0"] == 1.09 + assert ts.get(1).member_count == 0 + + def test_assign_prototype_equal_returns_existing_id(self): + ts = TypeSystem.from_prototypes( + InteractionClass.BOND, + [{"k": 300.0, "r0": 1.09}], + criterion=bond_default_criterion(), + ) + tid = ts.assign({"k": torch.tensor(300.0), "r0": torch.tensor(1.09)}) + assert tid == 0 + + def test_assign_out_of_budget_soft_fails(self): + """Out-of-budget params return UNMATCHED_TYPE_ID; assign never spawns.""" + ts = TypeSystem.from_prototypes( + InteractionClass.BOND, + [{"k": 300.0, "r0": 1.09}], + criterion=bond_default_criterion(), + ) + tid = ts.assign({"k": torch.tensor(300.0), "r0": torch.tensor(1.5)}) + assert tid == UNMATCHED_TYPE_ID + assert ts.n_types == 1 + + def test_assign_many(self): + ts = TypeSystem.from_prototypes( + InteractionClass.BOND, + [{"k": 300.0, "r0": 1.09}, {"k": 500.0, "r0": 1.5}], + criterion=bond_default_criterion(), + ) + params = { + "k": torch.tensor([300.0, 500.0, 100.0]), + "r0": torch.tensor([1.09, 1.5, 2.0]), + } + ids = ts.assign_many(params) + assert ids.tolist() == [0, 1, UNMATCHED_TYPE_ID] + + def test_prototypes_table_list(self): + ts = TypeSystem.from_prototypes( + InteractionClass.LJ, + [{"epsilon": 0.1, "sigma": 3.0}], + ) + table = ts.prototypes_table() + assert isinstance(table, list) + assert table[0]["sigma"] == 3.0 + + def test_duplicate_type_id_rejected(self): + import pytest + + with pytest.raises(ValueError, match="duplicate"): + TypeSystem( + InteractionClass.BOND, + [ + TypeRecord(type_id=0, prototype={"k": 1.0, "r0": 1.0}), + TypeRecord(type_id=0, prototype={"k": 2.0, "r0": 1.0}), + ], + ) diff --git a/tests/test_molrep/test_embedding/__init__.py b/tests/test_molrep/test_embedding/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/test_molrep/test_embedding/conftest.py b/tests/test_molrep/test_embedding/conftest.py new file mode 100644 index 0000000..a82885a --- /dev/null +++ b/tests/test_molrep/test_embedding/conftest.py @@ -0,0 +1,25 @@ +"""Shared fixtures for the :mod:`molrep.embedding` unit suite.""" + +from __future__ import annotations + +from collections.abc import Iterator + +import pytest +import torch + +from molix import config + + +@pytest.fixture +def fp64() -> Iterator[None]: + """Run the case under the global fp64 precision, restoring the previous one. + + Every layer in :mod:`molrep.embedding` bakes ``config["ftype"]`` into its + parameters and buffers at construction time (the contract documented in + :mod:`molix.config`), so the precision has to be switched *before* the + module is built and handed back afterwards. + """ + previous = config["ftype"] + config.set_precision("fp64") + yield + config.set_precision("fp64" if previous == torch.float64 else "fp32") diff --git a/tests/test_molrep/test_embedding/test_covalent.py b/tests/test_molrep/test_embedding/test_covalent.py new file mode 100644 index 0000000..c58f69a --- /dev/null +++ b/tests/test_molrep/test_embedding/test_covalent.py @@ -0,0 +1,45 @@ +"""Tests for molrep.embedding.covalent module.""" + +import pytest +import torch + +from molrep.embedding.covalent import covalent_radii + + +class TestCovalentRadii: + """Test the Z-indexed covalent-radius table.""" + + def test_length_is_max_z_plus_one(self): + """Table is indexable by Z directly, so it needs a row for Z=0.""" + assert covalent_radii(max_z=10).shape == (11,) + + def test_known_values(self): + """Spot-check against Cordero et al. (2008), the table MACE also uses.""" + table = covalent_radii(max_z=8) + assert table[1] == pytest.approx(0.31, abs=1e-6) # H + assert table[6] == pytest.approx(0.76, abs=1e-6) # C + assert table[8] == pytest.approx(0.66, abs=1e-6) # O + + def test_dummy_slot_zero(self): + """Index 0 is a dummy-atom placeholder, matching the reference table.""" + assert covalent_radii(max_z=4)[0] == pytest.approx(0.2, abs=1e-9) + + def test_all_positive(self): + """Every real element has a positive radius.""" + assert bool((covalent_radii(max_z=96)[1:] > 0).all()) + + def test_rejects_empty_table(self): + """A table with no elements is a caller error, not an empty tensor.""" + with pytest.raises(ValueError, match="max_z"): + covalent_radii(max_z=0) + + def test_dtype_follows_config(self): + """The table is built in the project dtype so it can be a module buffer.""" + from molix import config + + assert covalent_radii(max_z=4).dtype == config.ftype + + def test_is_a_plain_tensor(self): + """Callers register it as a buffer; it must not carry grad.""" + assert not covalent_radii(max_z=4).requires_grad + assert isinstance(covalent_radii(max_z=4), torch.Tensor) diff --git a/tests/test_molrep/test_embedding/test_cutoff.py b/tests/test_molrep/test_embedding/test_cutoff.py index 8284682..3fd233e 100644 --- a/tests/test_molrep/test_embedding/test_cutoff.py +++ b/tests/test_molrep/test_embedding/test_cutoff.py @@ -1,5 +1,7 @@ """Tests for molrep.embedding.cutoff module.""" +import math + import pytest import torch @@ -10,6 +12,12 @@ PolynomialCutoffSpec, ) +#: A radius with no exact fp32 representation — ``float32(5.1)`` is +#: ``5.099999904632568``, off by ~9.5e-8. Any buffer that silently lands in +#: fp32 shifts the whole envelope by that much, far above the 1e-12 position +#: tolerance an fp64 run is asking for. +R_CUT_NOT_FP32_EXACT = 5.1 + class TestCosineCutoffSpec: """Test CosineCutoffSpec configuration.""" @@ -133,8 +141,38 @@ def test_dtype_consistency(self): # Float64 dist_f64 = torch.tensor([1.0, 2.0], dtype=torch.float64) out_f64 = cutoff(dist_f64) - # Note: Cutoff may cast to float internally - assert out_f64.dtype in [torch.float32, torch.float64] + assert out_f64.dtype == torch.float64 + + def test_r_cut_buffer_honours_the_fp64_precision(self, fp64): + """The ``r_cut`` buffer is fp64 when the module is built under fp64. + + ``config["ftype"]`` is the single source of truth for the working + precision; an fp32 ``r_cut`` inside an otherwise-fp64 model silently + demotes every distance ratio it participates in. + """ + cutoff = CosineCutoff(r_cut=5.0) + + assert cutoff.r_cut.dtype == torch.float64 + + def test_r_cut_keeps_full_precision_for_a_non_fp32_representable_radius(self, fp64): + """An fp32 ``r_cut`` truncates the requested radius by ~1e-7 Å.""" + cutoff = CosineCutoff(r_cut=R_CUT_NOT_FP32_EXACT) + + assert float(cutoff.r_cut) == pytest.approx(R_CUT_NOT_FP32_EXACT, abs=1e-12) + + def test_forward_matches_the_fp64_envelope(self, fp64): + """c(r) matches the double-precision formula at the fp64 tolerance. + + The tolerance sits well above fp64 round-off (~1e-16) and well below + the ~1e-8 error an fp32-truncated ``r_cut`` introduces. + """ + cutoff = CosineCutoff(r_cut=R_CUT_NOT_FP32_EXACT) + r = torch.tensor([0.5, 2.55, 4.0], dtype=torch.float64) + + got = cutoff(r) + + expected = 0.5 * (torch.cos(math.pi * r / R_CUT_NOT_FP32_EXACT) + 1.0) + torch.testing.assert_close(got, expected, rtol=0.0, atol=1e-10) def test_broadcasting(self): """Test broadcasting behavior.""" @@ -195,3 +233,50 @@ def test_invalid_exponent(self): PolynomialCutoffSpec(r_cut=5.0, exponent=0) with pytest.raises(ValueError): PolynomialCutoffSpec(r_cut=5.0, exponent=-1) + + def test_r_cut_buffer_honours_the_fp64_precision(self, fp64): + """The ``r_cut`` buffer is fp64 when the module is built under fp64. + + Same construction-time contract as :class:`CosineCutoff` — this + envelope is the NequIP/Allegro default, so an fp32 ``r_cut`` here + demotes the radial ratio on the main representation path. + """ + cutoff = PolynomialCutoff(r_cut=5.0, exponent=6) + + assert cutoff.r_cut.dtype == torch.float64 + + def test_r_cut_keeps_full_precision_for_a_non_fp32_representable_radius(self, fp64): + """An fp32 ``r_cut`` truncates the requested radius by ~1e-7 Å.""" + cutoff = PolynomialCutoff(r_cut=R_CUT_NOT_FP32_EXACT, exponent=6) + + assert float(cutoff.r_cut) == pytest.approx(R_CUT_NOT_FP32_EXACT, abs=1e-12) + + +class TestPolynomialCutoffEnvelope: + """Test the static envelope used with a per-edge cutoff radius.""" + + def test_matches_forward_for_the_module_radius(self): + """The module's forward is the static envelope at its own r_cut.""" + cutoff = PolynomialCutoff(r_cut=5.0, exponent=5) + r = torch.linspace(0.0, 6.0, 25) + assert torch.allclose(cutoff(r), PolynomialCutoff.envelope(r, 5.0, 5)) + + def test_accepts_a_per_element_radius(self): + """ZBL needs one cutoff per edge, from the pair's covalent radii.""" + r = torch.tensor([1.0, 1.0, 1.0]) + per_edge = torch.tensor([0.8, 2.0, 5.0]) + out = PolynomialCutoff.envelope(r, per_edge, 5) + assert float(out[0]) == 0.0 # beyond its own cutoff + assert float(out[1]) > 0.0 + assert float(out[2]) > float(out[1]) # further inside a wider cutoff + + def test_is_one_at_zero_and_zero_beyond(self): + """Envelope endpoints: 1 at contact, exactly 0 past the radius.""" + assert float(PolynomialCutoff.envelope(torch.zeros(1), 3.0, 5)) == pytest.approx(1.0) + assert float(PolynomialCutoff.envelope(torch.tensor([3.5]), 3.0, 5)) == 0.0 + + def test_derivative_vanishes_at_the_cutoff(self): + """Smooth shutoff is what keeps autograd forces continuous.""" + r = torch.tensor([2.9999], requires_grad=True) + PolynomialCutoff.envelope(r, 3.0, 5).backward() + assert abs(float(r.grad)) < 1e-6 diff --git a/tests/test_molrep/test_embedding/test_mace.py b/tests/test_molrep/test_embedding/test_mace.py new file mode 100644 index 0000000..dd1fbac --- /dev/null +++ b/tests/test_molrep/test_embedding/test_mace.py @@ -0,0 +1,315 @@ +"""Tests for molrep.embedding.mace module.""" + +import math + +import pytest +import torch + +from molrep.embedding.angular import SphericalHarmonics +from molrep.embedding.cutoff import CosineCutoff +from molrep.embedding.mace import EmbeddingBlock, EmbeddingSpec +from molrep.embedding.node import DiscreteEmbeddingSpec, JointEmbedding +from molrep.embedding.radial import BesselRBF +from molrep.utils.equivariance import ( + random_rotation_matrix, + rotate_vectors, + rotation_matrix_z, +) + + +class TestEmbeddingSpec: + """Test EmbeddingSpec configuration.""" + + def test_requires_at_least_one_node_attr_spec(self): + """``node_attr_specs`` has ``min_length=1`` — an empty list is rejected.""" + with pytest.raises(ValueError): + EmbeddingSpec(node_attr_specs=[], num_features=16, r_max=5.0) + + +class TestEmbeddingBlock: + """Test EmbeddingBlock initialization and forward pass.""" + + @pytest.fixture + def embedding_config(self): + """Common configuration for embedding tests.""" + return { + "num_species": 5, + "num_features": 16, + "r_max": 5.0, + "num_bessel": 8, + "l_max": 2, + } + + @pytest.fixture + def node_attr_specs(self, embedding_config): + """Node attribute specifications.""" + return [ + DiscreteEmbeddingSpec( + input_key="Z", + num_classes=embedding_config["num_species"], + emb_dim=embedding_config["num_features"], + ) + ] + + @pytest.fixture + def embedding_block(self, node_attr_specs, embedding_config): + """Create an EmbeddingBlock instance.""" + return EmbeddingBlock( + node_attr_specs=node_attr_specs, + num_features=embedding_config["num_features"], + r_max=embedding_config["r_max"], + num_bessel=embedding_config["num_bessel"], + l_max=embedding_config["l_max"], + ) + + def test_initialization(self, embedding_block, embedding_config): + """Test that EmbeddingBlock initializes all components correctly.""" + # Check node_embedding + assert hasattr(embedding_block, "node_embedding") + assert isinstance(embedding_block.node_embedding, JointEmbedding) + + # Check radial_embedding + assert hasattr(embedding_block, "radial_embedding") + assert isinstance(embedding_block.radial_embedding, BesselRBF) + assert embedding_block.radial_embedding.config.r_cut == embedding_config["r_max"] + assert embedding_block.radial_embedding.config.num_radial == embedding_config["num_bessel"] + + # Check spherical_harmonics + assert hasattr(embedding_block, "spherical_harmonics") + assert isinstance(embedding_block.spherical_harmonics, SphericalHarmonics) + assert embedding_block.spherical_harmonics.l_max == embedding_config["l_max"] + + # Check cutoff_fn + assert hasattr(embedding_block, "cutoff_fn") + assert isinstance(embedding_block.cutoff_fn, CosineCutoff) + assert embedding_block.cutoff_fn.config.r_cut == embedding_config["r_max"] + + def test_config_storage(self, embedding_block, embedding_config): + """Test that configuration is properly stored.""" + assert hasattr(embedding_block, "config") + config = embedding_block.config + assert config.num_features == embedding_config["num_features"] + assert config.r_max == embedding_config["r_max"] + assert config.num_bessel == embedding_config["num_bessel"] + assert config.l_max == embedding_config["l_max"] + + def test_forward_output_shapes(self, embedding_block, embedding_config): + """Test forward pass returns correct output shapes.""" + n_atoms = 4 + n_edges = 6 + + # Create input data + z = torch.randint(0, embedding_config["num_species"], (n_atoms,)) + edge_dist = torch.rand(n_edges) * embedding_config["r_max"] + edge_diff = torch.randn(n_edges, 3) + + # Normalize edge_diff to match edge_dist + edge_diff = ( + edge_diff / torch.norm(edge_diff, dim=-1, keepdim=True) * edge_dist.unsqueeze(-1) + ) + + # Forward pass + node_feats, edge_attrs, edge_feats = embedding_block( + Z=z, + edge_dist=edge_dist, + edge_diff=edge_diff, + ) + + # Check shapes + assert node_feats.shape == (n_atoms, embedding_config["num_features"]) + + # Spherical harmonics dimension: (2*l_max + 1)^2 for l_max=2 is 9 + expected_sh_dim = (embedding_config["l_max"] + 1) ** 2 + assert edge_attrs.shape == (n_edges, expected_sh_dim) + + assert edge_feats.shape == (n_edges, embedding_config["num_bessel"]) + + def test_node_embedding_component(self, embedding_block, embedding_config): + """Test node_embedding component works independently.""" + n_atoms = 5 + z = torch.randint(0, embedding_config["num_species"], (n_atoms,)) + + # Call node_embedding directly + node_feats = embedding_block.node_embedding(Z=z) + + assert node_feats.shape == (n_atoms, embedding_config["num_features"]) + assert node_feats.dtype == torch.float32 + + def test_radial_embedding_component(self, embedding_block, embedding_config): + """Test radial_embedding component works independently.""" + n_edges = 10 + edge_dist = torch.rand(n_edges) * embedding_config["r_max"] + + # Call radial_embedding directly + edge_radial = embedding_block.radial_embedding(edge_dist) + + assert edge_radial.shape == (n_edges, embedding_config["num_bessel"]) + assert edge_radial.dtype == torch.float32 + + def test_spherical_harmonics_component(self, embedding_block, embedding_config): + """Test spherical_harmonics component works independently.""" + n_edges = 8 + # Create normalized direction vectors + edge_dir = torch.randn(n_edges, 3) + edge_dir = edge_dir / torch.norm(edge_dir, dim=-1, keepdim=True) + + # Call spherical_harmonics directly + edge_attrs = embedding_block.spherical_harmonics(edge_dir) + + expected_sh_dim = (embedding_config["l_max"] + 1) ** 2 + assert edge_attrs.shape == (n_edges, expected_sh_dim) + assert edge_attrs.dtype == torch.float32 + + def test_cutoff_component(self, embedding_block, embedding_config): + """Test cutoff_fn component works independently.""" + n_edges = 12 + edge_dist = torch.rand(n_edges) * embedding_config["r_max"] + + # Call cutoff_fn directly + cutoff_values = embedding_block.cutoff_fn(edge_dist) + + assert cutoff_values.shape == (n_edges,) + assert cutoff_values.dtype == torch.float32 + # Cutoff should be in [0, 1] + assert (cutoff_values >= 0.0).all() + assert (cutoff_values <= 1.0).all() + + def test_edge_feats_includes_cutoff(self, embedding_block, embedding_config): + """Test that edge_feats properly applies cutoff to radial basis.""" + n_edges = 6 + edge_dist = torch.rand(n_edges) * embedding_config["r_max"] + edge_diff = torch.randn(n_edges, 3) + edge_diff = ( + edge_diff / torch.norm(edge_diff, dim=-1, keepdim=True) * edge_dist.unsqueeze(-1) + ) + + z = torch.randint(0, embedding_config["num_species"], (3,)) + + # Get outputs + _, _, edge_feats = embedding_block( + Z=z, + edge_dist=edge_dist, + edge_diff=edge_diff, + ) + + # Compute expected edge_feats manually + edge_radial = embedding_block.radial_embedding(edge_dist) + cutoff_values = embedding_block.cutoff_fn(edge_dist) + expected_edge_feats = edge_radial * cutoff_values.unsqueeze(-1) + + # Check they match + assert torch.allclose(edge_feats, expected_edge_feats, atol=1e-6) + + def test_cutoff_at_boundary(self, embedding_block, embedding_config): + """Test cutoff behavior at r_max boundary.""" + # Distance at cutoff should give near-zero cutoff value + edge_dist = torch.tensor([embedding_config["r_max"]]) + cutoff_value = embedding_block.cutoff_fn(edge_dist) + + # Cosine cutoff should be near 0 at r_max + assert cutoff_value.item() < 0.01 + + # Distance at 0 should give cutoff value of 1 + bond_dist_zero = torch.tensor([0.0]) + cutoff_value_zero = embedding_block.cutoff_fn(bond_dist_zero) + assert abs(cutoff_value_zero.item() - 1.0) < 0.01 + + +class TestEmbeddingBlockEquivariance: + """Test equivariance properties of EmbeddingBlock. + + Migrated from ``tests/test_molzoo/test_mace.py`` by + mace-subpackage-restructure-02-core (the block now lives in molrep, and + the molzoo module/package name clash forced the file's removal). + """ + + @pytest.fixture + def embedding_block(self): + """Create an EmbeddingBlock for equivariance testing.""" + node_attr_specs = [ + DiscreteEmbeddingSpec( + input_key="Z", + num_classes=5, + emb_dim=16, + ) + ] + return EmbeddingBlock( + node_attr_specs=node_attr_specs, + num_features=16, + r_max=5.0, + num_bessel=8, + l_max=2, + ) + + def test_spherical_harmonics_equivariance(self, embedding_block): + """Test that spherical harmonics are equivariant under rotation. + + Rotating the bond vectors should rotate the spherical harmonics accordingly. + """ + n_atoms = 4 + n_edges = 6 + + # Create input data + z = torch.randint(0, 5, (n_atoms,)) + edge_diff = torch.randn(n_edges, 3) + edge_dist = torch.norm(edge_diff, dim=-1) + + # Forward pass + _, edge_attrs1, _ = embedding_block( + Z=z, + edge_dist=edge_dist, + edge_diff=edge_diff, + ) + + # Rotate bond vectors + angle = math.pi / 2 + rot_matrix = rotation_matrix_z(angle, dtype=edge_diff.dtype) + bond_diff_rot = rotate_vectors(edge_diff, rot_matrix) + + # Forward pass on rotated + _, edge_attrs2, _ = embedding_block( + Z=z, + edge_dist=edge_dist, + edge_diff=bond_diff_rot, + ) + + # l=0 component should be invariant + assert torch.allclose(edge_attrs1[:, 0], edge_attrs2[:, 0], atol=1e-5) + + # Overall norm should be preserved + norm1 = edge_attrs1.norm(dim=-1) + norm2 = edge_attrs2.norm(dim=-1) + assert torch.allclose(norm1, norm2, rtol=1e-4, atol=1e-4) + + def test_radial_features_invariance(self, embedding_block): + """Test that radial features are rotation invariant. + + Rotating bond vectors should not change radial features (distances). + """ + n_atoms = 4 + n_edges = 6 + + z = torch.randint(0, 5, (n_atoms,)) + edge_diff = torch.randn(n_edges, 3) + edge_dist = torch.norm(edge_diff, dim=-1) + + # Forward pass + _, _, edge_feats1 = embedding_block( + Z=z, + edge_dist=edge_dist, + edge_diff=edge_diff, + ) + + # Rotate bond vectors + rot_matrix = random_rotation_matrix(dtype=edge_diff.dtype) + bond_diff_rot = rotate_vectors(edge_diff, rot_matrix) + + # Forward pass on rotated + _, _, edge_feats2 = embedding_block( + Z=z, + edge_dist=edge_dist, + edge_diff=bond_diff_rot, + ) + + # Radial features should be identical (rotation invariant) + assert torch.allclose(edge_feats1, edge_feats2, rtol=1e-5, atol=1e-5) diff --git a/tests/test_molrep/test_embedding/test_mlp.py b/tests/test_molrep/test_embedding/test_mlp.py new file mode 100644 index 0000000..0ea8193 --- /dev/null +++ b/tests/test_molrep/test_embedding/test_mlp.py @@ -0,0 +1,85 @@ +"""Tests for molrep.embedding.mlp module.""" + +import math + +import pytest +import torch +import torch.nn.functional as F + +from molrep.embedding.mlp import MomentNormalizedMLP, normalize2mom + + +class TestNormalize2Mom: + """Test the activation second-moment normalisation constant.""" + + def test_scaled_activation_has_unit_second_moment(self): + """``E[(c·act(z))²] = 1`` for ``z ~ N(0, 1)`` — the defining property. + + Checked on an independent draw, so the tolerance covers the Monte-Carlo + error of both the constant (1e6 samples) and this estimate. + """ + constant = normalize2mom(F.silu) + z = torch.randn(2_000_000, dtype=torch.float64) + assert float((constant * F.silu(z)).pow(2).mean()) == pytest.approx(1.0, abs=1e-2) + + def test_is_deterministic(self): + """A fixed draw means two processes agree — weights must transfer.""" + assert normalize2mom(F.silu) == normalize2mom(F.silu) + + def test_identity_activation_is_unscaled(self): + """``E[z²] = 1`` already, so the identity needs no rescaling.""" + assert normalize2mom(lambda x: x) == pytest.approx(1.0, abs=1e-2) + + def test_larger_activation_gets_smaller_constant(self): + """A activation with bigger output needs a smaller constant.""" + assert normalize2mom(lambda x: 2.0 * x) < normalize2mom(lambda x: x) + + +class TestMomentNormalizedMLP: + """Test the e3nn-compatible scalar MLP.""" + + def test_output_shape(self): + """Maps the last dimension from ``channels[0]`` to ``channels[-1]``.""" + mlp = MomentNormalizedMLP([10, 64, 512]) + assert mlp(torch.randn(7, 10)).shape == (7, 512) + + def test_layer_names_mirror_the_reference(self): + """Official weights transfer by direct copy, so names must match.""" + names = {n for n, _ in MomentNormalizedMLP([10, 64, 64, 512]).named_parameters()} + assert names == {f"layer{i}.weight" for i in range(3)} + + def test_has_no_biases(self): + """e3nn's FullyConnectedNet is bias-free; a bias would break transfer.""" + assert not any("bias" in n for n, _ in MomentNormalizedMLP([4, 8, 2]).named_parameters()) + + def test_maps_zero_to_zero(self): + """Bias-free and SiLU(0)=0, so a dead edge contributes nothing.""" + mlp = MomentNormalizedMLP([6, 12, 3]) + assert float(mlp(torch.zeros(2, 6)).abs().max()) == 0.0 + + def test_single_layer_is_a_bare_scaled_linear(self): + """With two channel entries there is no hidden layer and no activation.""" + mlp = MomentNormalizedMLP([4, 1]).double() + x = torch.randn(3, 4, dtype=torch.float64) + want = x @ (mlp.layer0.weight * (1.0 / math.sqrt(4))) + assert torch.allclose(mlp(x), want, atol=1e-12) + + def test_weight_scaling_is_one_over_sqrt_fan_in(self): + """The ``1/√fan_in`` factor is applied at forward, not baked into W.""" + mlp = MomentNormalizedMLP([9, 1]).double() + with torch.no_grad(): + mlp.layer0.weight.fill_(1.0) + x = torch.ones(1, 9, dtype=torch.float64) + assert float(mlp(x)) == pytest.approx(9.0 / 3.0, abs=1e-12) + + def test_rejects_too_few_channels(self): + """A channel list needs at least an input and an output width.""" + with pytest.raises(ValueError, match="channels"): + MomentNormalizedMLP([8]) + + def test_is_differentiable(self): + """It generates tensor-product weights, so gradients must flow.""" + mlp = MomentNormalizedMLP([4, 8, 2]) + x = torch.randn(3, 4, requires_grad=True) + mlp(x).sum().backward() + assert x.grad is not None and float(x.grad.abs().sum()) > 0.0 diff --git a/tests/test_molrep/test_embedding/test_node.py b/tests/test_molrep/test_embedding/test_node.py index 1e4b86a..c39f8eb 100644 --- a/tests/test_molrep/test_embedding/test_node.py +++ b/tests/test_molrep/test_embedding/test_node.py @@ -8,6 +8,8 @@ DiscreteEmbeddingSpec, JointEmbedding, JointEmbeddingSpec, + JointFeatureEmbedding, + JointFeatureSpec, ) @@ -119,3 +121,117 @@ def test_gradient_flow(self): assert attr.grad is not None assert not torch.isnan(attr.grad).any() + + @pytest.mark.parametrize( + "spec_kinds", + [("discrete",), ("continuous",), ("discrete", "continuous")], + ids=["discrete", "continuous", "both"], + ) + def test_all_parameters_honour_the_fp64_precision(self, fp64, spec_kinds): + """Every parameter is fp64 when the module is built under fp64. + + ``config["ftype"]`` is the single source of truth for the working + precision; a layer that ignores it leaves the module mixed-precision + and its first forward dies on a dtype-mismatched matmul. + """ + available = { + "discrete": DiscreteEmbeddingSpec(input_key="Z", num_classes=10, emb_dim=16), + "continuous": ContinuousEmbeddingSpec(input_key="pos", in_dim=3, emb_dim=16), + } + joint = JointEmbedding( + embedding_specs=[available[kind] for kind in spec_kinds], + out_dim=32, + ) + + assert {p.dtype for p in joint.parameters()} == {torch.float64} + + def test_forward_runs_under_the_fp64_precision(self, fp64): + """A module built at fp64 consumes fp64 inputs and emits fp64.""" + specs = [ + DiscreteEmbeddingSpec(input_key="Z", num_classes=10, emb_dim=16), + ContinuousEmbeddingSpec(input_key="pos", in_dim=3, emb_dim=32), + ] + joint = JointEmbedding(embedding_specs=specs, out_dim=64) + + n_atoms = 5 + z = torch.zeros(n_atoms, dtype=torch.long) + pos = torch.zeros(n_atoms, 3, dtype=torch.float64) + + output = joint(Z=z, pos=pos) + assert output.shape == (n_atoms, 64) + assert output.dtype == torch.float64 + + +class TestJointFeatureSpec: + """Test JointFeatureSpec configuration.""" + + def test_valid_categorical_spec(self): + """A categorical feature carries its vocabulary size.""" + spec = JointFeatureSpec( + name="total_charge", kind="categorical", emb_dim=4, num_classes=5, offset=2 + ) + assert spec.name == "total_charge" + assert spec.kind == "categorical" + assert spec.num_classes == 5 + assert spec.offset == 2 + assert spec.per == "graph" + + def test_valid_continuous_spec(self): + """A continuous feature carries its input dimension.""" + spec = JointFeatureSpec(name="total_spin", kind="continuous", emb_dim=4, in_dim=1) + assert spec.kind == "continuous" + assert spec.in_dim == 1 + assert spec.use_bias is True + + def test_invalid_emb_dim(self): + """Test validation for emb_dim.""" + with pytest.raises(ValueError): + JointFeatureSpec(name="total_spin", kind="continuous", emb_dim=0, in_dim=1) + + +class TestJointFeatureEmbedding: + """Test JointFeatureEmbedding module.""" + + #: One spec per ``kind``; ``per="graph"`` matches the OMOL charge/spin use. + SPECS = { + "categorical": JointFeatureSpec( + name="total_charge", kind="categorical", emb_dim=4, num_classes=5, offset=2 + ), + "continuous": JointFeatureSpec(name="total_spin", kind="continuous", emb_dim=4, in_dim=1), + } + + @pytest.mark.parametrize( + "spec_kinds", + [("categorical",), ("continuous",), ("categorical", "continuous")], + ids=["categorical", "continuous", "both"], + ) + def test_all_parameters_honour_the_fp64_precision(self, fp64, spec_kinds): + """Every parameter is fp64 when the module is built under fp64.""" + embedding = JointFeatureEmbedding( + feature_specs=[self.SPECS[kind] for kind in spec_kinds], + out_dim=8, + ) + + assert {p.dtype for p in embedding.parameters()} == {torch.float64} + + def test_forward_runs_under_the_fp64_precision(self, fp64): + """A module built at fp64 fuses fp64 per-graph features into per-atom ones.""" + embedding = JointFeatureEmbedding( + feature_specs=[self.SPECS["categorical"], self.SPECS["continuous"]], + out_dim=8, + ) + + n_atoms = 3 + output = embedding( + batch=torch.zeros(n_atoms, dtype=torch.long), + total_charge=torch.zeros(1, dtype=torch.long), + total_spin=torch.ones(1, dtype=torch.float64), + ) + + assert output.shape == (n_atoms, 8) + assert output.dtype == torch.float64 + + def test_empty_specs_error(self): + """Test error when no feature specs are provided.""" + with pytest.raises(ValueError, match="feature_specs must be non-empty."): + JointFeatureEmbedding(feature_specs=[], out_dim=8) diff --git a/tests/test_molrep/test_embedding/test_radial.py b/tests/test_molrep/test_embedding/test_radial.py index 35968e8..a4cbad7 100644 --- a/tests/test_molrep/test_embedding/test_radial.py +++ b/tests/test_molrep/test_embedding/test_radial.py @@ -3,7 +3,7 @@ import pytest import torch -from molrep.embedding.radial import BesselRBF, BesselRBFSpec +from molrep.embedding.radial import AgnesiTransform, BesselRBF, BesselRBFSpec class TestBesselRBFSpec: @@ -126,3 +126,55 @@ def test_dtype_consistency(self): out_f64 = rbf(dist_f64) # Note: BesselRBF casts to float internally assert out_f64.dtype in [torch.float32, torch.float64] + + +class TestAgnesiTransform: + """Test the Agnesi element-pair distance transform.""" + + def test_maps_into_the_unit_interval(self): + """The transform compresses r onto (0, 1] so the basis stays bounded.""" + transform = AgnesiTransform().double() + r = torch.linspace(0.1, 12.0, 50, dtype=torch.float64) + z = torch.full((50,), 8, dtype=torch.long) + u = transform(r, z, z) + assert bool((u > 0).all()) and bool((u <= 1.0).all()) + + def test_is_monotonically_decreasing(self): + """Longer distances map to smaller transformed coordinates.""" + transform = AgnesiTransform().double() + r = torch.linspace(0.3, 8.0, 40, dtype=torch.float64) + z = torch.full((40,), 6, dtype=torch.long) + u = transform(r, z, z) + assert bool((u[1:] - u[:-1] < 0).all()) + + def test_is_symmetric_in_the_pair(self): + """r_0 uses the mean covalent radius, so swapping the pair is a no-op.""" + transform = AgnesiTransform().double() + r = torch.tensor([1.0, 2.0], dtype=torch.float64) + forward = transform(r, torch.tensor([1, 8]), torch.tensor([8, 1])) + reverse = transform(r, torch.tensor([8, 1]), torch.tensor([1, 8])) + assert torch.allclose(forward, reverse) + + def test_element_pair_sets_the_length_scale(self): + """A larger pair radius means the same r is compressed less.""" + transform = AgnesiTransform().double() + r = torch.tensor([1.5, 1.5], dtype=torch.float64) + small = transform(r[:1], torch.tensor([1]), torch.tensor([1])) # H-H + large = transform(r[:1], torch.tensor([55]), torch.tensor([55])) # Cs-Cs + assert float(large) > float(small) + + def test_is_differentiable(self): + """It sits on the force path, so it must carry gradient.""" + transform = AgnesiTransform().double() + r = torch.tensor([1.0, 2.5], dtype=torch.float64, requires_grad=True) + transform(r, torch.tensor([8, 1]), torch.tensor([1, 8])).sum().backward() + assert r.grad is not None and bool((r.grad < 0).all()) + + def test_parameters_are_frozen_by_default(self): + """a / q / p are buffers unless the caller opts into training them.""" + assert not any(p.requires_grad for p in AgnesiTransform().parameters()) + + def test_trainable_promotes_shape_parameters(self): + """``trainable=True`` exposes a / q / p for fine-tuning.""" + names = {n for n, _ in AgnesiTransform(trainable=True).named_parameters()} + assert {"a", "q", "p"} <= names diff --git a/tests/test_molrep/test_embedding/test_support.py b/tests/test_molrep/test_embedding/test_support.py new file mode 100644 index 0000000..c8feb90 --- /dev/null +++ b/tests/test_molrep/test_embedding/test_support.py @@ -0,0 +1,138 @@ +"""Tests for :class:`molrep.embedding.support.ChemicalSupportIndex`. + +Spec: learnable-classical-ff-09-provenance. +""" + +from __future__ import annotations + +import ast +from pathlib import Path + +import pytest +import torch + +from molrep.embedding.support import ChemicalSupportIndex + + +class TestChemicalSupportIndex: + """kNN L2 bank membership and coverage.""" + + def test_contains_exact_bank_members(self): + """Bank points lie inside the support radius (distance 0).""" + bank = torch.tensor( + [ + [0.0, 0.0], + [1.0, 0.0], + [0.0, 1.0], + ], + dtype=torch.float64, + ) + index = ChemicalSupportIndex(bank, radius=0.1, k=1) + query = bank.clone() + assert torch.equal(index.contains(query), torch.tensor([True, True, True])) + + def test_contains_outside_radius(self): + """Points farther than radius are out of support.""" + bank = torch.tensor([[0.0, 0.0]], dtype=torch.float64) + index = ChemicalSupportIndex(bank, radius=0.5, k=1) + query = torch.tensor( + [ + [0.0, 0.0], # dist 0 + [0.4, 0.0], # dist 0.4 <= 0.5 + [1.0, 0.0], # dist 1.0 > 0.5 + ], + dtype=torch.float64, + ) + assert torch.equal( + index.contains(query), + torch.tensor([True, True, False]), + ) + + def test_knn_distances_and_indices(self): + """Nearest neighbour is the expected bank row under L2.""" + bank = torch.tensor( + [ + [0.0, 0.0], + [10.0, 0.0], + [0.0, 10.0], + ], + dtype=torch.float64, + ) + index = ChemicalSupportIndex(bank, radius=1.0, k=1) + query = torch.tensor([[0.1, 0.0]], dtype=torch.float64) + dist, nn_idx = index.knn(query) + assert dist.shape == (1, 1) + assert nn_idx.shape == (1, 1) + assert int(nn_idx[0, 0]) == 0 + assert float(dist[0, 0]) == pytest.approx(0.1, abs=1e-8) + + def test_min_distance(self): + bank = torch.tensor([[0.0, 0.0], [3.0, 0.0]], dtype=torch.float64) + index = ChemicalSupportIndex(bank, radius=1.0, k=1) + query = torch.tensor([[1.0, 0.0], [3.0, 4.0]], dtype=torch.float64) + d = index.min_distance(query) + assert d.shape == (2,) + assert float(d[0]) == pytest.approx(1.0, abs=1e-8) + assert float(d[1]) == pytest.approx(4.0, abs=1e-8) + + def test_coverage_fraction(self): + """coverage = (# queries in support) / (# queries).""" + bank = torch.tensor([[0.0], [1.0], [2.0]], dtype=torch.float64) + index = ChemicalSupportIndex(bank, radius=0.0, k=1) + # type-id style 1-D bank: exact match only + query = torch.tensor([[0.0], [1.0], [99.0], [2.0]], dtype=torch.float64) + assert index.coverage_fraction(query) == pytest.approx(0.75) + + def test_from_type_ids_exact_membership(self): + """Discrete type_id sets map to a 1-D L2 bank with radius 0.""" + index = ChemicalSupportIndex.from_type_ids({0, 2, 5}, radius=0.0) + ids = torch.tensor([0, 1, 2, 5, 7], dtype=torch.long) + assert torch.equal( + index.contains_type_ids(ids), + torch.tensor([True, False, True, True, False]), + ) + assert index.coverage_fraction_type_ids(ids) == pytest.approx(3 / 5) + + def test_empty_query_coverage_is_one(self): + """Empty prediction set is fully covered by convention.""" + bank = torch.tensor([[0.0, 0.0]], dtype=torch.float64) + index = ChemicalSupportIndex(bank, radius=1.0, k=1) + empty = torch.empty(0, 2, dtype=torch.float64) + assert index.coverage_fraction(empty) == 1.0 + + def test_k_greater_than_one(self): + bank = torch.tensor( + [ + [0.0, 0.0], + [1.0, 0.0], + [0.0, 1.0], + ], + dtype=torch.float64, + ) + index = ChemicalSupportIndex(bank, radius=10.0, k=2) + query = torch.tensor([[0.0, 0.0]], dtype=torch.float64) + dist, nn_idx = index.knn(query) + assert dist.shape == (1, 2) + assert set(nn_idx[0].tolist()) == {0, 1} or set(nn_idx[0].tolist()) == {0, 2} + assert float(dist[0, 0]) == pytest.approx(0.0, abs=1e-8) + + def test_rejects_bad_bank_rank(self): + with pytest.raises(ValueError, match="2-D"): + ChemicalSupportIndex(torch.tensor([1.0, 2.0]), radius=1.0) + + def test_no_molpot_import_in_support_module(self): + """molrep embedding support must not depend on molpot (ac-006 spirit).""" + path = ( + Path(__file__).resolve().parents[3] + / "src" + / "molrep" + / "embedding" + / "support.py" + ) + tree = ast.parse(path.read_text()) + for node in ast.walk(tree): + if isinstance(node, ast.Import): + for alias in node.names: + assert not alias.name.startswith("molpot"), alias.name + if isinstance(node, ast.ImportFrom) and node.module: + assert not node.module.startswith("molpot"), node.module diff --git a/tests/test_molrep/test_heads/__init__.py b/tests/test_molrep/test_heads/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/test_molrep/test_heads/test_type_system_labeler.py b/tests/test_molrep/test_heads/test_type_system_labeler.py new file mode 100644 index 0000000..b997c13 --- /dev/null +++ b/tests/test_molrep/test_heads/test_type_system_labeler.py @@ -0,0 +1,99 @@ +"""Tests for TypeSystemLabeler and MultiTypeHead (learnable-classical-ff-06).""" + +from __future__ import annotations + +import torch + +from molrep.condensation import ( + Condenser, + InteractionClass, + TypeSystem, + bond_default_criterion, +) +from molrep.heads import Labeler, MultiTypeHead, TypeHead, TypeSystemLabeler + + +class TestTypeSystemLabeler: + def test_labeler_protocol_surface(self): + ts = TypeSystem.from_prototypes( + InteractionClass.BOND, + [{"k": 300.0, "r0": 1.09}, {"k": 500.0, "r0": 1.5}], + criterion=bond_default_criterion(), + ) + labeler = TypeSystemLabeler(ts) + assert isinstance(labeler, Labeler) + assert labeler.num_types == 2 + assert set(labeler.type_map.keys()) == {0, 1} + assert all(isinstance(v, str) for v in labeler.type_map.values()) + + def test_label_params_in_range(self): + ts = TypeSystem.from_prototypes( + InteractionClass.BOND, + [{"k": 300.0, "r0": 1.09}, {"k": 500.0, "r0": 1.5}], + criterion=bond_default_criterion(), + ) + labeler = TypeSystemLabeler(ts) + params = { + "k": torch.tensor([300.0, 500.0]), + "r0": torch.tensor([1.09, 1.5]), + } + ids = labeler.label(params) + assert ids.dtype == torch.long + assert ids.tolist() == [0, 1] + assert all(0 <= int(i) < labeler.num_types for i in ids) + + def test_label_from_condenser_assignment(self): + result = Condenser().merge( + [{"k": torch.tensor([300.0, 300.0]), "r0": torch.tensor([1.09, 1.09])}], + interaction=InteractionClass.BOND, + criterion=bond_default_criterion(), + ) + labeler = TypeSystemLabeler(result.type_system) + ids = labeler.label(result.assignment.type_ids) + assert ids.tolist() == [0, 0] + + +class TestTypeHeadMultiSystem: + def test_existing_type_head_api(self): + head = TypeHead(hidden_dim=4, num_types=5) + out = head(torch.ones(3, 4)) + assert out.shape == torch.Size([3, 5]) + idx = head.decode(out) + assert idx.shape == torch.Size([3]) + idx2, conf = head.decode_with_confidence(out) + assert idx2.shape == conf.shape == torch.Size([3]) + + def test_from_type_system(self): + ts = TypeSystem.from_prototypes( + InteractionClass.BOND, + [{"k": 1.0, "r0": 1.0}, {"k": 2.0, "r0": 1.1}], + ) + head = TypeHead.from_type_system(8, ts) + assert head.num_types == 2 + logits = head(torch.randn(4, 8)) + assert logits.shape == (4, 2) + + def test_multi_type_head_additive(self): + bond_ts = TypeSystem.from_prototypes( + InteractionClass.BOND, + [{"k": 1.0, "r0": 1.0}], + ) + angle_ts = TypeSystem.from_prototypes( + InteractionClass.ANGLE, + [{"k": 10.0, "theta0": 1.9}, {"k": 20.0, "theta0": 2.0}], + ) + multi = MultiTypeHead.from_type_systems( + 4, + {"bond": bond_ts, "angle": angle_ts}, + ) + assert multi.num_types == {"bond": 1, "angle": 2} + logits = multi( + { + "bond": torch.randn(2, 4), + "angle": torch.randn(3, 4), + } + ) + assert logits["bond"].shape == (2, 1) + assert logits["angle"].shape == (3, 2) + decoded = multi.decode(logits) + assert decoded["bond"].shape == (2,) diff --git a/tests/test_molrep/test_interaction/__init__.py b/tests/test_molrep/test_interaction/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/test_molrep/test_interaction/test_element.py b/tests/test_molrep/test_interaction/test_element.py index 41ea5bd..1405b6b 100644 --- a/tests/test_molrep/test_interaction/test_element.py +++ b/tests/test_molrep/test_interaction/test_element.py @@ -44,8 +44,10 @@ def test_initialization(self): update = ElementUpdate(hidden_dim=64, num_species=10) assert update.config.hidden_dim == 64 assert update.config.num_species == 10 - assert hasattr(update, "_linear_indexed") + # One backend only: the sorted-index `indexed_linear` kernel measured + # slower than `naive` once the required argsort round-trip is counted. assert hasattr(update, "_linear_naive") + assert not hasattr(update, "_linear_indexed") def test_forward_shape(self): """Test output shape.""" diff --git a/tests/test_molrep/test_interaction/test_mace/__init__.py b/tests/test_molrep/test_interaction/test_mace/__init__.py new file mode 100644 index 0000000..f5bd897 --- /dev/null +++ b/tests/test_molrep/test_interaction/test_mace/__init__.py @@ -0,0 +1 @@ +"""MACE interaction block unit tests.""" diff --git a/tests/test_molrep/test_interaction/test_mace/test_block.py b/tests/test_molrep/test_interaction/test_mace/test_block.py new file mode 100644 index 0000000..77ef169 --- /dev/null +++ b/tests/test_molrep/test_interaction/test_mace/test_block.py @@ -0,0 +1,198 @@ +"""Tests for molrep.interaction.mace.block module.""" + +import cuequivariance_torch as cuet +import pytest +import torch +import torch.nn as nn + +from molrep.interaction.mace.block import InteractionBlock, InteractionSpec +from molrep.interaction.mace.conv import ConvTP +from molrep.interaction.radial import RadialWeightMLP + +N_NODES = 4 +FEATURES = 8 +NUM_BESSEL = 5 +L_MAX = 1 +SH_DIM = (L_MAX + 1) ** 2 # 1x0e + 1x1o +MESSAGE_DIM = FEATURES * SH_DIM # mixed-l message irreps: 8x0e + 8x1o + + +@pytest.fixture +def graph(): + """A small fp64 directed graph, ``edge_index`` in ``(E, 2)`` layout.""" + torch.manual_seed(0) + edge_index = torch.tensor([[0, 1, 2, 3, 0], [1, 2, 3, 0, 2]]).t().contiguous() # (E, 2) + n_edges = edge_index.shape[0] + return { + "node_feats": torch.randn(N_NODES, FEATURES, dtype=torch.float64), + "edge_attrs": torch.randn(n_edges, SH_DIM, dtype=torch.float64), + "edge_feats": torch.randn(n_edges, NUM_BESSEL, dtype=torch.float64), + "edge_index": edge_index, + } + + +def _block(avg_num_neighbors: float) -> InteractionBlock: + """Build an fp64 block; the seed makes two calls weight-identical.""" + torch.manual_seed(0) + return InteractionBlock( + num_features=FEATURES, + num_bessel=NUM_BESSEL, + l_max=L_MAX, + avg_num_neighbors=avg_num_neighbors, + ).double() + + +class TestInteractionSpec: + """Test InteractionSpec configuration.""" + + def test_stores_avg_num_neighbors(self): + """The normalisation constant round-trips through the pydantic spec.""" + spec = InteractionSpec( + num_features=FEATURES, + num_bessel=NUM_BESSEL, + l_max=L_MAX, + avg_num_neighbors=2.0, + ) + assert spec.avg_num_neighbors == 2.0 + + +class TestInteractionBlock: + """Test the promoted MACE interaction block.""" + + def test_forward_output_shape(self, graph): + """The message carries the mixed-l irreps ``(N, num_features * (l_max+1)**2)``.""" + node_feats, _ = _block(1.0)(**graph) + assert node_feats.shape == (N_NODES, MESSAGE_DIM) + + def test_skip_connection_is_the_input_tensor(self, graph): + """``sc`` is the untouched input object, not a copy — the product block needs it.""" + _, sc = _block(1.0)(**graph) + assert sc is graph["node_feats"] + + def test_avg_num_neighbors_scales_message(self, graph): + """Doubling ``avg_num_neighbors`` halves the message (post-conv linear is bias-free).""" + one, _ = _block(1.0)(**graph) + two, _ = _block(2.0)(**graph) + assert torch.allclose(two * 2.0, one, rtol=0.0, atol=1e-12) + + def test_double_precision_roundtrip(self, graph): + """`.double()` must work: cuEq bakes its dtype in at construction.""" + node_feats, _ = _block(1.0)(**graph) + assert node_feats.dtype == torch.float64 + + def test_state_dict_keys_are_the_weight_transfer_contract(self): + """Sub-layer names must stay put — official MACE weights load by key.""" + assert sorted(_block(1.0).state_dict().keys()) == [ + "conv_tp.cue_tp.f.m.graphs.0.graph.c0", + "conv_tp.cue_tp.f.m.graphs.0.graph.c1", + "linear.f.m.graphs.0.graph.c0", + "linear.f.m.graphs.0.graph.c1", + "linear.weight", + "node_linear.f.m.graphs.0.graph.c0", + "node_linear.weight", + "radial_mlp.mlp.0.bias", + "radial_mlp.mlp.0.weight", + "radial_mlp.mlp.2.bias", + "radial_mlp.mlp.2.weight", + "radial_mlp.mlp.4.bias", + "radial_mlp.mlp.4.weight", + ] + + # -- migrated from tests/test_molzoo/test_mace.py by + # mace-subpackage-restructure-02-core (module/package name clash forced + # that file's removal; these cases have no equivalent above) ---------- + + @pytest.fixture + def interaction_config(self): + """Wider configuration inherited with the migrated cases.""" + return { + "num_features": 64, + "num_bessel": 8, + "l_max": 2, + "avg_num_neighbors": 10.0, + } + + @pytest.fixture + def interaction_block(self, interaction_config): + """Create an InteractionBlock instance.""" + return InteractionBlock(**interaction_config) + + def test_initialization_wires_components(self, interaction_block, interaction_config): + """Every sub-block is present and of the expected type.""" + # conv_tp must be created first to provide weight_numel + assert isinstance(interaction_block.conv_tp, ConvTP) + assert isinstance(interaction_block.node_linear, cuet.Linear) + assert isinstance(interaction_block.radial_mlp, RadialWeightMLP) + assert isinstance(interaction_block.linear, cuet.Linear) + assert interaction_block.avg_num_neighbors == interaction_config["avg_num_neighbors"] + + def test_initialization_order_fix(self, interaction_config): + """conv_tp is initialized before radial_mlp (regression guard). + + Building the radial MLP first raised ``AttributeError: self.conv_tp``, + because the MLP's output width is ``conv_tp.weight_numel``. + """ + block = InteractionBlock(**interaction_config) + assert hasattr(block.conv_tp, "weight_numel") + assert block.radial_mlp.mlp[-1].out_features == block.conv_tp.weight_numel + + def test_config_storage(self, interaction_block, interaction_config): + """Constructor kwargs round-trip through the pydantic ``config``.""" + config = interaction_block.config + assert config.num_features == interaction_config["num_features"] + assert config.num_bessel == interaction_config["num_bessel"] + assert config.l_max == interaction_config["l_max"] + assert config.avg_num_neighbors == interaction_config["avg_num_neighbors"] + + def test_radial_mlp_architecture(self, interaction_block, interaction_config): + """``Linear → SiLU → Linear → SiLU → Linear`` with MACE's widths.""" + mlp = interaction_block.radial_mlp.mlp + num_features = interaction_config["num_features"] + num_bessel = interaction_config["num_bessel"] + weight_numel = interaction_block.conv_tp.weight_numel + + assert len(mlp) == 5 + assert isinstance(mlp[0], nn.Linear) + assert isinstance(mlp[1], nn.SiLU) + assert isinstance(mlp[2], nn.Linear) + assert isinstance(mlp[3], nn.SiLU) + assert isinstance(mlp[4], nn.Linear) + + assert mlp[0].in_features == num_bessel + assert mlp[0].out_features == num_features + assert mlp[2].in_features == num_features + assert mlp[2].out_features == num_features + assert mlp[4].in_features == num_features + assert mlp[4].out_features == weight_numel + + def test_cuequivariance_integration(self, interaction_block): + """The convolution is a cuEq ``ChannelWiseTensorProduct``, not a port.""" + cue_tp = interaction_block.conv_tp.cue_tp + assert isinstance(cue_tp, cuet.ChannelWiseTensorProduct) + assert hasattr(cue_tp, "irreps_in1") + assert hasattr(cue_tp, "irreps_in2") + assert hasattr(cue_tp, "irreps_out") + + @pytest.mark.parametrize("l_max", [1, 2, 3]) + def test_different_l_max_values(self, interaction_config, l_max): + """Construction succeeds for every angular order MACE ships.""" + block = InteractionBlock( + num_features=interaction_config["num_features"], + num_bessel=interaction_config["num_bessel"], + l_max=l_max, + avg_num_neighbors=interaction_config["avg_num_neighbors"], + ) + assert block.config.l_max == l_max + assert hasattr(block.conv_tp, "weight_numel") + + @pytest.mark.parametrize("num_features", [32, 64, 128]) + def test_different_num_features(self, interaction_config, num_features): + """The radial MLP hidden width tracks ``num_features``.""" + block = InteractionBlock( + num_features=num_features, + num_bessel=interaction_config["num_bessel"], + l_max=interaction_config["l_max"], + avg_num_neighbors=interaction_config["avg_num_neighbors"], + ) + assert block.radial_mlp.mlp[0].out_features == num_features + assert block.radial_mlp.mlp[2].in_features == num_features diff --git a/tests/test_molrep/test_interaction/test_mace/test_conv.py b/tests/test_molrep/test_interaction/test_mace/test_conv.py new file mode 100644 index 0000000..e0bb3f0 --- /dev/null +++ b/tests/test_molrep/test_interaction/test_mace/test_conv.py @@ -0,0 +1,196 @@ +"""Tests for molrep.interaction.mace.conv module.""" + +import math + +import torch + +from molrep.interaction.mace.conv import ConvTP, ConvTPSpec +from molrep.utils.equivariance import ( + check_equivariance, + rotate_irreps_features_simple, + rotate_vectors, + rotation_matrix_z, +) + + +class TestConvTPSpec: + """Test ConvTPSpec configuration.""" + + def test_valid_config(self): + """Test creation with valid parameters.""" + spec = ConvTPSpec( + in_irreps="64x0e", + out_irreps="64x0e", + sh_irreps="1x0e + 1x1o", + ) + assert spec.in_irreps == "64x0e" + assert spec.out_irreps == "64x0e" + assert spec.sh_irreps == "1x0e + 1x1o" + + +class TestConvTP: + """Test ConvTP tensor product layer.""" + + def test_initialization(self): + """Test ConvTP initialization.""" + tp = ConvTP( + in_irreps="64x0e", + out_irreps="64x0e", + sh_irreps="1x0e + 1x1o", + ) + assert tp.config.in_irreps == "64x0e" + assert tp.config.out_irreps == "64x0e" + + def test_forward_shape(self): + """Test output shape.""" + tp = ConvTP( + in_irreps="16x0e", + out_irreps="16x0e", + sh_irreps="1x0e + 1x1o", + ) + + n_nodes = 10 + n_edges = 30 + node_features = torch.randn(n_nodes, 16) + edge_angular = torch.randn(n_edges, 4) # 1 + 3 = 4 dims + edge_index = torch.randint(0, n_nodes, (n_edges, 2)) + + # ConvTP weights usually depend on radial embedding + # cuet.Linear handles the weights if passed correctly. + # Wait, ConvTP expects tp_weights. + # Let's check how many weights are needed. + # In ConvTP, self.cue_tp is initialized with weight_dim. + weight_dim = tp.cue_tp.weight_numel + tp_weights = torch.randn(n_edges, weight_dim) + + output = tp(node_features, edge_angular, edge_index, tp_weights) + + # Messages should be per node after aggregation, not per edge + assert output.shape[0] == n_nodes + assert output.shape[1] == 16 + + def test_different_irreps(self): + """Test with different input/output irreps.""" + tp = ConvTP( + in_irreps="32x0e", + out_irreps="64x0e", + sh_irreps="1x0e + 1x1o + 1x2e", + ) + assert tp.config.in_irreps == "32x0e" + assert tp.config.out_irreps == "64x0e" + + def test_equivariance_scalars(self): + """Test equivariance with scalar features only. + + For scalar inputs, tensor product with spherical harmonics should + maintain equivariance properties. + """ + tp = ConvTP( + in_irreps="16x0e", + out_irreps="16x0e", + sh_irreps="1x0e + 1x1o", + ) + + n_nodes = 5 + n_edges = 10 + node_features = torch.randn(n_nodes, 16) + edge_angular = torch.randn(n_edges, 4) # 1 + 3 dims + edge_index = torch.randint(0, n_nodes, (n_edges, 2)) + + weight_dim = tp.weight_numel + tp_weights = torch.randn(n_edges, weight_dim) + + # Forward pass + output1 = tp(node_features, edge_angular, edge_index, tp_weights) + + # Rotate edge angular features (spherical harmonics) + angle = math.pi / 2 + rot_matrix = rotation_matrix_z(angle, dtype=edge_angular.dtype) + + # For spherical harmonics, we need to rotate the vector components (l=1) + # l=0 is invariant, l=1 needs rotation + edge_angular_rotated = edge_angular.clone() + edge_angular_rotated[:, 1:4] = rotate_vectors(edge_angular[:, 1:4], rot_matrix) + + # Forward on rotated + output2 = tp(node_features, edge_angular_rotated, edge_index, tp_weights) + + # For scalar outputs, they should be approximately equal (rotation invariant) + assert torch.allclose(output1, output2, rtol=1e-3, atol=1e-3) + + def test_equivariance_vectors(self): + """Test equivariance with vector features. + + ConvTP should satisfy: TP(R·h, R·Y) = R·TP(h, Y) + where h are node features, Y are spherical harmonics, R is rotation. + """ + # Vector input and output + in_irreps = "4x1o" + out_irreps = "4x1o" + sh_irreps = "1x0e + 1x1o" + + tp = ConvTP( + in_irreps=in_irreps, + out_irreps=out_irreps, + sh_irreps=sh_irreps, + ) + + n_nodes = 5 + n_edges = 10 + node_features = torch.randn(n_nodes, 12) # 4 vectors * 3 + edge_angular = torch.randn(n_edges, 4) # 1 + 3 + edge_index = torch.randint(0, n_nodes, (n_edges, 2)) + + weight_dim = tp.weight_numel + tp_weights = torch.randn(n_edges, weight_dim) + + # Forward pass + output1 = tp(node_features, edge_angular, edge_index, tp_weights) + + # Rotate everything + angle = math.pi / 2 + rot_matrix = rotation_matrix_z(angle, dtype=node_features.dtype) + + # Rotate node features + node_features_rot = rotate_irreps_features_simple(node_features, rot_matrix, in_irreps) + + # Rotate edge angular (spherical harmonics) + edge_angular_rot = edge_angular.clone() + edge_angular_rot[:, 1:4] = rotate_vectors(edge_angular[:, 1:4], rot_matrix) + + # Forward on rotated inputs + output2 = tp(node_features_rot, edge_angular_rot, edge_index, tp_weights) + + # Rotate output1 + output1_rot = rotate_irreps_features_simple(output1, rot_matrix, out_irreps) + + # Check equivariance + assert check_equivariance(output1_rot, output2, rtol=1e-3, atol=1e-3) + + def test_differentiable(self): + """Test that gradients flow through ConvTP.""" + tp = ConvTP( + in_irreps="8x0e", + out_irreps="8x0e", + sh_irreps="1x0e + 1x1o", + ) + + n_nodes = 5 + n_edges = 10 + node_features = torch.randn(n_nodes, 8, requires_grad=True) + edge_angular = torch.randn(n_edges, 4, requires_grad=True) + edge_index = torch.randint(0, n_nodes, (n_edges, 2)) + + weight_dim = tp.weight_numel + tp_weights = torch.randn(n_edges, weight_dim, requires_grad=True) + + output = tp(node_features, edge_angular, edge_index, tp_weights) + loss = output.sum() + loss.backward() + + assert node_features.grad is not None + assert edge_angular.grad is not None + assert tp_weights.grad is not None + assert not torch.isnan(node_features.grad).any() + assert not torch.isnan(edge_angular.grad).any() + assert not torch.isnan(tp_weights.grad).any() diff --git a/tests/test_molrep/test_interaction/test_mace/test_density.py b/tests/test_molrep/test_interaction/test_mace/test_density.py new file mode 100644 index 0000000..10b95ed --- /dev/null +++ b/tests/test_molrep/test_interaction/test_mace/test_density.py @@ -0,0 +1,153 @@ +"""Tests for molrep.interaction.mace.density module.""" + +import pytest +import torch + +from molrep.interaction.mace.density import DensityInteraction, DensityResidualInteraction + +N_ELEMENTS = 3 +N_NODES = 4 +NUM_RADIAL = 5 +FEATURES = 8 +SH = "1x0e+1x1o" +TARGET = f"{FEATURES}x0e+{FEATURES}x1o" + + +@pytest.fixture +def graph(): + """A small directed graph plus the edge/node inputs the blocks consume.""" + torch.manual_seed(0) + edge_index = torch.tensor([[0, 1, 2, 3, 0], [1, 2, 3, 0, 2]]).t().contiguous() # (E, 2) + n_edges = edge_index.shape[0] + node_attrs = torch.zeros(N_NODES, N_ELEMENTS, dtype=torch.float64) + node_attrs[torch.arange(N_NODES), torch.tensor([0, 1, 2, 0])] = 1.0 + return { + "node_attrs": node_attrs, + "node_feats": torch.randn(N_NODES, FEATURES, dtype=torch.float64), + "edge_attrs": torch.randn(n_edges, 4, dtype=torch.float64), + "edge_feats": torch.randn(n_edges, NUM_RADIAL, dtype=torch.float64), + "edge_index": edge_index, + } + + +def _common(): + return dict( + node_attrs_irreps=f"{N_ELEMENTS}x0e", + edge_attrs_irreps=SH, + edge_feats_irreps=f"{NUM_RADIAL}x0e", + target_irreps=TARGET, + radial_mlp=[8], + ) + + +@pytest.fixture +def first_layer(): + return DensityInteraction( + node_feats_irreps=f"{FEATURES}x0e", edge_irreps=f"{FEATURES}x0e", **_common() + ).double() + + +@pytest.fixture +def residual_layer(): + return DensityResidualInteraction( + node_feats_irreps=f"{FEATURES}x0e", + edge_irreps=f"{FEATURES}x0e", + hidden_irreps=f"{FEATURES}x0e", + **_common(), + ).double() + + +class TestDensityInteraction: + """Test the first-layer density-normalised interaction.""" + + def test_output_shape_is_multiplet_layout(self, first_layer, graph): + """reshape_irreps emits ``(N, ir_dim, mul)`` for the product block.""" + out, _ = first_layer(**graph) + assert out.shape == (N_NODES, 4, FEATURES) + + def test_returns_no_skip_connection(self, first_layer, graph): + """The first layer has no residual to carry — MACE returns ``None``.""" + assert first_layer(**graph)[1] is None + + def test_double_precision_roundtrip(self, graph): + """`.double()` must work: cuEq bakes its dtype in at construction.""" + layer = DensityInteraction( + node_feats_irreps=f"{FEATURES}x0e", edge_irreps=f"{FEATURES}x0e", **_common() + ).double() + out, _ = layer(**graph) + assert out.dtype == torch.float64 + + def test_skip_tp_weight_survives_dtype_change(self, graph): + """Rebuilding skip_tp on `.double()` must preserve its trained weight.""" + layer = DensityInteraction( + node_feats_irreps=f"{FEATURES}x0e", edge_irreps=f"{FEATURES}x0e", **_common() + ) + before = layer.skip_tp.weight.detach().clone() + layer = layer.double() + assert torch.allclose(layer.skip_tp.weight.detach(), before.double()) + + def test_density_normalisation_damps_high_coordination(self, first_layer, graph): + """A node gathering more edges is divided by a larger density.""" + dense = dict(graph) + # Route every edge at node 2 so its density is the largest in the graph. + dense["edge_index"] = torch.tensor([[0, 1, 3, 0, 1], [2, 2, 2, 2, 2]]).t() # (E, 2) + out, _ = first_layer(**dense) + assert torch.isfinite(out).all() + + def test_isolated_node_is_finite(self, first_layer, graph): + """Zero density must not divide by zero — the ``+1`` guarantees it.""" + lonely = dict(graph) + lonely["edge_index"] = torch.tensor([[0, 1], [1, 0]]).t() # (E, 2) + lonely["edge_attrs"] = graph["edge_attrs"][:2] + lonely["edge_feats"] = graph["edge_feats"][:2] + out, _ = first_layer(**lonely) + assert torch.isfinite(out).all() + # Nodes 2 and 3 receive nothing, so their message is exactly zero. + assert torch.count_nonzero(out[2]) == 0 + assert torch.count_nonzero(out[3]) == 0 + + +class TestDensityResidualInteraction: + """Test the residual density-normalised interaction.""" + + def test_output_shape_is_multiplet_layout(self, residual_layer, graph): + """Message keeps the same multiplet layout as the first layer.""" + out, _ = residual_layer(**graph) + assert out.shape == (N_NODES, 4, FEATURES) + + def test_returns_skip_connection(self, residual_layer, graph): + """The residual variant hands a skip connection to the product block.""" + _, skip = residual_layer(**graph) + assert skip.shape == (N_NODES, FEATURES) + + def test_skip_depends_only_on_node_inputs(self, residual_layer, graph): + """``skip_tp`` reads node features and attrs — not the edges.""" + _, skip_a = residual_layer(**graph) + rewired = dict(graph) + rewired["edge_index"] = torch.tensor([[3, 2, 1, 0, 2], [0, 3, 2, 1, 0]]).t() # (E, 2) + _, skip_b = residual_layer(**rewired) + assert torch.allclose(skip_a, skip_b) + + def test_permutation_equivariance(self, residual_layer, graph): + """Relabelling atoms permutes the output identically.""" + perm = torch.tensor([2, 0, 3, 1]) + inverse = torch.argsort(perm) + out, skip = residual_layer(**graph) + + permuted = dict(graph) + permuted["node_attrs"] = graph["node_attrs"][perm] + permuted["node_feats"] = graph["node_feats"][perm] + permuted["edge_index"] = inverse[graph["edge_index"]] + out_p, skip_p = residual_layer(**permuted) + + assert torch.allclose(out_p, out[perm], atol=1e-12) + assert torch.allclose(skip_p, skip[perm], atol=1e-12) + + def test_cutoff_scales_messages(self, residual_layer, graph): + """A zero cutoff kills every message but leaves the skip untouched.""" + n_edges = graph["edge_index"].shape[0] + zero = torch.zeros(n_edges, 1, dtype=torch.float64) + out, skip = residual_layer(cutoff=zero, **graph) + _, skip_ref = residual_layer(**graph) + assert torch.count_nonzero(out) == 0 + assert torch.allclose(skip, skip_ref) diff --git a/tests/test_molrep/test_interaction/test_product.py b/tests/test_molrep/test_interaction/test_product.py index 7da3e5b..1d3478a 100644 --- a/tests/test_molrep/test_interaction/test_product.py +++ b/tests/test_molrep/test_interaction/test_product.py @@ -1,208 +1,8 @@ """Tests for molrep.interaction.product module.""" -import math - import torch -from molrep.interaction.product import ( - ConvTP, - ConvTPSpec, - EquivariantPolynomialTP, -) -from molrep.utils.equivariance import ( - check_equivariance, - rotate_irreps_features_simple, - rotate_vectors, - rotation_matrix_z, -) - - -class TestConvTPSpec: - """Test ConvTPSpec configuration.""" - - def test_valid_config(self): - """Test creation with valid parameters.""" - spec = ConvTPSpec( - in_irreps="64x0e", - out_irreps="64x0e", - sh_irreps="1x0e + 1x1o", - ) - assert spec.in_irreps == "64x0e" - assert spec.out_irreps == "64x0e" - assert spec.sh_irreps == "1x0e + 1x1o" - - -class TestConvTP: - """Test ConvTP tensor product layer.""" - - def test_initialization(self): - """Test ConvTP initialization.""" - tp = ConvTP( - in_irreps="64x0e", - out_irreps="64x0e", - sh_irreps="1x0e + 1x1o", - ) - assert tp.config.in_irreps == "64x0e" - assert tp.config.out_irreps == "64x0e" - - def test_forward_shape(self): - """Test output shape.""" - tp = ConvTP( - in_irreps="16x0e", - out_irreps="16x0e", - sh_irreps="1x0e + 1x1o", - ) - - n_nodes = 10 - n_edges = 30 - node_features = torch.randn(n_nodes, 16) - edge_angular = torch.randn(n_edges, 4) # 1 + 3 = 4 dims - edge_index = torch.randint(0, n_nodes, (n_edges, 2)) - - # ConvTP weights usually depend on radial embedding - # cuet.Linear handles the weights if passed correctly. - # Wait, ConvTP expects tp_weights. - # Let's check how many weights are needed. - # In ConvTP, self.cue_tp is initialized with weight_dim. - weight_dim = tp.cue_tp.weight_numel - tp_weights = torch.randn(n_edges, weight_dim) - - output = tp(node_features, edge_angular, edge_index, tp_weights) - - # Messages should be per node after aggregation, not per edge - assert output.shape[0] == n_nodes - assert output.shape[1] == 16 - - def test_different_irreps(self): - """Test with different input/output irreps.""" - tp = ConvTP( - in_irreps="32x0e", - out_irreps="64x0e", - sh_irreps="1x0e + 1x1o + 1x2e", - ) - assert tp.config.in_irreps == "32x0e" - assert tp.config.out_irreps == "64x0e" - - def test_equivariance_scalars(self): - """Test equivariance with scalar features only. - - For scalar inputs, tensor product with spherical harmonics should - maintain equivariance properties. - """ - tp = ConvTP( - in_irreps="16x0e", - out_irreps="16x0e", - sh_irreps="1x0e + 1x1o", - ) - - n_nodes = 5 - n_edges = 10 - node_features = torch.randn(n_nodes, 16) - edge_angular = torch.randn(n_edges, 4) # 1 + 3 dims - edge_index = torch.randint(0, n_nodes, (n_edges, 2)) - - weight_dim = tp.weight_numel - tp_weights = torch.randn(n_edges, weight_dim) - - # Forward pass - output1 = tp(node_features, edge_angular, edge_index, tp_weights) - - # Rotate edge angular features (spherical harmonics) - angle = math.pi / 2 - rot_matrix = rotation_matrix_z(angle, dtype=edge_angular.dtype) - - # For spherical harmonics, we need to rotate the vector components (l=1) - # l=0 is invariant, l=1 needs rotation - edge_angular_rotated = edge_angular.clone() - edge_angular_rotated[:, 1:4] = rotate_vectors(edge_angular[:, 1:4], rot_matrix) - - # Forward on rotated - output2 = tp(node_features, edge_angular_rotated, edge_index, tp_weights) - - # For scalar outputs, they should be approximately equal (rotation invariant) - assert torch.allclose(output1, output2, rtol=1e-3, atol=1e-3) - - def test_equivariance_vectors(self): - """Test equivariance with vector features. - - ConvTP should satisfy: TP(R·h, R·Y) = R·TP(h, Y) - where h are node features, Y are spherical harmonics, R is rotation. - """ - # Vector input and output - in_irreps = "4x1o" - out_irreps = "4x1o" - sh_irreps = "1x0e + 1x1o" - - tp = ConvTP( - in_irreps=in_irreps, - out_irreps=out_irreps, - sh_irreps=sh_irreps, - ) - - n_nodes = 5 - n_edges = 10 - node_features = torch.randn(n_nodes, 12) # 4 vectors * 3 - edge_angular = torch.randn(n_edges, 4) # 1 + 3 - edge_index = torch.randint(0, n_nodes, (n_edges, 2)) - - weight_dim = tp.weight_numel - tp_weights = torch.randn(n_edges, weight_dim) - - # Forward pass - output1 = tp(node_features, edge_angular, edge_index, tp_weights) - - # Rotate everything - angle = math.pi / 2 - rot_matrix = rotation_matrix_z(angle, dtype=node_features.dtype) - - # Rotate node features - node_features_rot = rotate_irreps_features_simple(node_features, rot_matrix, in_irreps) - - # Rotate edge angular (spherical harmonics) - edge_angular_rot = edge_angular.clone() - edge_angular_rot[:, 1:4] = rotate_vectors(edge_angular[:, 1:4], rot_matrix) - - # Forward on rotated inputs - output2 = tp(node_features_rot, edge_angular_rot, edge_index, tp_weights) - - # Rotate output1 - output1_rot = rotate_irreps_features_simple(output1, rot_matrix, out_irreps) - - # Check equivariance - assert check_equivariance(output1_rot, output2, rtol=1e-3, atol=1e-3) - - def test_differentiable(self): - """Test that gradients flow through ConvTP.""" - tp = ConvTP( - in_irreps="8x0e", - out_irreps="8x0e", - sh_irreps="1x0e + 1x1o", - ) - - n_nodes = 5 - n_edges = 10 - node_features = torch.randn(n_nodes, 8, requires_grad=True) - edge_angular = torch.randn(n_edges, 4, requires_grad=True) - edge_index = torch.randint(0, n_nodes, (n_edges, 2)) - - weight_dim = tp.weight_numel - tp_weights = torch.randn(n_edges, weight_dim, requires_grad=True) - - output = tp(node_features, edge_angular, edge_index, tp_weights) - loss = output.sum() - loss.backward() - - assert node_features.grad is not None - assert edge_angular.grad is not None - assert tp_weights.grad is not None - assert not torch.isnan(node_features.grad).any() - assert not torch.isnan(edge_angular.grad).any() - assert not torch.isnan(tp_weights.grad).any() - - -# =========================================================================== -# EquivariantPolynomialTP -# =========================================================================== +from molrep.interaction.product import EquivariantPolynomialTP class TestEquivariantPolynomialTP: diff --git a/tests/test_molrep/test_perception/__init__.py b/tests/test_molrep/test_perception/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/test_molrep/test_perception/test_forcefield.py b/tests/test_molrep/test_perception/test_forcefield.py new file mode 100644 index 0000000..10fa7db --- /dev/null +++ b/tests/test_molrep/test_perception/test_forcefield.py @@ -0,0 +1,173 @@ +"""SymbolicForceField records + match_molecule (learnable-classical-ff-07).""" + +from __future__ import annotations + +import ast +from pathlib import Path + +import pytest +import torch + +from molrep.condensation import InteractionClass, TypeRecord, TypeSystem +from molrep.perception import ( + ClassPatternRegistry, + DiscreteClassRecord, + FakeSmartsMatcher, + SymbolicForceField, + SymbolicPattern, +) + +FORCEFIELD_PATH = ( + Path(__file__).resolve().parents[3] / "src" / "molrep" / "perception" / "forcefield.py" +) +PERCEPTION_ROOT = FORCEFIELD_PATH.parent + + +def _bond_type_system() -> TypeSystem: + return TypeSystem( + InteractionClass.BOND, + [ + TypeRecord( + type_id=0, + prototype={"k": 300.0, "r0": 1.09}, + label="CT-CT", + ), + TypeRecord( + type_id=1, + prototype={"k": 320.0, "r0": 1.41}, + label="CT-OH", + ), + ], + ) + + +class TestSymbolicForceFieldRecords: + def test_records_pair_prototypes_with_patterns(self): + ts = _bond_type_system() + reg = ClassPatternRegistry() + reg.bind( + InteractionClass.BOND, + 0, + SymbolicPattern("[#6]-[#6]", arity=2), + ) + reg.bind( + InteractionClass.BOND, + 1, + SymbolicPattern("[#6]-[#8]", arity=2), + ) + ff = SymbolicForceField({InteractionClass.BOND: ts}, reg) + recs = ff.records() + assert len(recs) == 2 + assert all(isinstance(r, DiscreteClassRecord) for r in recs) + by_id = {r.type_id: r for r in recs} + assert by_id[0].smarts == "[#6]-[#6]" + assert by_id[0].prototype["r0"] == 1.09 + assert by_id[0].label == "CT-CT" + assert by_id[1].smarts == "[#6]-[#8]" + assert by_id[1].smirks is None + + def test_unbound_type_has_none_smarts(self): + ts = _bond_type_system() + reg = ClassPatternRegistry() + reg.bind( + InteractionClass.BOND, + 0, + SymbolicPattern("[#6]-[#6]", arity=2), + ) + ff = SymbolicForceField({InteractionClass.BOND: ts}, reg) + recs = {r.type_id: r for r in ff.records()} + assert recs[0].smarts == "[#6]-[#6]" + assert recs[1].smarts is None + + def test_smirks_kind_fills_smirks_field(self): + ts = TypeSystem( + InteractionClass.BOND, + [TypeRecord(type_id=0, prototype={"k": 1.0, "r0": 1.0})], + ) + reg = ClassPatternRegistry() + reg.bind( + InteractionClass.BOND, + 0, + SymbolicPattern("[C:1][O:2]>>[C:1][O:2]", arity=2, kind="smirks"), + ) + ff = SymbolicForceField({InteractionClass.BOND: ts}, reg) + rec = ff.records()[0] + assert rec.smirks == "[C:1][O:2]>>[C:1][O:2]" + assert rec.smarts is None + + +class TestSymbolicForceFieldMatchMolecule: + def test_match_molecule_assigns_types_via_fake(self): + ts = _bond_type_system() + reg = ClassPatternRegistry() + reg.bind( + InteractionClass.BOND, + 0, + SymbolicPattern("[#6]-[#6]", arity=2), + ) + reg.bind( + InteractionClass.BOND, + 1, + SymbolicPattern("[#6]-[#8]", arity=2), + ) + ff = SymbolicForceField({InteractionClass.BOND: ts}, reg) + + matcher = FakeSmartsMatcher( + { + "[#6]-[#6]": torch.tensor([[0], [1]], dtype=torch.long), + "[#6]-[#8]": torch.tensor([[1, 2], [2, 3]], dtype=torch.long), + } + ) + assigned = ff.match_molecule(mol=object(), matcher=matcher) + assert InteractionClass.BOND in assigned + matches = assigned[InteractionClass.BOND]["matches"] + type_ids = assigned[InteractionClass.BOND]["type_ids"] + assert matches.shape[0] == 2 # arity + assert matches.shape[1] == 3 # 1 + 2 hits + assert type_ids.tolist() == [0, 1, 1] + # first column is the C-C hit + assert matches[:, 0].tolist() == [0, 1] + + def test_no_hits_omits_interaction(self): + ts = _bond_type_system() + reg = ClassPatternRegistry() + reg.bind( + InteractionClass.BOND, + 0, + SymbolicPattern("[#6]-[#6]", arity=2), + ) + ff = SymbolicForceField({InteractionClass.BOND: ts}, reg) + matcher = FakeSmartsMatcher() # all empty + assigned = ff.match_molecule(mol=None, matcher=matcher) + assert assigned == {} + + def test_forcefield_module_has_no_molpot_imports(self): + src = FORCEFIELD_PATH.read_text(encoding="utf-8") + tree = ast.parse(src) + for node in ast.walk(tree): + if isinstance(node, ast.Import): + for alias in node.names: + assert not alias.name.startswith("molpot"), alias.name + elif isinstance(node, ast.ImportFrom): + mod = node.module or "" + assert not mod.startswith("molpot"), mod + + def test_perception_package_has_no_molpot_energy_imports(self): + for path in PERCEPTION_ROOT.glob("*.py"): + tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path)) + for node in ast.walk(tree): + if isinstance(node, ast.ImportFrom): + mod = node.module or "" + assert not mod.startswith("molpot"), f"{path.name}: {mod}" + elif isinstance(node, ast.Import): + for alias in node.names: + assert not alias.name.startswith("molpot"), f"{path.name}: {alias.name}" + + def test_type_system_interaction_mismatch_raises(self): + ts = TypeSystem( + InteractionClass.ANGLE, + [TypeRecord(type_id=0, prototype={"k": 50.0, "theta0": 1.9})], + ) + reg = ClassPatternRegistry() + with pytest.raises(ValueError, match="mapped under"): + SymbolicForceField({InteractionClass.BOND: ts}, reg) diff --git a/tests/test_molrep/test_perception/test_matcher.py b/tests/test_molrep/test_perception/test_matcher.py new file mode 100644 index 0000000..a310d69 --- /dev/null +++ b/tests/test_molrep/test_perception/test_matcher.py @@ -0,0 +1,103 @@ +"""SmartsMatcher Protocol + Fake / Molpy matchers (learnable-classical-ff-07).""" + +from __future__ import annotations + +import ast +from pathlib import Path + +import pytest +import torch + +from molrep.perception import ( + FakeSmartsMatcher, + MolpySmartsMatcher, + SmartsMatcher, + SymbolicPattern, +) + +PERCEPTION_ROOT = Path(__file__).resolve().parents[3] / "src" / "molrep" / "perception" + + +class TestFakeSmartsMatcher: + def test_returns_configured_hits_shape(self): + hits = torch.tensor([[0, 1], [1, 2]], dtype=torch.long) # [2, 2] + matcher = FakeSmartsMatcher({"[#6]-[#6]": hits}) + pat = SymbolicPattern("[#6]-[#6]", arity=2) + out = matcher.match(mol=None, pattern=pat) + assert out.shape == (2, 2) + assert torch.equal(out, hits) + assert out.dtype == torch.long + + def test_unknown_pattern_returns_empty(self): + matcher = FakeSmartsMatcher() + pat = SymbolicPattern("[#8]", arity=1) + out = matcher.match(mol=object(), pattern=pat) + assert out.shape == (1, 0) + assert out.dtype == torch.long + + def test_arity_mismatch_raises(self): + matcher = FakeSmartsMatcher({"[#6]-[#6]": torch.tensor([[0], [1], [2]], dtype=torch.long)}) + pat = SymbolicPattern("[#6]-[#6]", arity=2) + with pytest.raises(ValueError, match="arity"): + matcher.match(None, pat) + + def test_satisfies_protocol(self): + matcher = FakeSmartsMatcher() + assert isinstance(matcher, SmartsMatcher) + + +class TestMolpySmartsMatcher: + def test_satisfies_protocol(self): + try: + matcher = MolpySmartsMatcher() + except ImportError: + pytest.skip("molpy.SmartsPattern unavailable") + assert isinstance(matcher, SmartsMatcher) + + def test_import_path_is_molpy_only(self): + """Source of matcher.py must not reference bare molrs imports.""" + src = (PERCEPTION_ROOT / "matcher.py").read_text(encoding="utf-8") + tree = ast.parse(src) + for node in ast.walk(tree): + if isinstance(node, ast.Import): + for alias in node.names: + assert not alias.name.startswith("molrs"), alias.name + elif isinstance(node, ast.ImportFrom): + mod = node.module or "" + assert not mod.startswith("molrs"), mod + assert "from molpy" in src or "import molpy" in src + + def test_functional_match_if_available(self): + """Optional smoke: match carbon atoms in ethanol when molpy works.""" + try: + import molpy as mp + from molpy import SmartsPattern # noqa: F401 + + matcher = MolpySmartsMatcher() + mol = mp.io.read_smiles("CCO") + except Exception as exc: # pragma: no cover - env-dependent + pytest.skip(f"molpy SMARTS smoke unavailable: {exc}") + + pat = SymbolicPattern("[#6]", arity=1) + out = matcher.match(mol, pat) + assert out.ndim == 2 + assert out.shape[0] == 1 + assert out.shape[1] >= 1 # ethanol has ≥1 carbon (heavy-only graph) + assert out.dtype == torch.long + + +class TestPerceptionMolrsForbidden: + def test_no_molrs_imports_in_package(self): + py_files = sorted(PERCEPTION_ROOT.glob("*.py")) + assert py_files, "perception package missing" + for path in py_files: + tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path)) + for node in ast.walk(tree): + if isinstance(node, ast.Import): + for alias in node.names: + assert not alias.name.startswith("molrs"), ( + f"{path.name} imports {alias.name}" + ) + elif isinstance(node, ast.ImportFrom): + mod = node.module or "" + assert not mod.startswith("molrs"), f"{path.name} imports from {mod}" diff --git a/tests/test_molrep/test_perception/test_patterns.py b/tests/test_molrep/test_perception/test_patterns.py new file mode 100644 index 0000000..c86f09c --- /dev/null +++ b/tests/test_molrep/test_perception/test_patterns.py @@ -0,0 +1,93 @@ +"""SymbolicPattern + DiscreteClassRecord validation (learnable-classical-ff-07).""" + +from __future__ import annotations + +import pytest + +from molrep.condensation import InteractionClass +from molrep.perception import DiscreteClassRecord, SymbolicPattern + + +class TestSymbolicPattern: + def test_valid_smarts_arity_2(self): + pat = SymbolicPattern(pattern="[#6]-[#6]", arity=2) + assert pat.pattern == "[#6]-[#6]" + assert pat.arity == 2 + assert pat.kind == "smarts" + + def test_empty_pattern_raises(self): + with pytest.raises(ValueError, match="non-empty"): + SymbolicPattern(pattern="", arity=1) + + def test_whitespace_only_pattern_raises(self): + with pytest.raises(ValueError, match="non-empty"): + SymbolicPattern(pattern=" ", arity=2) + + def test_arity_0_raises(self): + with pytest.raises(ValueError, match="arity"): + SymbolicPattern(pattern="[#6]", arity=0) + + def test_arity_5_raises(self): + with pytest.raises(ValueError, match="arity"): + SymbolicPattern(pattern="[#6]", arity=5) + + def test_valid_arities(self): + for a in (1, 2, 3, 4): + pat = SymbolicPattern(pattern="[#6]", arity=a) + assert pat.arity == a + + def test_smirks_kind(self): + pat = SymbolicPattern( + pattern="[C:1][O:2]>>[C:1][O:2]", + arity=2, + kind="smirks", + ) + assert pat.kind == "smirks" + + def test_atom_maps_frozen_copy(self): + maps = {0: "1", 1: "2"} + pat = SymbolicPattern(pattern="[C:1][O:2]", arity=2, atom_maps=maps) + maps[0] = "99" + assert pat.atom_maps is not None + assert pat.atom_maps[0] == "1" + + +class TestDiscreteClassRecord: + def test_construct_unbound(self): + rec = DiscreteClassRecord( + interaction=InteractionClass.BOND, + type_id=0, + prototype={"k": 300.0, "r0": 1.09}, + ) + assert rec.smarts is None + assert rec.smirks is None + assert rec.prototype["k"] == 300.0 + + def test_construct_with_smarts(self): + rec = DiscreteClassRecord( + interaction=InteractionClass.BOND, + type_id=1, + prototype={"k": 200.0, "r0": 1.5}, + smarts="[#6]-[#8]", + label="C-O", + ) + assert rec.smarts == "[#6]-[#8]" + assert rec.label == "C-O" + + def test_negative_type_id_raises(self): + with pytest.raises(ValueError, match="type_id"): + DiscreteClassRecord( + interaction=InteractionClass.ANGLE, + type_id=-1, + prototype={"k": 50.0, "theta0": 1.9}, + ) + + def test_prototype_is_copied(self): + proto = {"k": 1.0, "r0": 1.0} + rec = DiscreteClassRecord( + interaction=InteractionClass.BOND, + type_id=0, + prototype=proto, + ) + proto["k"] = 999.0 + assert rec.prototype["k"] == 1.0 diff --git a/tests/test_molrep/test_perception/test_registry.py b/tests/test_molrep/test_perception/test_registry.py new file mode 100644 index 0000000..e806bdd --- /dev/null +++ b/tests/test_molrep/test_perception/test_registry.py @@ -0,0 +1,74 @@ +"""ClassPatternRegistry bind/get/conflict (learnable-classical-ff-07).""" + +from __future__ import annotations + +import pytest + +from molrep.condensation import InteractionClass +from molrep.perception import ClassPatternRegistry, SymbolicPattern + + +class TestClassPatternRegistry: + def test_bind_then_get(self): + reg = ClassPatternRegistry() + pat = SymbolicPattern("[#6]-[#6]", arity=2) + reg.bind(InteractionClass.BOND, 0, pat) + got = reg.get(InteractionClass.BOND, 0) + assert got is pat or got == pat + assert got.pattern == "[#6]-[#6]" + + def test_idempotent_same_pattern(self): + reg = ClassPatternRegistry() + pat = SymbolicPattern("[#6]", arity=1) + reg.bind(InteractionClass.LJ, 0, pat) + reg.bind(InteractionClass.LJ, 0, pat) + assert len(reg) == 1 + + def test_conflicting_bind_raises(self): + reg = ClassPatternRegistry() + reg.bind( + InteractionClass.BOND, + 0, + SymbolicPattern("[#6]-[#6]", arity=2), + ) + with pytest.raises(ValueError, match="conflicting"): + reg.bind( + InteractionClass.BOND, + 0, + SymbolicPattern("[#6]-[#8]", arity=2), + ) + + def test_reverse_lookup(self): + reg = ClassPatternRegistry() + pat = SymbolicPattern("[#6]-[#8]", arity=2) + reg.bind(InteractionClass.BOND, 3, pat) + interaction, type_id = reg.reverse_lookup(pat) + assert interaction is InteractionClass.BOND + assert type_id == 3 + interaction2, type_id2 = reg.reverse_lookup("[#6]-[#8]") + assert (interaction2, type_id2) == (InteractionClass.BOND, 3) + + def test_pattern_reuse_across_keys_raises(self): + reg = ClassPatternRegistry() + pat = SymbolicPattern("[#6]-[#6]", arity=2) + reg.bind(InteractionClass.BOND, 0, pat) + with pytest.raises(ValueError, match="already bound"): + reg.bind(InteractionClass.BOND, 1, pat) + + def test_get_missing_raises(self): + reg = ClassPatternRegistry() + with pytest.raises(KeyError, match="no pattern"): + reg.get(InteractionClass.ANGLE, 0) + + def test_get_optional_none(self): + reg = ClassPatternRegistry() + assert reg.get_optional(InteractionClass.ANGLE, 0) is None + + def test_contains(self): + reg = ClassPatternRegistry() + pat = SymbolicPattern("[#7]", arity=1) + reg.bind(InteractionClass.CHARGE, 0, pat) + assert (InteractionClass.CHARGE, 0) in reg + assert "[#7]" in reg + assert pat in reg + assert (InteractionClass.BOND, 0) not in reg diff --git a/tests/test_molrep/test_readout/test_mace.py b/tests/test_molrep/test_readout/test_mace.py new file mode 100644 index 0000000..c6e934e --- /dev/null +++ b/tests/test_molrep/test_readout/test_mace.py @@ -0,0 +1,413 @@ +"""Tests for molrep.readout.mace module.""" + +import cuequivariance_torch as cuet +import pytest +import torch +import torch.nn as nn + +from molix import config +from molrep.interaction.contraction import SymmetricContraction +from molrep.readout.mace import ( + LinearReadout, + NonLinearBiasReadout, + NonLinearReadout, + ProductHead, + ProductHeadSpec, + _ScalarO3Linear, +) +from molrep.readout.projection import BasisProjection + +FEATURES = 8 +MLP_DIM = 4 +N_NODES = 4 + +#: Configuration inherited with the cases migrated from +#: ``tests/test_molzoo/test_mace.py`` (mace-subpackage-restructure-02-core). +#: ``hidden_dim`` is the mixed-l message width ``num_features * (l_max+1)**2`` +#: = 64 * 9 = 576 for num_features=64, l_max=2. +MIGRATED_CONFIG = { + "hidden_dim": 576, + "out_dim": 64, + "num_radial": 8, + "l_max": 2, + "max_body_order": 2, + "num_species": 118, +} +MIGRATED_NUM_FEATURES = MIGRATED_CONFIG["hidden_dim"] // (MIGRATED_CONFIG["l_max"] + 1) ** 2 + +#: Multiplicities giving ``in_mul * out_mul = 16384`` weights — a sample large +#: enough that the N(0, 1) moments of a correct e3nn-style init land inside the +#: hard-coded bands of +#: :meth:`TestScalarO3Linear.test_weight_init_is_standard_normal`. +NORMAL_SAMPLE_MUL = 128 + + +class TestProductHeadSpec: + """Test ProductHeadSpec configuration.""" + + def test_valid_config(self): + """Test creation with valid parameters.""" + spec = ProductHeadSpec( + hidden_dim=64, + out_dim=1, + num_radial=8, + l_max=2, + max_body_order=2, + num_species=10, + ) + assert spec.hidden_dim == 64 + assert spec.out_dim == 1 + assert spec.num_radial == 8 + assert spec.l_max == 2 + assert spec.max_body_order == 2 + assert spec.num_species == 10 + + def test_invalid_hidden_dim(self): + """Test validation for hidden_dim.""" + with pytest.raises(ValueError): + ProductHeadSpec(hidden_dim=0, out_dim=1) + + def test_invalid_out_dim(self): + """Test validation for out_dim.""" + with pytest.raises(ValueError): + ProductHeadSpec(hidden_dim=64, out_dim=0) + + +class TestProductHead: + """Test ProductHead prediction layer.""" + + def test_initialization(self): + """Test ProductHead initialization.""" + # hidden_dim is the mixed-l message dim = num_features * (l_max+1)**2. + # For num_features=8, l_max=2 -> 8 * 9 = 72. + head = ProductHead( + hidden_dim=72, + out_dim=1, + num_radial=8, + l_max=2, + max_body_order=2, + num_species=10, + ) + assert head.config.hidden_dim == 72 + assert head.config.out_dim == 1 + + def test_forward_shape(self): + """Test output shape.""" + head = ProductHead( + hidden_dim=72, + out_dim=1, + num_radial=8, + l_max=2, + max_body_order=2, + num_species=10, + ) + + n_nodes = 20 + node_features = torch.randn(n_nodes, 72, dtype=config.ftype) + atom_types = torch.randint(0, 10, (n_nodes,), dtype=torch.long) + + output = head(node_features, atom_types) + assert output.shape == (n_nodes, 1) + + def test_different_output_dims(self): + """Test with different output dimensions.""" + # num_features=4, l_max=2 (default) -> hidden_dim = 4 * 9 = 36. + for out_dim in [1, 3, 5]: + head = ProductHead( + hidden_dim=36, + out_dim=out_dim, + num_species=5, + ) + + node_features = torch.randn(10, 36, dtype=config.ftype) + atom_types = torch.randint(0, 5, (10,), dtype=torch.long) + + output = head(node_features, atom_types) + assert output.shape == (10, out_dim) + + def test_differentiable(self): + """Test that gradients flow through head.""" + head = ProductHead( + hidden_dim=36, + out_dim=1, + num_species=5, + ) + + node_features = torch.randn(10, 36, requires_grad=True, dtype=config.ftype) + atom_types = torch.randint(0, 5, (10,), dtype=torch.long) + + output = head(node_features, atom_types) + loss = output.sum() + loss.backward() + + assert node_features.grad is not None + assert not torch.isnan(node_features.grad).any() + + # -- migrated from tests/test_molzoo/test_mace.py by + # mace-subpackage-restructure-02-core (module/package name clash forced + # that file's removal; these cases have no equivalent above) ---------- + + @pytest.fixture + def migrated_head(self): + """A ProductHead on the wider migrated configuration.""" + return ProductHead(**MIGRATED_CONFIG) + + def test_initialization_wires_components(self, migrated_head): + """Every sub-block is present and of the expected type.""" + assert isinstance(migrated_head.symmetric_contraction, SymmetricContraction) + assert isinstance(migrated_head.basis_projection, BasisProjection) + assert isinstance(migrated_head.linear, nn.Linear) + # The contraction emits invariant scalars (num_features), so the readout + # linear maps num_features -> out_dim (not the full mixed-l hidden_dim). + assert migrated_head.linear.in_features == MIGRATED_NUM_FEATURES + assert migrated_head.linear.out_features == MIGRATED_CONFIG["out_dim"] + + def test_symmetric_contraction_config(self, migrated_head): + """Constructor kwargs reach the SymmetricContraction spec.""" + sc = migrated_head.symmetric_contraction + assert sc.config.hidden_dim == MIGRATED_CONFIG["hidden_dim"] + assert sc.config.num_species == MIGRATED_CONFIG["num_species"] + assert sc.config.max_body_order == MIGRATED_CONFIG["max_body_order"] + + def test_basis_projection_config(self, migrated_head): + """Constructor kwargs reach the BasisProjection spec.""" + bp = migrated_head.basis_projection + assert bp.config.hidden_dim == MIGRATED_CONFIG["hidden_dim"] + assert bp.config.num_radial == MIGRATED_CONFIG["num_radial"] + assert bp.config.l_max == MIGRATED_CONFIG["l_max"] + assert bp.config.max_body_order == MIGRATED_CONFIG["max_body_order"] + + def test_forward_preserves_dtype(self, migrated_head): + """The head does not silently up/down-cast the working precision.""" + node_features = torch.randn(10, MIGRATED_CONFIG["hidden_dim"], dtype=config.ftype) + atom_types = torch.randint(0, MIGRATED_CONFIG["num_species"], (10,), dtype=torch.long) + + assert migrated_head(node_features, atom_types).dtype == config.ftype + + @pytest.mark.parametrize("max_body_order", [1, 2, 3]) + def test_different_max_body_orders(self, max_body_order): + """Body order propagates to both the contraction and the projection.""" + head = ProductHead(**{**MIGRATED_CONFIG, "max_body_order": max_body_order}) + assert head.symmetric_contraction.config.max_body_order == max_body_order + assert head.basis_projection.config.max_body_order == max_body_order + + @pytest.mark.parametrize("num_species", [10, 50, 118]) + def test_different_num_species(self, num_species): + """The element table size propagates to the contraction.""" + head = ProductHead(**{**MIGRATED_CONFIG, "num_species": num_species}) + assert head.symmetric_contraction.config.num_species == num_species + + def test_symmetric_contraction_component(self, migrated_head): + """The contraction maps the mixed-l message to invariant scalars.""" + n_nodes = 15 + node_features = torch.randn(n_nodes, MIGRATED_CONFIG["hidden_dim"], dtype=config.ftype) + atom_types = torch.randint(0, MIGRATED_CONFIG["num_species"], (n_nodes,), dtype=torch.long) + + basis = migrated_head.symmetric_contraction(node_features, atom_types) + assert basis.shape == (n_nodes, MIGRATED_NUM_FEATURES) + + def test_basis_projection_component(self, migrated_head): + """``BasisProjection`` is an identity in the current implementation.""" + basis = torch.randn(15, MIGRATED_CONFIG["hidden_dim"], dtype=config.ftype) + assert torch.equal(migrated_head.basis_projection(basis), basis) + + def test_linear_component(self, migrated_head): + """The readout linear consumes the contracted scalars.""" + features = torch.randn(15, MIGRATED_NUM_FEATURES, dtype=config.ftype) + assert migrated_head.linear(features).shape == (15, MIGRATED_CONFIG["out_dim"]) + + def test_cuequivariance_integration(self, migrated_head): + """The contraction is a cuEq ``SymmetricContraction``, not a port.""" + cue_sc = migrated_head.symmetric_contraction.symmetric_contraction + assert isinstance(cue_sc, cuet.SymmetricContraction) + assert hasattr(cue_sc, "contraction_degree") + assert hasattr(cue_sc, "num_elements") + + +class TestProductHeadEquivariance: + """Test equivariance properties of ProductHead. + + Migrated from ``tests/test_molzoo/test_mace.py`` by + mace-subpackage-restructure-02-core. + """ + + @pytest.fixture + def product_head(self): + """A small ProductHead (num_features=32, l_max=1 -> hidden_dim=128).""" + return ProductHead( + hidden_dim=128, + out_dim=32, + num_radial=8, + l_max=1, + max_body_order=2, + num_species=10, + ) + + def test_permutation_equivariance(self, product_head): + """Relabelling atoms permutes the per-atom output the same way.""" + n_nodes = 15 + torch.manual_seed(0) + node_features = torch.randn(n_nodes, 128, dtype=config.ftype) + atom_types = torch.randint(0, 10, (n_nodes,), dtype=torch.long) + + output1 = product_head(node_features, atom_types) + + perm = torch.randperm(n_nodes) + output2 = product_head(node_features[perm], atom_types[perm]) + + assert torch.allclose(output1[perm], output2, rtol=1e-5, atol=1e-5) + + +@pytest.fixture +def scalar_feats(): + """Per-atom scalar node features ``(N, FEATURES)`` matching ``FEATURES x0e``.""" + torch.manual_seed(0) + return torch.randn(N_NODES, FEATURES, dtype=config.ftype) + + +class TestScalarO3Linear: + """Test the scalar-only ``o3.Linear`` port behind :class:`NonLinearBiasReadout`. + + Initialisation follows e3nn's ``o3.Linear`` convention: the weight is drawn + from the **global** RNG as standard normal ``N(0, 1)`` and the path + normalisation ``1/sqrt(in_mul)`` is applied in ``forward`` (``self._alpha``), + not folded into the init; the bias stays zero. + """ + + def test_weight_is_not_zero_initialised(self): + """A fresh layer must not start from an all-zero weight. + + Zero weights make every downstream assertion on an untrained model + vacuous — the readout emits a constant per-atom energy, so the forces + are identically zero. + """ + torch.manual_seed(0) + layer = _ScalarO3Linear(FEATURES, MLP_DIM) + + assert not torch.equal(layer.weight, torch.zeros_like(layer.weight)) + + def test_bias_is_zero_initialised(self): + """The bias stays at zero — e3nn's ``o3.Linear`` convention.""" + torch.manual_seed(0) + layer = _ScalarO3Linear(FEATURES, MLP_DIM) + + assert torch.equal(layer.bias, torch.zeros_like(layer.bias)) + + def test_weight_init_is_reproducible_under_a_fixed_seed(self): + """Two constructions under the same seed give bit-identical weights.""" + torch.manual_seed(0) + first = _ScalarO3Linear(FEATURES, MLP_DIM) + torch.manual_seed(0) + second = _ScalarO3Linear(FEATURES, MLP_DIM) + + assert torch.equal(first.weight, second.weight) + + def test_weight_init_differs_across_seeds(self): + """The weight is drawn from the global RNG, so a new seed changes it.""" + torch.manual_seed(0) + first = _ScalarO3Linear(FEATURES, MLP_DIM) + torch.manual_seed(1) + second = _ScalarO3Linear(FEATURES, MLP_DIM) + + assert not torch.equal(first.weight, second.weight) + + def test_weight_init_is_standard_normal(self): + """Sample moments of the init sit in the N(0, 1) bands at n = 16384. + + The ``1/sqrt(in_mul)`` path normalisation is a ``forward`` factor + (``_alpha``), so the stored weight itself is unit-variance. + """ + torch.manual_seed(0) + layer = _ScalarO3Linear(NORMAL_SAMPLE_MUL, NORMAL_SAMPLE_MUL) + + assert layer.weight.numel() == 16384 + assert abs(layer.weight.mean().item()) < 0.1 + assert 0.9 < layer.weight.std().item() < 1.1 + + def test_state_dict_load_overwrites_the_random_init(self): + """Loading a checkpoint replaces the init exactly — no blending. + + Guards that giving the layer a random init cannot perturb any + loaded-weights path (the foundation-model checkpoints). + """ + torch.manual_seed(0) + source = _ScalarO3Linear(FEATURES, MLP_DIM) + torch.manual_seed(1) + target = _ScalarO3Linear(FEATURES, MLP_DIM) + assert not torch.equal(source.weight, target.weight) + + target.load_state_dict(source.state_dict()) + + assert torch.equal(target.weight, source.weight) + assert torch.equal(target.bias, source.bias) + + +class TestLinearReadout: + """Test the MACE ``LinearReadoutBlock`` port.""" + + def test_forward_shape(self, scalar_feats): + """Projects scalar node features to one number per atom.""" + head = LinearReadout(irreps_in=f"{FEATURES}x0e") + assert head(scalar_feats).shape == (N_NODES, 1) + + def test_state_dict_keys_are_the_weight_transfer_contract(self): + """Sub-layer name ``linear`` must stay put — official weights load by key.""" + head = LinearReadout(irreps_in=f"{FEATURES}x0e") + assert sorted(head.state_dict().keys()) == [ + "linear.f.m.graphs.0.graph.c0", + "linear.weight", + ] + + +class TestNonLinearReadout: + """Test the bias-free gated readout (MACE-MP / MatPES last layer).""" + + def test_forward_shape(self, scalar_feats): + """Two equivariant linears with a SiLU gate give one number per atom.""" + head = NonLinearReadout(irreps_in=f"{FEATURES}x0e", mlp_dim=MLP_DIM) + assert head(scalar_feats).shape == (N_NODES, 1) + + def test_state_dict_keys_are_the_weight_transfer_contract(self): + """``linear_1`` / ``linear_2`` and no bias entries — MACE's bias-free variant.""" + head = NonLinearReadout(irreps_in=f"{FEATURES}x0e", mlp_dim=MLP_DIM) + assert sorted(head.state_dict().keys()) == [ + "linear_1.f.m.graphs.0.graph.c0", + "linear_1.weight", + "linear_2.f.m.graphs.0.graph.c0", + "linear_2.weight", + ] + + +class TestNonLinearBiasReadout: + """Test the biased three-layer readout (MACE-OMOL variant).""" + + def test_forward_shape(self, scalar_feats): + """``Linear → SiLU → linear_mid → SiLU → linear_2`` gives one number per atom.""" + head = NonLinearBiasReadout(irreps_in=f"{FEATURES}x0e", mlp_dim=MLP_DIM) + assert head(scalar_feats).shape == (N_NODES, 1) + + def test_untrained_output_is_not_constant_across_atoms(self, scalar_feats): + """An *untrained* readout must still resolve two different atoms. + + A constant output here means the model's interaction energy is + position-independent and its forces are identically zero, which + silently voids every force / parity assertion built on a fresh model. + """ + torch.manual_seed(0) + head = NonLinearBiasReadout(irreps_in=f"{FEATURES}x0e", mlp_dim=MLP_DIM) + + energies = head(scalar_feats) + + assert not torch.allclose(energies[0], energies[1]) + + def test_state_dict_keys_are_the_weight_transfer_contract(self): + """``linear_mid`` plus the two biases distinguish this from NonLinearReadout.""" + head = NonLinearBiasReadout(irreps_in=f"{FEATURES}x0e", mlp_dim=MLP_DIM) + assert sorted(head.state_dict().keys()) == [ + "linear_1.f.m.graphs.0.graph.c0", + "linear_1.weight", + "linear_2.bias", + "linear_2.weight", + "linear_mid.bias", + "linear_mid.weight", + ] diff --git a/tests/test_molrep/test_readout/test_pooling.py b/tests/test_molrep/test_readout/test_pooling.py index 703c605..dbce4a9 100644 --- a/tests/test_molrep/test_readout/test_pooling.py +++ b/tests/test_molrep/test_readout/test_pooling.py @@ -111,10 +111,8 @@ def test_different_feature_dims(self): assert graph_features.shape[1] == feat_dim def test_empty_graph(self): - """Test behavior with graph that might have no atoms.""" - pooling = ScatterPooling(strategy="sum") - - # Graph 0: 5 atoms, Graph 1: 0 atoms (skip), Graph 2: 3 atoms + """A graph with no atoms pools to an all-zero row, not an error.""" + # Graph 0: 5 atoms, Graph 1: 0 atoms, Graph 2: 3 atoms. node_features = torch.randn(8, 16) batch = torch.cat( [ @@ -123,13 +121,18 @@ def test_empty_graph(self): ] ) - try: - graph_features = pooling(node_features, batch) - # Implementation might handle this differently - assert graph_features.shape[1] == 16 - except Exception: - # Empty graphs might not be supported - pytest.skip("Empty graphs not supported in this implementation") + summed = ScatterPooling(strategy="sum")(node_features, batch) + assert summed.shape == (3, 16) + assert torch.allclose(summed[0], node_features[:5].sum(dim=0)) + assert torch.equal(summed[1], torch.zeros(16)) + assert torch.allclose(summed[2], node_features[5:].sum(dim=0)) + + # Mean divides by a clamped count, so the empty graph stays finite. + averaged = ScatterPooling(strategy="mean")(node_features, batch) + assert torch.allclose(averaged[0], node_features[:5].mean(dim=0)) + assert torch.equal(averaged[1], torch.zeros(16)) + assert torch.allclose(averaged[2], node_features[5:].mean(dim=0)) + assert torch.isfinite(averaged).all() def test_differentiable(self): """Test that gradients flow through pooling.""" diff --git a/tests/test_molrep/test_readout/test_product.py b/tests/test_molrep/test_readout/test_product.py deleted file mode 100644 index d6892e9..0000000 --- a/tests/test_molrep/test_readout/test_product.py +++ /dev/null @@ -1,109 +0,0 @@ -"""Tests for molrep.readout.product module.""" - -import pytest -import torch - -from molix import config -from molrep.readout.product import ProductHead, ProductHeadSpec - - -class TestProductHeadSpec: - """Test ProductHeadSpec configuration.""" - - def test_valid_config(self): - """Test creation with valid parameters.""" - spec = ProductHeadSpec( - hidden_dim=64, - out_dim=1, - num_radial=8, - l_max=2, - max_body_order=2, - num_species=10, - ) - assert spec.hidden_dim == 64 - assert spec.out_dim == 1 - assert spec.num_radial == 8 - assert spec.l_max == 2 - assert spec.max_body_order == 2 - assert spec.num_species == 10 - - def test_invalid_hidden_dim(self): - """Test validation for hidden_dim.""" - with pytest.raises(ValueError): - ProductHeadSpec(hidden_dim=0, out_dim=1) - - def test_invalid_out_dim(self): - """Test validation for out_dim.""" - with pytest.raises(ValueError): - ProductHeadSpec(hidden_dim=64, out_dim=0) - - -class TestProductHead: - """Test ProductHead prediction layer.""" - - def test_initialization(self): - """Test ProductHead initialization.""" - # hidden_dim is the mixed-l message dim = num_features * (l_max+1)**2. - # For num_features=8, l_max=2 -> 8 * 9 = 72. - head = ProductHead( - hidden_dim=72, - out_dim=1, - num_radial=8, - l_max=2, - max_body_order=2, - num_species=10, - ) - assert head.config.hidden_dim == 72 - assert head.config.out_dim == 1 - - def test_forward_shape(self): - """Test output shape.""" - head = ProductHead( - hidden_dim=72, - out_dim=1, - num_radial=8, - l_max=2, - max_body_order=2, - num_species=10, - ) - - n_nodes = 20 - node_features = torch.randn(n_nodes, 72, dtype=config.ftype) - atom_types = torch.randint(0, 10, (n_nodes,), dtype=torch.long) - - output = head(node_features, atom_types) - assert output.shape == (n_nodes, 1) - - def test_different_output_dims(self): - """Test with different output dimensions.""" - # num_features=4, l_max=2 (default) -> hidden_dim = 4 * 9 = 36. - for out_dim in [1, 3, 5]: - head = ProductHead( - hidden_dim=36, - out_dim=out_dim, - num_species=5, - ) - - node_features = torch.randn(10, 36, dtype=config.ftype) - atom_types = torch.randint(0, 5, (10,), dtype=torch.long) - - output = head(node_features, atom_types) - assert output.shape == (10, out_dim) - - def test_differentiable(self): - """Test that gradients flow through head.""" - head = ProductHead( - hidden_dim=36, - out_dim=1, - num_species=5, - ) - - node_features = torch.randn(10, 36, requires_grad=True, dtype=config.ftype) - atom_types = torch.randint(0, 5, (10,), dtype=torch.long) - - output = head(node_features, atom_types) - loss = output.sum() - loss.backward() - - assert node_features.grad is not None - assert not torch.isnan(node_features.grad).any() diff --git a/tests/test_molrep/test_reexport_compat.py b/tests/test_molrep/test_reexport_compat.py new file mode 100644 index 0000000..f0bed13 --- /dev/null +++ b/tests/test_molrep/test_reexport_compat.py @@ -0,0 +1,140 @@ +"""Back-compat guard for the mace-subpackage-restructure-01 re-export shims. + +Every legacy import path must resolve to the *same object* as its new +``mace`` home (not a re-implementation), and the two package ``__all__`` +lists must stay byte-identical to the pre-move lists so downstream +``from molrep.interaction import X`` keeps working. + +Scope is the degraded chain step: ``residual`` / ``product_basis`` / +``element`` are deliberately NOT moved here (deferred to -01b), so they +are absent from these assertions on purpose. +""" + +import importlib + +import molrep.interaction +import molrep.interaction.density +import molrep.interaction.mace.conv +import molrep.interaction.mace.density +import molrep.interaction.product +import molrep.readout +import molrep.readout.mace +import molrep.readout.product +import molrep.readout.scalar + +# Pre-move lists, copied verbatim from src/molrep/{interaction,readout}/__init__.py +# at parent commit cf60f99. The move must not add, drop, or reorder a name. +EXPECTED_INTERACTION_ALL = [ + "DensityInteraction", + "DensityResidualInteraction", + "GatedNonlinearity", + "ResidualInteraction", + "EquivariantProductBasis", + "RadialMLP", + "MessageAggregation", + "MessageAggregationSpec", + "EquivariantLinear", + "ConvTP", + "ConvTPSpec", + "irreps_from_l_max", + "sh_irreps_from_l_max", + "SymmetricContraction", + "SymmetricContractionSpec", + "ElementUpdate", + "ElementUpdateSpec", + "RadialWeightMLP", + "RadialWeightMLPSpec", +] + +EXPECTED_READOUT_ALL = [ + "masked_sum_pooling", + "masked_mean_pooling", + "BasisProjection", + "BasisProjectionSpec", + "ProductHead", + "ProductHeadSpec", +] + + +class TestInteractionDensityShim: + """``molrep.interaction.density`` re-exports ``molrep.interaction.mace.density``.""" + + def test_density_interaction_is_same_object(self): + """Legacy ``DensityInteraction`` must be the moved class itself.""" + assert ( + molrep.interaction.density.DensityInteraction + is molrep.interaction.mace.density.DensityInteraction + ) + + def test_density_residual_interaction_is_same_object(self): + """Legacy ``DensityResidualInteraction`` must be the moved class itself.""" + assert ( + molrep.interaction.density.DensityResidualInteraction + is molrep.interaction.mace.density.DensityResidualInteraction + ) + + def test_skip_tp_method_constant_is_same_value(self): + """The ``skip_tp`` method constant travels with the module.""" + assert ( + molrep.interaction.density.SKIP_TP_METHOD + is molrep.interaction.mace.density.SKIP_TP_METHOD + ) + + +class TestInteractionConvShim: + """``molrep.interaction.product`` re-exports the split-out ConvTP pair.""" + + def test_conv_tp_is_same_object(self): + """Legacy ``ConvTP`` must be the class now living in ``mace.conv``.""" + assert molrep.interaction.product.ConvTP is molrep.interaction.mace.conv.ConvTP + + def test_conv_tp_spec_is_same_object(self): + """Legacy ``ConvTPSpec`` must be the class now living in ``mace.conv``.""" + assert molrep.interaction.product.ConvTPSpec is molrep.interaction.mace.conv.ConvTPSpec + + +class TestReadoutScalarShim: + """``molrep.readout.scalar`` re-exports ``molrep.readout.mace``.""" + + def test_linear_readout_is_same_object(self): + """Legacy ``LinearReadout`` must be the moved class itself.""" + assert molrep.readout.scalar.LinearReadout is molrep.readout.mace.LinearReadout + + def test_non_linear_readout_is_same_object(self): + """Legacy ``NonLinearReadout`` must be the moved class itself.""" + assert molrep.readout.scalar.NonLinearReadout is molrep.readout.mace.NonLinearReadout + + def test_non_linear_bias_readout_is_same_object(self): + """Legacy ``NonLinearBiasReadout`` must be the moved class itself.""" + assert ( + molrep.readout.scalar.NonLinearBiasReadout is molrep.readout.mace.NonLinearBiasReadout + ) + + +class TestReadoutProductShim: + """``molrep.readout.product`` re-exports ``molrep.readout.mace``.""" + + def test_product_head_is_same_object(self): + """Legacy ``ProductHead`` must be the moved class itself.""" + assert molrep.readout.product.ProductHead is molrep.readout.mace.ProductHead + + def test_product_head_spec_is_same_object(self): + """Legacy ``ProductHeadSpec`` must be the moved class itself.""" + assert molrep.readout.product.ProductHeadSpec is molrep.readout.mace.ProductHeadSpec + + +class TestPackageExports: + """Package-level ``__all__`` lists are frozen across the move.""" + + def test_interaction_all_is_unchanged(self): + """``molrep.interaction.__all__`` must match the pre-move list verbatim.""" + assert molrep.interaction.__all__ == EXPECTED_INTERACTION_ALL + + def test_readout_all_is_unchanged(self): + """``molrep.readout.__all__`` must match the pre-move list verbatim.""" + assert molrep.readout.__all__ == EXPECTED_READOUT_ALL + + def test_legacy_packages_import_clean(self): + """A cold import of every touched package raises no ImportError.""" + for name in ("molrep", "molrep.interaction", "molrep.readout"): + assert importlib.import_module(name) is not None diff --git a/tests/test_molzoo/test_chem/__init__.py b/tests/test_molzoo/test_chem/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/test_molzoo/test_chem/test_encoder.py b/tests/test_molzoo/test_chem/test_encoder.py new file mode 100644 index 0000000..e64ddb7 --- /dev/null +++ b/tests/test_molzoo/test_chem/test_encoder.py @@ -0,0 +1,102 @@ +"""Tests for molzoo.chem ChemPerception recipe.""" + +from __future__ import annotations + +import ast +from pathlib import Path + +import torch +from tensordict import TensorDict + +from molzoo.chem import ChemPerception, ChemPerceptionSpec +from molzoo.chem.encoder import ChemPerception as ChemPerceptionDirect + + +class TestChemPerceptionSpec: + def test_defaults(self): + spec = ChemPerceptionSpec() + assert spec.atom_dim > 0 + assert spec.bond_dim > 0 + assert spec.num_elements > 0 + + +class TestChemPerception: + def test_constructs_and_forwards(self): + model = ChemPerception( + atom_dim=8, + bond_dim=8, + angle_dim=8, + proper_dim=8, + improper_dim=8, + num_elements=20, + ) + batch = TensorDict( + { + "atoms": TensorDict( + { + "Z": torch.tensor([8, 1, 1], dtype=torch.long), + "batch": torch.zeros(3, dtype=torch.long), + }, + batch_size=[3], + ), + "bonds": TensorDict( + { + "atomi": torch.tensor([0, 0], dtype=torch.long), + "atomj": torch.tensor([1, 2], dtype=torch.long), + }, + batch_size=[2], + ), + "angles": TensorDict( + { + "atomi": torch.tensor([1], dtype=torch.long), + "atomj": torch.tensor([0], dtype=torch.long), + "atomk": torch.tensor([2], dtype=torch.long), + }, + batch_size=[1], + ), + }, + batch_size=[], + ) + out = model(batch) + assert out["atoms", "chem_features"].shape == (3, 8) + assert out["bonds", "chem_features"].shape == (2, 8) + assert out["angles", "chem_features"].shape == (1, 8) + # No energy keys from a pure perception recipe + nested = {str(k) for k in out.keys(include_nested=True)} + assert not any("energy" in k for k in nested) + + def test_from_spec(self): + spec = ChemPerceptionSpec( + atom_dim=4, + bond_dim=4, + angle_dim=4, + proper_dim=4, + improper_dim=4, + num_elements=10, + ) + model = ChemPerception(spec=spec) + assert model.config.atom_dim == 4 + assert isinstance(model, ChemPerceptionDirect) + + def test_lazy_export_from_molzoo(self): + import molzoo + + assert "ChemPerception" in molzoo.__all__ + assert "ChemPerceptionSpec" in molzoo.__all__ + # Attribute access resolves via lazy table + assert molzoo.ChemPerception is ChemPerception + assert molzoo.ChemPerceptionSpec is ChemPerceptionSpec + + +class TestNoMolpotImport: + def test_source_tree_has_no_molpot(self): + root = Path(__file__).resolve().parents[3] / "src" / "molzoo" / "chem" + assert root.is_dir() + for path in root.rglob("*.py"): + tree = ast.parse(path.read_text()) + for node in ast.walk(tree): + if isinstance(node, ast.Import): + for alias in node.names: + assert not alias.name.startswith("molpot"), path + elif isinstance(node, ast.ImportFrom) and node.module: + assert not node.module.startswith("molpot"), path diff --git a/tests/test_molzoo/test_imports.py b/tests/test_molzoo/test_imports.py new file mode 100644 index 0000000..eb9eb98 --- /dev/null +++ b/tests/test_molzoo/test_imports.py @@ -0,0 +1,200 @@ +"""Tests for the ``molzoo`` public import surface (``src/molzoo/__init__.py``). + +Two contracts live in that file and nowhere else: + +* the **names** the rest of the world imports from the top level — + ``scripts/matpes_port/run_nve.py:45``, ``benchmarks/bench_mace_matpes.py:38`` + and ``benchmarks/bench_trainer_throughput.py:26`` bind them by hand, so the + ``mace`` sub-package cutover must not move a single one; +* the **lazy policy**: no model symbol is imported at module level, so + ``import molzoo`` costs nothing beyond ``typing`` and the cuEquivariance + stack only appears once a model symbol is actually touched + (``.claude/specs/mace-subpackage-restructure-06-wire.md`` §Design 2). + +The lazy assertions run in a **subprocess**. This pytest session imported +cuEquivariance long before this module was collected (any molzoo model test +does), so an in-process ``sys.modules`` assertion is vacuously green. + +:class:`TestMolzooMaceReexports` additionally pins the *identity* of the +``molzoo.mace`` re-export surface against its promoted ``molrep`` home. It was +moved here from ``tests/test_molrep/test_reexport_compat.py`` by +``mace-subpackage-restructure-06-wire`` ac-006 (no ``molzoo`` import may appear +under ``tests/test_molrep/``); the assertions are unchanged. +""" + +from __future__ import annotations + +import json +import os +import subprocess +import sys +from pathlib import Path + +import pytest + +import molrep.embedding.mace +import molrep.interaction.mace.block +import molzoo +import molzoo.mace + +#: Repo root — ``tests/test_molzoo/test_imports.py`` → ``tests/test_molzoo`` → +#: ``tests`` → root. Used to put ``src/`` on the probe interpreter's path +#: whether or not the package happens to be installed. +_REPO_ROOT = Path(__file__).resolve().parents[2] + +#: The import package root, prepended to ``PYTHONPATH`` for the probe. +_SRC = _REPO_ROOT / "src" + +#: Third-party module whose presence in ``sys.modules`` marks "the equivariance +#: stack got imported". It is the expensive one the lazy policy exists to defer. +_CUEQ = "cuequivariance_torch" + +#: Seconds the probe interpreter may take. Importing the cuEquivariance stack +#: is the slow half and takes ~10 s cold; the budget is deliberately loose +#: because a *hang* (not a slow import) is what this guards against. +_PROBE_TIMEOUT = 600.0 + +#: Probe script: a clean interpreter reports whether :data:`_CUEQ` is loaded +#: right after ``import molzoo`` and again after touching a model symbol. +#: It writes JSON to ``argv[1]`` rather than stdout because importing this +#: stack prints to stdout ("opt_einsum_fx not available."), which would have to +#: be parsed around. +_PROBE = """ +import json +import sys +from pathlib import Path + +import molzoo + +report = {"at_import": "cuequivariance_torch" in sys.modules} +molzoo.MACE +report["after_attribute_access"] = "cuequivariance_torch" in sys.modules +Path(sys.argv[1]).write_text(json.dumps(report)) +""" + + +@pytest.fixture(scope="module") +def lazy_probe(tmp_path_factory: pytest.TempPathFactory) -> dict[str, bool]: + """Run :data:`_PROBE` once in a clean interpreter and return its report. + + Args: + tmp_path_factory: pytest's session-scoped temporary directory factory — + the probe's only filesystem contact. + + Returns: + ``{"at_import": bool, "after_attribute_access": bool}``: whether + :data:`_CUEQ` was in ``sys.modules`` after ``import molzoo`` and after + ``molzoo.MACE``. + """ + report_path = tmp_path_factory.mktemp("molzoo_lazy") / "probe.json" + environment = dict(os.environ) + environment["PYTHONPATH"] = os.pathsep.join( + part for part in (str(_SRC), environment.get("PYTHONPATH", "")) if part + ) + completed = subprocess.run( + [sys.executable, "-c", _PROBE, str(report_path)], + capture_output=True, + text=True, + env=environment, + cwd=str(_REPO_ROOT), + timeout=_PROBE_TIMEOUT, + check=False, + ) + assert completed.returncode == 0, ( + f"probe interpreter failed ({completed.returncode}):\n{completed.stderr}" + ) + report: dict[str, bool] = json.loads(report_path.read_text()) + return report + + +class TestPublicImportSurface: + """The names and the import cost of the ``molzoo`` public surface. + + Both levels of it: ``molzoo/__init__.py`` and the ``molzoo.mace`` + re-export package, which runs the same PEP 562 policy. + """ + + def test_top_level_exports_the_symbols_the_scripts_bind(self) -> None: + """The six names ``run_nve.py`` / ``bench_mace_matpes.py`` import.""" + from molzoo import ( + MACE, + MACEMatpes, + MACEOMol, + MACESpec, + load_matpes_state_dict, + load_omol_state_dict, + ) + + resolved = ( + MACE, + MACESpec, + MACEMatpes, + MACEOMol, + load_matpes_state_dict, + load_omol_state_dict, + ) + assert all(symbol is not None for symbol in resolved) + + def test_mace_sub_package_re_exports_the_research_encoder(self) -> None: + """``bench_trainer_throughput.py:26`` imports through the package.""" + from molzoo.mace import MACE + + assert isinstance(MACE, type) + + def test_importing_molzoo_does_not_import_cuequivariance( + self, lazy_probe: dict[str, bool] + ) -> None: + """``import molzoo`` must not pay for the equivariance stack.""" + assert lazy_probe["at_import"] is False + + def test_touching_a_model_symbol_imports_cuequivariance( + self, lazy_probe: dict[str, bool] + ) -> None: + """Lazy, not absent: the symbol still resolves to the real class.""" + assert lazy_probe["after_attribute_access"] is True + + def test_unknown_attribute_raises_attribute_error(self) -> None: + """``__getattr__`` must not turn a typo into an import error.""" + with pytest.raises(AttributeError): + molzoo.NotAThing + + def test_dir_reports_exactly_the_public_surface(self) -> None: + """``__dir__`` keeps completion and ``from molzoo import *`` alive.""" + assert set(dir(molzoo)) == set(molzoo.__all__) + + def test_mace_subpackage_dir_covers_its_public_surface(self) -> None: + """``molzoo.mace.__dir__`` reports exactly the sub-package's ``__all__``. + + The mirror of :meth:`test_dir_reports_exactly_the_public_surface` one + level down, and the ac-001 lock. ``molzoo.mace`` imports only its + torch-free config models eagerly, so without ``__dir__`` sixteen of its + eighteen exported names would be absent from ``dir()`` — completion + would hide every lazy symbol while ``from molzoo.mace import *`` still + bound it. + + ac-001's gate is the subset form ``set(__all__) <= set(dir())``; the + equality asserted here is what ``src/molzoo/mace/__init__.py``'s + ``__dir__`` actually guarantees and is strictly stronger, since it also + catches a ``dir()`` that grew a name ``__all__`` does not export. + """ + assert set(dir(molzoo.mace)) == set(molzoo.mace.__all__) + + +class TestMolzooMaceReexports: + """``molzoo.mace`` re-exports the promoted embedding / interaction blocks.""" + + def test_embedding_block_is_same_object(self) -> None: + """``molzoo.mace.EmbeddingBlock`` must be the promoted molrep class.""" + assert molzoo.mace.EmbeddingBlock is molrep.embedding.mace.EmbeddingBlock + + def test_embedding_spec_is_same_object(self) -> None: + """``molzoo.mace.EmbeddingSpec`` must be the promoted molrep class.""" + assert molzoo.mace.EmbeddingSpec is molrep.embedding.mace.EmbeddingSpec + + def test_interaction_block_is_same_object(self) -> None: + """``molzoo.mace.InteractionBlock`` must be the promoted molrep class.""" + assert molzoo.mace.InteractionBlock is molrep.interaction.mace.block.InteractionBlock + + def test_interaction_spec_is_same_object(self) -> None: + """``molzoo.mace.InteractionSpec`` must be the promoted molrep class.""" + assert molzoo.mace.InteractionSpec is molrep.interaction.mace.block.InteractionSpec diff --git a/tests/test_molzoo/test_mace.py b/tests/test_molzoo/test_mace.py deleted file mode 100644 index 96ab3ae..0000000 --- a/tests/test_molzoo/test_mace.py +++ /dev/null @@ -1,792 +0,0 @@ -"""Tests for molzoo.mace module.""" - -import math - -import cuequivariance_torch as cuet -import pytest -import torch -import torch.nn as nn - -from molrep.embedding.angular import SphericalHarmonics -from molrep.embedding.cutoff import CosineCutoff -from molrep.embedding.node import DiscreteEmbeddingSpec, JointEmbedding -from molrep.embedding.radial import BesselRBF -from molrep.interaction.contraction import SymmetricContraction -from molrep.interaction.product import ConvTP -from molrep.interaction.radial import RadialWeightMLP -from molrep.readout.product import ProductHead -from molrep.readout.projection import BasisProjection -from molrep.utils.equivariance import ( - random_rotation_matrix, - rotate_vectors, - rotation_matrix_z, -) -from molzoo.mace import ( - EmbeddingBlock, - InteractionBlock, -) - - -class TestEmbeddingBlock: - """Test EmbeddingBlock initialization and forward pass.""" - - @pytest.fixture - def embedding_config(self): - """Common configuration for embedding tests.""" - return { - "num_species": 5, - "num_features": 16, - "r_max": 5.0, - "num_bessel": 8, - "l_max": 2, - } - - @pytest.fixture - def node_attr_specs(self, embedding_config): - """Node attribute specifications.""" - return [ - DiscreteEmbeddingSpec( - input_key="Z", - num_classes=embedding_config["num_species"], - emb_dim=embedding_config["num_features"], - ) - ] - - @pytest.fixture - def embedding_block(self, node_attr_specs, embedding_config): - """Create an EmbeddingBlock instance.""" - return EmbeddingBlock( - node_attr_specs=node_attr_specs, - num_features=embedding_config["num_features"], - r_max=embedding_config["r_max"], - num_bessel=embedding_config["num_bessel"], - l_max=embedding_config["l_max"], - ) - - def test_initialization(self, embedding_block, embedding_config): - """Test that EmbeddingBlock initializes all components correctly.""" - # Check node_embedding - assert hasattr(embedding_block, "node_embedding") - assert isinstance(embedding_block.node_embedding, JointEmbedding) - - # Check radial_embedding - assert hasattr(embedding_block, "radial_embedding") - assert isinstance(embedding_block.radial_embedding, BesselRBF) - assert embedding_block.radial_embedding.config.r_cut == embedding_config["r_max"] - assert embedding_block.radial_embedding.config.num_radial == embedding_config["num_bessel"] - - # Check spherical_harmonics - assert hasattr(embedding_block, "spherical_harmonics") - assert isinstance(embedding_block.spherical_harmonics, SphericalHarmonics) - assert embedding_block.spherical_harmonics.l_max == embedding_config["l_max"] - - # Check cutoff_fn - assert hasattr(embedding_block, "cutoff_fn") - assert isinstance(embedding_block.cutoff_fn, CosineCutoff) - assert embedding_block.cutoff_fn.config.r_cut == embedding_config["r_max"] - - def test_config_storage(self, embedding_block, embedding_config): - """Test that configuration is properly stored.""" - assert hasattr(embedding_block, "config") - config = embedding_block.config - assert config.num_features == embedding_config["num_features"] - assert config.r_max == embedding_config["r_max"] - assert config.num_bessel == embedding_config["num_bessel"] - assert config.l_max == embedding_config["l_max"] - - def test_forward_output_shapes(self, embedding_block, embedding_config): - """Test forward pass returns correct output shapes.""" - n_atoms = 4 - n_edges = 6 - - # Create input data - z = torch.randint(0, embedding_config["num_species"], (n_atoms,)) - edge_dist = torch.rand(n_edges) * embedding_config["r_max"] - edge_diff = torch.randn(n_edges, 3) - - # Normalize edge_diff to match edge_dist - edge_diff = ( - edge_diff / torch.norm(edge_diff, dim=-1, keepdim=True) * edge_dist.unsqueeze(-1) - ) - - # Forward pass - node_feats, edge_attrs, edge_feats = embedding_block( - Z=z, - edge_dist=edge_dist, - edge_diff=edge_diff, - ) - - # Check shapes - assert node_feats.shape == (n_atoms, embedding_config["num_features"]) - - # Spherical harmonics dimension: (2*l_max + 1)^2 for l_max=2 is 9 - expected_sh_dim = (embedding_config["l_max"] + 1) ** 2 - assert edge_attrs.shape == (n_edges, expected_sh_dim) - - assert edge_feats.shape == (n_edges, embedding_config["num_bessel"]) - - def test_node_embedding_component(self, embedding_block, embedding_config): - """Test node_embedding component works independently.""" - n_atoms = 5 - z = torch.randint(0, embedding_config["num_species"], (n_atoms,)) - - # Call node_embedding directly - node_feats = embedding_block.node_embedding(Z=z) - - assert node_feats.shape == (n_atoms, embedding_config["num_features"]) - assert node_feats.dtype == torch.float32 - - def test_radial_embedding_component(self, embedding_block, embedding_config): - """Test radial_embedding component works independently.""" - n_edges = 10 - edge_dist = torch.rand(n_edges) * embedding_config["r_max"] - - # Call radial_embedding directly - edge_radial = embedding_block.radial_embedding(edge_dist) - - assert edge_radial.shape == (n_edges, embedding_config["num_bessel"]) - assert edge_radial.dtype == torch.float32 - - def test_spherical_harmonics_component(self, embedding_block, embedding_config): - """Test spherical_harmonics component works independently.""" - n_edges = 8 - # Create normalized direction vectors - edge_dir = torch.randn(n_edges, 3) - edge_dir = edge_dir / torch.norm(edge_dir, dim=-1, keepdim=True) - - # Call spherical_harmonics directly - edge_attrs = embedding_block.spherical_harmonics(edge_dir) - - expected_sh_dim = (embedding_config["l_max"] + 1) ** 2 - assert edge_attrs.shape == (n_edges, expected_sh_dim) - assert edge_attrs.dtype == torch.float32 - - def test_cutoff_component(self, embedding_block, embedding_config): - """Test cutoff_fn component works independently.""" - n_edges = 12 - edge_dist = torch.rand(n_edges) * embedding_config["r_max"] - - # Call cutoff_fn directly - cutoff_values = embedding_block.cutoff_fn(edge_dist) - - assert cutoff_values.shape == (n_edges,) - assert cutoff_values.dtype == torch.float32 - # Cutoff should be in [0, 1] - assert (cutoff_values >= 0.0).all() - assert (cutoff_values <= 1.0).all() - - def test_edge_feats_includes_cutoff(self, embedding_block, embedding_config): - """Test that edge_feats properly applies cutoff to radial basis.""" - n_edges = 6 - edge_dist = torch.rand(n_edges) * embedding_config["r_max"] - edge_diff = torch.randn(n_edges, 3) - edge_diff = ( - edge_diff / torch.norm(edge_diff, dim=-1, keepdim=True) * edge_dist.unsqueeze(-1) - ) - - z = torch.randint(0, embedding_config["num_species"], (3,)) - - # Get outputs - _, _, edge_feats = embedding_block( - Z=z, - edge_dist=edge_dist, - edge_diff=edge_diff, - ) - - # Compute expected edge_feats manually - edge_radial = embedding_block.radial_embedding(edge_dist) - cutoff_values = embedding_block.cutoff_fn(edge_dist) - expected_edge_feats = edge_radial * cutoff_values.unsqueeze(-1) - - # Check they match - assert torch.allclose(edge_feats, expected_edge_feats, atol=1e-6) - - def test_cutoff_at_boundary(self, embedding_block, embedding_config): - """Test cutoff behavior at r_max boundary.""" - # Distance at cutoff should give near-zero cutoff value - edge_dist = torch.tensor([embedding_config["r_max"]]) - cutoff_value = embedding_block.cutoff_fn(edge_dist) - - # Cosine cutoff should be near 0 at r_max - assert cutoff_value.item() < 0.01 - - # Distance at 0 should give cutoff value of 1 - bond_dist_zero = torch.tensor([0.0]) - cutoff_value_zero = embedding_block.cutoff_fn(bond_dist_zero) - assert abs(cutoff_value_zero.item() - 1.0) < 0.01 - - -class TestInteractionBlock: - """Test InteractionBlock initialization and forward pass.""" - - @pytest.fixture - def interaction_config(self): - """Common configuration for interaction block tests.""" - return { - "num_features": 64, - "num_bessel": 8, - "l_max": 2, - "avg_num_neighbors": 10.0, - } - - @pytest.fixture - def interaction_block(self, interaction_config): - """Create an InteractionBlock instance.""" - return InteractionBlock(**interaction_config) - - def test_initialization(self, interaction_block, interaction_config): - """Test that InteractionBlock initializes all components correctly.""" - # Check conv_tp (must be created first to provide weight_numel) - assert hasattr(interaction_block, "conv_tp") - assert isinstance(interaction_block.conv_tp, ConvTP) - - # Check node_linear - assert hasattr(interaction_block, "node_linear") - assert isinstance(interaction_block.node_linear, cuet.Linear) - - # Check radial_mlp - assert hasattr(interaction_block, "radial_mlp") - assert isinstance(interaction_block.radial_mlp, RadialWeightMLP) - - # Check linear - assert hasattr(interaction_block, "linear") - assert isinstance(interaction_block.linear, cuet.Linear) - - # Check avg_num_neighbors - assert hasattr(interaction_block, "avg_num_neighbors") - assert interaction_block.avg_num_neighbors == interaction_config["avg_num_neighbors"] - - def test_initialization_order_fix(self, interaction_config): - """Test that conv_tp is initialized before radial_mlp (bug fix).""" - # This should NOT raise AttributeError about self.conv_tp - block = InteractionBlock(**interaction_config) - - # Verify conv_tp exists and has weight_numel - assert hasattr(block, "conv_tp") - assert hasattr(block.conv_tp, "weight_numel") - - # Verify MLP output dimension matches weight_numel - mlp_output_layer = block.radial_mlp.mlp[-1] - assert mlp_output_layer.out_features == block.conv_tp.weight_numel - - def test_config_storage(self, interaction_block, interaction_config): - """Test that configuration is properly stored.""" - assert hasattr(interaction_block, "config") - config = interaction_block.config - assert config.num_features == interaction_config["num_features"] - assert config.num_bessel == interaction_config["num_bessel"] - assert config.l_max == interaction_config["l_max"] - assert config.avg_num_neighbors == interaction_config["avg_num_neighbors"] - - def test_forward_output_shapes(self, interaction_block, interaction_config): - """Test forward pass returns correct output shapes.""" - n_nodes = 20 - n_edges = 50 - l_max = interaction_config["l_max"] - num_features = interaction_config["num_features"] - num_bessel = interaction_config["num_bessel"] - - # Calculate irreps dimension (uniform multiplicity after optimization). - # This is the *message* dim; the node state itself is pure scalar. - irreps_dim = sum(num_features * (2 * l + 1) for l in range(l_max + 1)) - - # Calculate spherical harmonics dimension - sh_dim = sum(2 * l + 1 for l in range(l_max + 1)) - - # Node state is pure scalar (num_features); messages are mixed-l. - node_feats = torch.randn(n_nodes, num_features) - edge_attrs = torch.randn(n_edges, sh_dim) - edge_feats = torch.randn(n_edges, num_bessel) - edge_index = torch.randint(0, n_nodes, (n_edges, 2)) - - # Forward pass - output_feats, skip_connection = interaction_block( - node_feats=node_feats, - edge_attrs=edge_attrs, - edge_feats=edge_feats, - edge_index=edge_index, - ) - - # Message is mixed-l (irreps_dim); skip is the scalar node state. - assert output_feats.shape == (n_nodes, irreps_dim) - assert skip_connection.shape == (n_nodes, num_features) - - def test_skip_connection_is_input(self, interaction_block, interaction_config): - """Test that skip connection returns the original input.""" - n_nodes = 10 - n_edges = 30 - l_max = interaction_config["l_max"] - num_features = interaction_config["num_features"] - - # Node state is pure scalar (num_features); messages are mixed-l. - sh_dim = sum(2 * l + 1 for l in range(l_max + 1)) - - # Create input data - node_feats = torch.randn(n_nodes, num_features) - edge_attrs = torch.randn(n_edges, sh_dim) - edge_feats = torch.randn(n_edges, interaction_config["num_bessel"]) - edge_index = torch.randint(0, n_nodes, (n_edges, 2)) - - # Forward pass - _, skip_connection = interaction_block( - node_feats=node_feats, - edge_attrs=edge_attrs, - edge_feats=edge_feats, - edge_index=edge_index, - ) - - # Skip connection should be identical to input - assert torch.equal(skip_connection, node_feats) - - def test_radial_mlp_architecture(self, interaction_block, interaction_config): - """Test radial_mlp MLP architecture.""" - mlp = interaction_block.radial_mlp.mlp - num_features = interaction_config["num_features"] - num_bessel = interaction_config["num_bessel"] - weight_numel = interaction_block.conv_tp.weight_numel - - # Check number of layers (Linear -> SiLU -> Linear -> SiLU -> Linear) - assert len(mlp) == 5 - - # Check layer types - assert isinstance(mlp[0], nn.Linear) - assert isinstance(mlp[1], nn.SiLU) - assert isinstance(mlp[2], nn.Linear) - assert isinstance(mlp[3], nn.SiLU) - assert isinstance(mlp[4], nn.Linear) - - # Check dimensions - assert mlp[0].in_features == num_bessel - assert mlp[0].out_features == num_features - assert mlp[2].in_features == num_features - assert mlp[2].out_features == num_features - assert mlp[4].in_features == num_features - assert mlp[4].out_features == weight_numel - - def test_cuequivariance_integration(self, interaction_block): - """Test that cuEquivariance ChannelWiseTensorProduct is used.""" - # Conv_tp should wrap a ChannelWiseTensorProduct - assert hasattr(interaction_block.conv_tp, "cue_tp") - assert isinstance(interaction_block.conv_tp.cue_tp, cuet.ChannelWiseTensorProduct) - - # Verify it has the expected ChannelWiseTensorProduct configuration - cue_tp = interaction_block.conv_tp.cue_tp - assert hasattr(cue_tp, "irreps_in1") - assert hasattr(cue_tp, "irreps_in2") - assert hasattr(cue_tp, "irreps_out") - - def test_different_l_max_values(self, interaction_config): - """Test initialization with different l_max values.""" - for l_max in [1, 2, 3]: - block = InteractionBlock( - num_features=interaction_config["num_features"], - num_bessel=interaction_config["num_bessel"], - l_max=l_max, - avg_num_neighbors=interaction_config["avg_num_neighbors"], - ) - - # Should initialize without errors - assert block.config.l_max == l_max - assert hasattr(block.conv_tp, "weight_numel") - - def test_different_num_features(self, interaction_config): - """Test initialization with different num_features values.""" - for num_features in [32, 64, 128]: - block = InteractionBlock( - num_features=num_features, - num_bessel=interaction_config["num_bessel"], - l_max=interaction_config["l_max"], - avg_num_neighbors=interaction_config["avg_num_neighbors"], - ) - - # MLP hidden dimension should match num_features - assert block.radial_mlp.mlp[0].out_features == num_features - assert block.radial_mlp.mlp[2].in_features == num_features - - def test_edge_index_format(self, interaction_block, interaction_config): - """Test that edge_index format (n_edges, 2) works correctly.""" - n_nodes = 15 - n_edges = 40 - l_max = interaction_config["l_max"] - num_features = interaction_config["num_features"] - - # Message dim is mixed-l; node state is pure scalar. - irreps_dim = sum(num_features * (2 * l + 1) for l in range(l_max + 1)) - sh_dim = sum(2 * l + 1 for l in range(l_max + 1)) - - # Create input data (scalar node state) - node_feats = torch.randn(n_nodes, num_features) - edge_attrs = torch.randn(n_edges, sh_dim) - edge_feats = torch.randn(n_edges, interaction_config["num_bessel"]) - edge_index = torch.randint(0, n_nodes, (n_edges, 2)) - - # Forward pass should work without errors - output, _ = interaction_block(node_feats, edge_attrs, edge_feats, edge_index) - assert output.shape == (n_nodes, irreps_dim) - - -class TestProductHead: - """Test ProductHead initialization and forward pass.""" - - @pytest.fixture - def product_config(self): - """Common configuration for product head tests.""" - return { - "hidden_dim": 576, # For num_features=64, l_max=2 (uniform): 64*1 + 64*3 + 64*5 = 576 - "out_dim": 64, - "num_radial": 8, - "l_max": 2, - "max_body_order": 2, - "num_species": 118, - } - - @pytest.fixture - def product_head(self, product_config): - """Create a ProductHead instance.""" - return ProductHead(**product_config) - - def test_initialization(self, product_head, product_config): - """Test that EquivariantProductBasisBlock initializes all components correctly.""" - # Check symmetric_contraction - assert hasattr(product_head, "symmetric_contraction") - assert isinstance(product_head.symmetric_contraction, SymmetricContraction) - - # Check basis_projection - assert hasattr(product_head, "basis_projection") - assert isinstance(product_head.basis_projection, BasisProjection) - - # Check linear - assert hasattr(product_head, "linear") - assert isinstance(product_head.linear, nn.Linear) - # The contraction emits invariant scalars (num_features), so the readout - # linear maps num_features -> out_dim (not the full mixed-l hidden_dim). - num_features = product_config["hidden_dim"] // (product_config["l_max"] + 1) ** 2 - assert product_head.linear.in_features == num_features - assert product_head.linear.out_features == product_config["out_dim"] - - def test_symmetric_contraction_config(self, product_head, product_config): - """Test that SymmetricContraction is configured correctly.""" - sc = product_head.symmetric_contraction - assert sc.config.hidden_dim == product_config["hidden_dim"] - assert sc.config.num_species == product_config["num_species"] - assert sc.config.max_body_order == product_config["max_body_order"] - - def test_basis_projection_config(self, product_head, product_config): - """Test that BasisProjection is configured correctly.""" - bp = product_head.basis_projection - assert bp.config.hidden_dim == product_config["hidden_dim"] - assert bp.config.num_radial == product_config["num_radial"] - assert bp.config.l_max == product_config["l_max"] - assert bp.config.max_body_order == product_config["max_body_order"] - - def test_forward_output_shapes(self, product_head, product_config): - """Test forward pass returns correct output shapes.""" - n_nodes = 20 - hidden_dim = product_config["hidden_dim"] - out_dim = product_config["out_dim"] - num_species = product_config["num_species"] - - # Create input data - node_features = torch.randn(n_nodes, hidden_dim) - atom_types = torch.randint(0, num_species, (n_nodes,)) - - # Forward pass - output = product_head(node_features, atom_types) - - # Check output shape - assert output.shape == (n_nodes, out_dim) - - def test_forward_dtype(self, product_head, product_config): - """Test that forward pass preserves dtype.""" - n_nodes = 10 - hidden_dim = product_config["hidden_dim"] - - # Test with float32 - node_features = torch.randn(n_nodes, hidden_dim, dtype=torch.float32) - atom_types = torch.randint(0, product_config["num_species"], (n_nodes,)) - - output = product_head(node_features, atom_types) - assert output.dtype == torch.float32 - - def test_different_max_body_orders(self, product_config): - """Test initialization with different max_body_order values.""" - for max_body_order in [1, 2, 3]: - head = ProductHead( - hidden_dim=product_config["hidden_dim"], - out_dim=product_config["out_dim"], - num_radial=product_config["num_radial"], - l_max=product_config["l_max"], - max_body_order=max_body_order, - num_species=product_config["num_species"], - ) - - # Should initialize without errors - assert head.symmetric_contraction.config.max_body_order == max_body_order - assert head.basis_projection.config.max_body_order == max_body_order - - def test_different_num_species(self, product_config): - """Test initialization with different num_species values.""" - for num_species in [10, 50, 118]: - head = ProductHead( - hidden_dim=product_config["hidden_dim"], - out_dim=product_config["out_dim"], - num_radial=product_config["num_radial"], - l_max=product_config["l_max"], - max_body_order=product_config["max_body_order"], - num_species=num_species, - ) - - assert head.symmetric_contraction.config.num_species == num_species - - def test_batch_processing(self, product_head, product_config): - """Test that the head processes batches correctly.""" - # Test with different batch sizes - for n_nodes in [5, 20, 100]: - node_features = torch.randn(n_nodes, product_config["hidden_dim"]) - atom_types = torch.randint(0, product_config["num_species"], (n_nodes,)) - - output = product_head(node_features, atom_types) - assert output.shape == (n_nodes, product_config["out_dim"]) - - def test_symmetric_contraction_component(self, product_head, product_config): - """Test symmetric_contraction component works independently.""" - n_nodes = 15 - hidden_dim = product_config["hidden_dim"] - - # Input is the mixed-l message (hidden_dim); the contraction outputs - # invariant scalars (num_features). - num_features = hidden_dim // (product_config["l_max"] + 1) ** 2 - node_features = torch.randn(n_nodes, hidden_dim) - atom_types = torch.randint(0, product_config["num_species"], (n_nodes,)) - - # Call symmetric_contraction directly - basis = product_head.symmetric_contraction(node_features, atom_types) - - # Contraction maps mixed-l input -> invariant scalar output. - assert basis.shape == (n_nodes, num_features) - - def test_basis_projection_component(self, product_head, product_config): - """Test basis_projection component works independently.""" - n_nodes = 15 - hidden_dim = product_config["hidden_dim"] - - # Create dummy basis features - basis = torch.randn(n_nodes, hidden_dim) - - # Call basis_projection directly (passthrough in current implementation) - features = product_head.basis_projection(basis) - - # Currently acts as identity - assert torch.equal(features, basis) - - def test_linear_component(self, product_head, product_config): - """Test linear component works independently.""" - n_nodes = 15 - out_dim = product_config["out_dim"] - - # The readout linear consumes the contracted scalars (num_features). - num_features = product_config["hidden_dim"] // (product_config["l_max"] + 1) ** 2 - features = torch.randn(n_nodes, num_features) - - # Call linear directly - output = product_head.linear(features) - - assert output.shape == (n_nodes, out_dim) - - def test_gradient_flow(self, product_head, product_config): - """Test that gradients flow through the head correctly.""" - n_nodes = 10 - hidden_dim = product_config["hidden_dim"] - - node_features = torch.randn(n_nodes, hidden_dim, requires_grad=True) - atom_types = torch.randint(0, product_config["num_species"], (n_nodes,)) - - # Forward pass - output = product_head(node_features, atom_types) - - # Backward pass - loss = output.sum() - loss.backward() - - # Check gradients exist - assert node_features.grad is not None - assert not torch.isnan(node_features.grad).any() - - def test_cuequivariance_integration(self, product_head): - """Test that cuEquivariance SymmetricContraction is used.""" - sc = product_head.symmetric_contraction - assert hasattr(sc, "symmetric_contraction") - assert isinstance(sc.symmetric_contraction, cuet.SymmetricContraction) - - # Verify configuration - cue_sc = sc.symmetric_contraction - assert hasattr(cue_sc, "contraction_degree") - assert hasattr(cue_sc, "num_elements") - - -class TestEmbeddingBlockEquivariance: - """Test equivariance properties of EmbeddingBlock.""" - - @pytest.fixture - def embedding_block(self): - """Create an EmbeddingBlock for equivariance testing.""" - node_attr_specs = [ - DiscreteEmbeddingSpec( - input_key="Z", - num_classes=5, - emb_dim=16, - ) - ] - return EmbeddingBlock( - node_attr_specs=node_attr_specs, - num_features=16, - r_max=5.0, - num_bessel=8, - l_max=2, - ) - - def test_spherical_harmonics_equivariance(self, embedding_block): - """Test that spherical harmonics are equivariant under rotation. - - Rotating the bond vectors should rotate the spherical harmonics accordingly. - """ - n_atoms = 4 - n_edges = 6 - - # Create input data - z = torch.randint(0, 5, (n_atoms,)) - edge_diff = torch.randn(n_edges, 3) - edge_dist = torch.norm(edge_diff, dim=-1) - - # Forward pass - _, edge_attrs1, _ = embedding_block( - Z=z, - edge_dist=edge_dist, - edge_diff=edge_diff, - ) - - # Rotate bond vectors - angle = math.pi / 2 - rot_matrix = rotation_matrix_z(angle, dtype=edge_diff.dtype) - bond_diff_rot = rotate_vectors(edge_diff, rot_matrix) - - # Forward pass on rotated - _, edge_attrs2, _ = embedding_block( - Z=z, - edge_dist=edge_dist, - edge_diff=bond_diff_rot, - ) - - # l=0 component should be invariant - assert torch.allclose(edge_attrs1[:, 0], edge_attrs2[:, 0], atol=1e-5) - - # Overall norm should be preserved - norm1 = edge_attrs1.norm(dim=-1) - norm2 = edge_attrs2.norm(dim=-1) - assert torch.allclose(norm1, norm2, rtol=1e-4, atol=1e-4) - - def test_radial_features_invariance(self, embedding_block): - """Test that radial features are rotation invariant. - - Rotating bond vectors should not change radial features (distances). - """ - n_atoms = 4 - n_edges = 6 - - z = torch.randint(0, 5, (n_atoms,)) - edge_diff = torch.randn(n_edges, 3) - edge_dist = torch.norm(edge_diff, dim=-1) - - # Forward pass - _, _, edge_feats1 = embedding_block( - Z=z, - edge_dist=edge_dist, - edge_diff=edge_diff, - ) - - # Rotate bond vectors - rot_matrix = random_rotation_matrix(dtype=edge_diff.dtype) - bond_diff_rot = rotate_vectors(edge_diff, rot_matrix) - - # Forward pass on rotated - _, _, edge_feats2 = embedding_block( - Z=z, - edge_dist=edge_dist, - edge_diff=bond_diff_rot, - ) - - # Radial features should be identical (rotation invariant) - assert torch.allclose(edge_feats1, edge_feats2, rtol=1e-5, atol=1e-5) - - -class TestInteractionBlockEquivariance: - """Test equivariance properties of InteractionBlock.""" - - @pytest.fixture - def interaction_block(self): - """Create an InteractionBlock for testing.""" - return InteractionBlock( - num_features=32, - num_bessel=8, - l_max=1, # Use l_max=1 for simpler testing - avg_num_neighbors=10.0, - ) - - def test_output_shape_consistency(self, interaction_block): - """Test that rotation doesn't change output shapes.""" - n_nodes = 10 - n_edges = 30 - l_max = 1 - num_features = 32 - - # Message dim is mixed-l; node state is pure scalar. - irreps_dim = sum(num_features * (2 * l + 1) for l in range(l_max + 1)) - sh_dim = sum(2 * l + 1 for l in range(l_max + 1)) - - node_feats = torch.randn(n_nodes, num_features) - edge_attrs = torch.randn(n_edges, sh_dim) - edge_feats = torch.randn(n_edges, 8) - edge_index = torch.randint(0, n_nodes, (n_edges, 2)) - - # Forward pass - output, _ = interaction_block(node_feats, edge_attrs, edge_feats, edge_index) - - assert output.shape == (n_nodes, irreps_dim) - - -class TestProductHeadEquivariance: - """Test equivariance properties of ProductHead.""" - - @pytest.fixture - def product_head(self): - """Create a ProductHead for testing.""" - return ProductHead( - hidden_dim=160, # 32*1 + 32*3 = 128 for l_max=1, out_dim=32 - out_dim=32, - num_radial=8, - l_max=1, - max_body_order=2, - num_species=10, - ) - - def test_permutation_equivariance(self, product_head): - """Test that ProductHead is equivariant to node permutations.""" - n_nodes = 15 - hidden_dim = 160 - - node_features = torch.randn(n_nodes, hidden_dim) - atom_types = torch.randint(0, 10, (n_nodes,)) - - # Forward pass - output1 = product_head(node_features, atom_types) - - # Permute nodes - perm = torch.randperm(n_nodes) - node_features_perm = node_features[perm] - atom_types_perm = atom_types[perm] - - # Forward on permuted - output2 = product_head(node_features_perm, atom_types_perm) - - # Outputs should match after inverse permutation - assert torch.allclose(output1[perm], output2, rtol=1e-5, atol=1e-5) diff --git a/tests/test_molzoo/test_mace/__init__.py b/tests/test_molzoo/test_mace/__init__.py new file mode 100644 index 0000000..55d99fc --- /dev/null +++ b/tests/test_molzoo/test_mace/__init__.py @@ -0,0 +1 @@ +"""Unit tests mirroring the ``molzoo.mace`` sub-package.""" diff --git a/tests/test_molzoo/test_mace/conftest.py b/tests/test_molzoo/test_mace/conftest.py new file mode 100644 index 0000000..045fe07 --- /dev/null +++ b/tests/test_molzoo/test_mace/conftest.py @@ -0,0 +1,301 @@ +"""Shared fixtures for the ``molzoo.mace`` sub-package tests. + +Only MACE-family constants, geometries, models and specs live here. Batch +construction goes through :func:`tests.conftest.make_graph_batch` — this +package deliberately does not grow a second batch builder (the three private +builders of the pre-cutover flat test files were folded into it by +``mace-subpackage-restructure-06-wire``, and none may come back). + +The tiny configurations below are the ones the pre-cutover flat variant tests +used, so ``MACEEncoder`` / ``MACEPotential`` can be diffed against the keyword +variants key-for-key. + +Everything a second module in this package needs is defined here exactly once: + +* :data:`CLUSTER_POS` / :data:`CLUSTER_Z` and the ``cluster`` / + ``periodic_cluster`` / ``omol_cluster`` / ``pair_batch`` fixtures — the + fixed geometries (no RNG anywhere); +* :func:`full_edge_index` — the ordered intra-graph pair list; +* :func:`wake_zero_init_readout` — pins the OMOL readout to a private + generator, so state-transfer parity does not ride on the global RNG; +* ``matpes_variant`` / ``omol_variant`` — the keyword-constructed foundation + models, used both as the subject of ``test_variants.py`` and as the + raw-tensor (``energy_forces``) oracle of ``test_potential.py`` / + ``test_encoder.py``. +""" + +from __future__ import annotations + +from collections.abc import Sequence +from typing import Any + +import pytest +import torch +from tensordict import TensorDict + +from molix import config +from molzoo.mace.spec import MACEMatpesSpec, MACEOMolSpec +from molzoo.mace.variants import MACEMatpes, MACEOMol +from tests.conftest import make_graph_batch + +#: Element table (z-table) shared by every model built in this package. +#: Ascending and duplicate-free — ``torch.searchsorted`` requires it. +ATOMIC_NUMBERS: list[int] = [1, 6, 8] + +#: Per-element reference energies ``E0`` in eV/atom, same order as +#: :data:`ATOMIC_NUMBERS`. +ATOMIC_ENERGIES: list[float] = [-13.6, -1029.0, -2041.0] + +#: Tiny MACE-MatPES hyper-parameters (l_max=1, 16 channels): fast on CPU and +#: structurally identical to the shipped model. Passed verbatim to both +#: ``MACEMatpesSpec`` and the keyword ``MACEMatpes`` constructor, which is what +#: makes the ``state_dict`` parity assertion meaningful. +TINY_MATPES_KWARGS: dict[str, Any] = { + "r_max": 5.0, + "num_bessel": 4, + "num_polynomial_cutoff": 5, + "l_max": 1, + "num_features": 16, + "max_hidden_l": 1, + "num_interactions": 2, + "correlation": 2, + "mlp_dim": 8, + "radial_mlp": [8], + "use_fallback": True, # CPU tests: fused kernels need a GPU + ops wheel +} + +#: Tiny MACE-OMOL hyper-parameters (l_max=1, 16 channels). ``use_fallback`` is +#: absent: :class:`~molzoo.mace.variants.MACEOMol` takes no such keyword, so the +#: spec must default/carry ``use_fallback=False`` for the two stacks to agree +#: key-wise. +TINY_OMOL_KWARGS: dict[str, Any] = { + "r_max": 5.0, + "num_bessel": 4, + "num_polynomial_cutoff": 5, + "l_max": 1, + "num_features": 16, + "num_interactions": 2, + "correlation": 2, + "mlp_dim": 8, + "edge_channels": 8, +} + +#: A five-atom cluster at literal coordinates (Å) — no RNG anywhere. +CLUSTER_POS: list[list[float]] = [ + [0.00, 0.00, 0.00], + [0.95, 0.00, 0.00], + [-0.24, 0.93, 0.00], + [0.00, 0.00, 1.40], + [1.20, 1.10, 0.60], +] + +#: Atomic numbers of :data:`CLUSTER_POS`, all inside the ``[1, 6, 8]`` table. +CLUSTER_Z: list[int] = [8, 1, 1, 6, 1] + +#: Graph membership of :data:`CLUSTER_POS` — one graph, ``B = 1``. +SINGLE_GRAPH: list[int] = [0] * len(CLUSTER_Z) + +#: Two clusters (4 + 3 atoms) in one batch — exercises the per-graph reduction. +PAIR_POS: list[list[float]] = CLUSTER_POS[:4] + [ + [5.00, 5.00, 5.00], + [5.95, 5.00, 5.00], + [5.00, 5.90, 5.00], +] +PAIR_Z: list[int] = [8, 1, 1, 6, 8, 1, 1] +PAIR_BATCH: list[int] = [0, 0, 0, 0, 1, 1, 1] + +#: Periodic shift added to the first two edges (Å) — ``unit_shifts @ cell``. +PERIODIC_SHIFT = 2.0 + + +def full_edge_index(batch: Sequence[int]) -> torch.Tensor: + """All ordered intra-graph atom pairs as ``(E, 2)`` ``[source, target]``. + + Args: + batch: Graph index per atom; ``len(batch)`` is the atom count. + + Returns: + ``(E, 2)`` edge index on the repo convention (``[:, 0]`` = source, + ``[:, 1]`` = target), self-pairs and cross-graph pairs excluded. + """ + n_atoms = len(batch) + pairs = [ + [i, j] for i in range(n_atoms) for j in range(n_atoms) if i != j and batch[i] == batch[j] + ] + return torch.tensor(pairs, dtype=torch.long) + + +def raw_tensors( + batch: TensorDict, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, int]: + """``(pos, Z, edge_index, atom_batch, num_graphs)`` off a post-collate batch. + + The raw-tensor seams (``energy_core`` / ``energy_forces``) take these five + in this order; every case that crosses from the batch schema to a flat call + goes through here rather than unpacking by hand. + + Args: + batch: Post-collate ``TensorDict`` with ``atoms`` / ``edges``. + + Returns: + Positions ``(N, 3)`` Å, atomic numbers ``(N,)``, edge index ``(E, 2)``, + graph index per atom ``(N,)`` and the graph count ``B``. + """ + return ( + batch["atoms", "pos"], + batch["atoms", "Z"], + batch["edges", "edge_index"], + batch["atoms", "batch"], + int(batch["atoms", "batch"].max()) + 1, + ) + + +def wake_zero_init_readout(model: torch.nn.Module) -> None: + """Pin the OMOL readout to a private generator, off the global RNG. + + ``molrep.readout.mace._ScalarO3Linear`` now draws its weight from ``N(0, 1)`` + on the **global** RNG and zero-initialises only its bias, so a fresh OMOL + ``NonLinearBiasReadout`` is already non-trivial and the assertions here are + no longer at risk of being vacuous. (They were: the weight used to be + zero-initialised too, which made an untrained model's interaction energy + position-independent and every force identically zero, so parity, physics + and equivalence assertions all reduced to ``0 == 0``.) + + The helper stays as belt-and-braces determinism for the state-transfer + parity tests: it fills both scalar linears from a private, fixed + ``torch.Generator``, so the values a test compares across two models depend + on that generator alone and not on how much global RNG each model happened + to consume before its readout was built. Checkpoint use is unaffected — + official weights overwrite these entries. + + Args: + model: An OMOL-configured model exposing ``readout.linear_mid`` / + ``readout.linear_2``. + """ + generator = torch.Generator().manual_seed(0) + with torch.no_grad(): + for linear in (model.readout.linear_mid, model.readout.linear_2): + linear.weight.normal_(generator=generator) + linear.bias.normal_(generator=generator) + + +@pytest.fixture(autouse=True) +def fp64(): + """Build every model in this package at fp64. + + cuEquivariance freezes its working precision at construction, so a model + built under the default fp32 and then ``.double()``-d still contracts in + float32 — ~1e-8 eV of noise, which is fine for inference but swamps the + invariance assertions here. Foundation-weight use is fp64 anyway, so the + tests exercise that path. + """ + previous = config["ftype"] + config.set_precision("fp64") + yield + config.set_precision("fp64" if previous == torch.float64 else "fp32") + + +@pytest.fixture +def tiny_matpes_spec() -> MACEMatpesSpec: + """Tiny :class:`molzoo.mace.spec.MACEMatpesSpec`.""" + return MACEMatpesSpec( + atomic_numbers=ATOMIC_NUMBERS, + atomic_energies=ATOMIC_ENERGIES, + **TINY_MATPES_KWARGS, + ) + + +@pytest.fixture +def tiny_omol_spec() -> MACEOMolSpec: + """Tiny :class:`molzoo.mace.spec.MACEOMolSpec`.""" + return MACEOMolSpec( + atomic_numbers=ATOMIC_NUMBERS, + atomic_energies=ATOMIC_ENERGIES, + use_fallback=False, # MACEOMol takes no such keyword; the spec defaults + **TINY_OMOL_KWARGS, + ) + + +@pytest.fixture +def cluster() -> TensorDict: + """A shift-free five-atom, one-graph batch (no charge/spin conditioning).""" + return make_graph_batch( + pos=torch.tensor(CLUSTER_POS, dtype=torch.float64), + Z=torch.tensor(CLUSTER_Z, dtype=torch.long), + edge_index=full_edge_index(SINGLE_GRAPH), + batch=torch.zeros(len(CLUSTER_Z), dtype=torch.long), + ) + + +@pytest.fixture +def periodic_cluster() -> TensorDict: + """The same cluster with a non-zero ``edges.shifts`` on the first two edges.""" + edge_index = full_edge_index(SINGLE_GRAPH) + shifts = torch.zeros(edge_index.shape[0], 3, dtype=torch.float64) + shifts[0, 0] = PERIODIC_SHIFT + shifts[1, 0] = -PERIODIC_SHIFT + return make_graph_batch( + pos=torch.tensor(CLUSTER_POS, dtype=torch.float64), + Z=torch.tensor(CLUSTER_Z, dtype=torch.long), + edge_index=edge_index, + batch=torch.zeros(len(CLUSTER_Z), dtype=torch.long), + shifts=shifts, + ) + + +@pytest.fixture +def pair_batch() -> TensorDict: + """Two separated clusters in one batch (``B = 2``).""" + return make_graph_batch( + pos=torch.tensor(PAIR_POS, dtype=torch.float64), + Z=torch.tensor(PAIR_Z, dtype=torch.long), + edge_index=full_edge_index(PAIR_BATCH), + batch=torch.tensor(PAIR_BATCH, dtype=torch.long), + ) + + +@pytest.fixture +def omol_cluster() -> TensorDict: + """The cluster with explicit OMOL conditioning (neutral closed-shell singlet).""" + return make_graph_batch( + pos=torch.tensor(CLUSTER_POS, dtype=torch.float64), + Z=torch.tensor(CLUSTER_Z, dtype=torch.long), + edge_index=full_edge_index(SINGLE_GRAPH), + batch=torch.zeros(len(CLUSTER_Z), dtype=torch.long), + graphs={ + "total_charge": torch.zeros(1, dtype=torch.long), + "total_spin": torch.ones(1, dtype=torch.long), + }, + ) + + +@pytest.fixture +def matpes_variant() -> MACEMatpes: + """Tiny :class:`~molzoo.mace.variants.MACEMatpes` in the keyword shape. + + ``atomic_energies`` is a **tensor**, as ``scripts/matpes_port/run_nve.py`` + passes it (``torch.tensor(cfg["atomic_energies"], dtype=config.ftype)``), + even though :class:`~molzoo.mace.spec.MACEMatpesSpec` holds a list. + """ + torch.manual_seed(0) + return MACEMatpes( + atomic_numbers=list(ATOMIC_NUMBERS), + atomic_energies=torch.tensor(ATOMIC_ENERGIES), + **TINY_MATPES_KWARGS, + ).eval() + + +@pytest.fixture +def omol_variant() -> MACEOMol: + """Tiny :class:`~molzoo.mace.variants.MACEOMol` with a woken readout. + + See :func:`wake_zero_init_readout` for why the readout is filled. + """ + torch.manual_seed(0) + model = MACEOMol( + atomic_numbers=list(ATOMIC_NUMBERS), + atomic_energies=torch.tensor(ATOMIC_ENERGIES), + **TINY_OMOL_KWARGS, + ).eval() + wake_zero_init_readout(model) + return model diff --git a/tests/test_molzoo/test_mace/test_checkpoint.py b/tests/test_molzoo/test_mace/test_checkpoint.py new file mode 100644 index 0000000..02fe0a7 --- /dev/null +++ b/tests/test_molzoo/test_mace/test_checkpoint.py @@ -0,0 +1,882 @@ +"""Tests for molzoo.mace.checkpoint — the one official-checkpoint key remap. + +:class:`~molzoo.mace.checkpoint.CheckpointRemap` replaces the two hand-copied +loaders of the pre-cutover flat modules (now the thin compatibility aliases +:func:`molzoo.mace.variants.load_matpes_state_dict` / +:func:`~molzoo.mace.variants.load_omol_state_dict`) with a single type whose only +family-dependent behaviour is the ``on_unexpected`` knob: MatPES raises on a +checkpoint key with no home, OMol returns it (that family ships auxiliary +heads this port deliberately does not model). Everything else — the +``.graph.c`` / ``output_mask`` skip, the longest-prefix rename, the +numel-conserving reshape of ``(1,)`` frozen scalars, and above all the +strictness doctrine — is shared and must stay identical. + +The doctrine is the point of most of this file. A checkpoint holds a *fitted +potential energy surface*: a tensor that is silently dropped yields a model +that runs, looks sane, and is quietly wrong (``mace_matpes.py:405-410``). So an +unfilled learnable parameter or a genuine shape disagreement raises under +*both* policies — the knob is not a laxness dial. ``test_rejects_missing_para +meter_without_weight_suffix`` is the regression lock for the second incident +(``mace_omol.py:437-439``): "learnable" means exactly ``nn.Parameter``, never a +``.weight`` / ``.bias`` name heuristic, which once excused ``bessel.freqs`` +from the check and left ~2e-7 of the official weights in the parity residual. + +Test data is built with the ``_official_state`` round trip — a model's own +``state_dict`` renamed into official cueq names and loaded back into a +differently initialised model — the trick the pre-cutover +``TestLoadMatpesStateDict._roundtrip_state`` used, now shared by that class +(migrated into this file by ``mace-subpackage-restructure-06-wire``) and by +:class:`TestCheckpointRemap`. The inverse name tables below are written out +**by hand** rather than derived from the tables under test, so a wrong entry in +the production table cannot cancel itself out in the round trip. + +Every model here is tiny (2 layers, 16 channels, ``l_max=1``), fp64 (autouse +``fp64`` fixture in ``conftest.py``), CPU, ``use_fallback=True`` (no fused +cuEquivariance wheel on CPU) and seeded — see +``.claude/specs/mace-subpackage-restructure-05-checkpoint.md`` §Testing +strategy. Units are eV / eV·Å throughout; the remap converts nothing. + +Reference: + Batatia et al. "MACE: Higher Order Equivariant Message Passing Neural + Networks for Fast and Accurate Force Fields" NeurIPS 2022. + https://arxiv.org/abs/2206.07697 +""" + +from __future__ import annotations + +import json +import os +from collections.abc import Mapping +from pathlib import Path + +import pytest +import torch +from tensordict import TensorDict + +from molzoo.mace.checkpoint import ( + MATPES_KEY_REMAP, + MATPES_REMAP, + OMOL_REMAP, + CheckpointRemap, +) +from molzoo.mace.potential import MACEPotential +from molzoo.mace.spec import MACEMatpesSpec, MACEOMolSpec +from molzoo.mace.variants import MACEMatpes, load_matpes_state_dict +from tests.test_molzoo.test_mace.conftest import ( + ATOMIC_ENERGIES, + ATOMIC_NUMBERS, + TINY_MATPES_KWARGS, +) + +#: Model name → official cueq name: the hand-written inverse of +#: ``MATPES_KEY_REMAP``, longest prefix first (see the module docstring for why +#: it is not derived from the table under test). +MATPES_OFFICIAL_NAMES: dict[str, str] = { + "node_embedding.": "node_embedding.linear.", + "bessel.freqs": "radial_embedding.bessel_fn.bessel_weights", + "distance_transform.": "radial_embedding.distance_transform.", + "pair_repulsion.": "pair_repulsion_fn.", + "z_table": "atomic_numbers", +} + +#: Model name → official cueq name for the OMol family (inverse of +#: ``OMOL_KEY_REMAP``). ``readout.`` is the single final readout molnex holds +#: where the official checkpoint numbers a one-entry ``readouts`` list. +OMOL_OFFICIAL_NAMES: dict[str, str] = { + "node_embedding.": "node_embedding.linear.", + "embedding_readout.": "embedding_readout.linear.", + "readout.": "readouts.0.", + "bessel.freqs": "radial_embedding.bessel_fn.bessel_weights", + "z_table": "atomic_numbers", +} + +#: A checkpoint key that matches no table entry and no model parameter — the +#: probe for the ``on_unexpected`` knob (same key as the flat loader's test). +MYSTERY_KEY = "interactions.0.mystery_layer.weight" + +#: A learnable weight every MatPES model has; deleting it must raise. +COVERED_PARAMETER = "interactions.0.linear_up.weight" + +#: Official name of ``bessel.freqs`` — an ``nn.Parameter`` ending in neither +#: ``.weight`` nor ``.bias``. Deleting it is the ``mace_omol.py:437-439`` +#: regression lock. +BESSEL_OFFICIAL_KEY = "radial_embedding.bessel_fn.bessel_weights" + +#: MACE stores some frozen scalars as ``(1,)`` where molnex holds a 0-d buffer. +FROZEN_SCALAR_VALUE = 0.75 + +#: Environment variable pointing at the offline official-weights directory. +WEIGHTS_DIR_ENV = "MOLNEX_MACE_WEIGHTS_DIR" + +#: The two architecture switches of the stock MatPES dump, spelled exactly as +#: ``matpes_r2scan_config.json`` spells them (read 2026-08-09 from the offline +#: weights directory): ZBL pair repulsion on, Agnesi distance transform — +#: capital ``A``. :class:`~molzoo.mace.spec.MACEMatpesSpec` hard-wires both, so +#: these are the only *present* values ``from_checkpoint`` may accept, and it +#: must accept them without caring about the case of the ``A``. +OFFICIAL_SWITCHES: dict[str, object] = {"pair_repulsion": True, "distance_transform": "Agnesi"} + +#: Hard-coded goldens: the dimensions of ``MACE-matpes-r2scan-omat-ft`` as +#: published (``hidden_irreps="128x0e+128x1o"``, ``MLP_irreps="16x0e"``). The +#: bit-parity oracle states them literally so ``from_checkpoint`` has to derive +#: the same numbers from the irreps strings on its own. +OFFICIAL_NUM_FEATURES = 128 +OFFICIAL_MAX_HIDDEN_L = 1 +OFFICIAL_MLP_DIM = 16 + +#: Offline OMol asset inside :data:`WEIGHTS_DIR_ENV`: the cueq ``state_dict`` of +#: ``MACE-omol-0-extra-large-1024``. The shipped ``OMOL-cueq.model`` is a +#: *pickled* ``mace.modules.models.ScaleShiftMACE`` — reading it needs ``mace`` +#: and ``e3nn`` imported, which CLAUDE.md forbids and which the toolchain does +#: not install — so it was dumped out of tree into the plain ``name -> tensor`` +#: twin that ``torch.load(weights_only=True)`` reads with no third-party class +#: at all, exactly the shape MatPES already ships +#: (``matpes_r2scan_cueq_state.pt``). The converter sits beside the asset as +#: ``convert_omol_to_cueq_state.py``. +OMOL_WEIGHTS_FILE = "omol_cueq_state.pt" + +#: Hard-coded golden: every ``nn.Parameter`` tensor of the OMol model the +#: official checkpoint has to fill. Pinning the count keeps +#: ``test_official_omol_load_fills_every_learnable`` from passing vacuously on +#: a model that lost half its blocks. +OMOL_PARAMETER_COUNT = 104 + +#: Hard-coded golden: the checkpoint keys with no home in molnex's OMol model. +#: ``OMOL_REMAP`` exists with ``on_unexpected="return"`` because the OMol family +#: *can* ship auxiliary heads this port does not model; the single-head +#: ``omol`` checkpoint actually shipped has none, so the snapshot is empty — +#: and any future key growing into this list is a deliberate decision, not a +#: silent one. +OMOL_UNEXPECTED_KEYS: list[str] = [] + +#: Hard-coded goldens (eV, eV/Å): energy, ``max|F|`` and ``F[0]`` of the +#: ``omol_cluster`` fixture under the official OMol weights, captured at +#: ``OMP_NUM_THREADS=1`` — see :class:`TestOfficialOMolWeights` for the full +#: provenance and :data:`OMOL_ENERGY_TOL` for why the thread count is named. +OMOL_GOLDEN_ENERGY = -3122.454566006894 +OMOL_GOLDEN_MAX_FORCE = 5.397722165018621 +OMOL_GOLDEN_FIRST_FORCE = (-5.397722165018621, -3.9298849841895853, -3.931997122634094) + +#: The repo's *numerical* tolerances (energy 1e-6 eV, force 1e-4 eV/Å), not the +#: exact ones, because this surface is **not** bit-reproducible across thread +#: counts: the CPU reduction order inside the cuEquivariance fallback follows +#: ``torch.get_num_threads()``. Measured here it is bit-identical within a +#: thread count and across repeat processes, but ``{1, 4, 16}`` threads and +#: this node's 48-thread default disagree by 3.2e-08 eV / 9.3e-08 eV/Å. +#: Tightening to the exact pair would make the case pass or fail on how many +#: cores the runner happens to have — a flake, not a stricter lock. +OMOL_ENERGY_TOL = 1e-6 +OMOL_FORCE_TOL = 1e-4 + +_TEST_PACKAGE = Path(__file__).resolve().parent + + +def _official_weights_absent(filename: str) -> bool: + """Whether the gated offline asset ``filename`` is unavailable. + + Args: + filename: File expected inside the :data:`WEIGHTS_DIR_ENV` directory. + + Returns: + ``True`` when the environment variable is unset or the file is missing, + which is the skip condition of the weights-gated cases. + """ + directory = os.environ.get(WEIGHTS_DIR_ENV) + return directory is None or not (Path(directory) / filename).is_file() + + +# -------------------------------------------------------------------------- +# Precondition guard (spec §Tasks): a same-named module next to the package +# makes ``tests.test_molzoo.test_mace`` ambiguous and collection silent. +# -------------------------------------------------------------------------- + + +def test_no_module_shadows_the_mace_test_package() -> None: + """A leftover ``test_mace.py`` next to ``test_mace/`` hides one of them.""" + assert not (_TEST_PACKAGE.parent / "test_mace.py").exists() + + +# -------------------------------------------------------------------------- +# Helpers +# -------------------------------------------------------------------------- + + +def _official_state(model: torch.nn.Module, names: Mapping[str, str]) -> dict[str, torch.Tensor]: + """Rename a model's own ``state_dict`` back into official cueq key names. + + Args: + model: Source model. + names: Molnex prefix → official prefix (longest match wins). + + Returns: + The same tensors under the names an official checkpoint would use. + """ + ordered = sorted(names.items(), key=lambda item: -len(item[0])) + out: dict[str, torch.Tensor] = {} + for key, value in model.state_dict().items(): + for prefix, official in ordered: + if key == prefix or key.startswith(prefix): + key = official + (key[len(prefix) :] if key != prefix else "") + break + out[key] = value + return out + + +def _filled(spec: MACEMatpesSpec | MACEOMolSpec, seed: int) -> MACEPotential: + """A tiny CPU potential whose every parameter is non-zero and seed-specific. + + Two models built from the same spec with different seeds must disagree in + every parameter, or a round-trip assertion could pass on tensors nobody + moved. The OMol readout is zero-initialised at construction + (``molrep.readout.mace._ScalarO3Linear``), so the explicit fill is what + keeps those keys meaningful here. + + Args: + spec: Validated variant configuration. + seed: Seed of both the construction RNG and the fill generator. + + Returns: + A ``MACEPotential`` on the pure-torch cuEquivariance path. + """ + torch.manual_seed(seed) + model = MACEPotential(spec, use_fallback=True) + generator = torch.Generator().manual_seed(seed) + with torch.no_grad(): + for parameter in model.parameters(): + parameter.normal_(generator=generator) + return model + + +def _write_checkpoint( + directory: Path, + spec: MACEMatpesSpec, + model: MACEPotential, + *, + hidden_irreps: str, + mlp_irreps: str, + drop: tuple[str, ...] = (), + switches: Mapping[str, object] | None = None, +) -> tuple[Path, Path]: + """Write an official-shaped MatPES config json + cueq ``state_dict``. + + Args: + directory: Destination directory (``tmp_path``). + spec: The spec ``model`` was built from; supplies the config values. + model: Source of the weights, dumped under official cueq names. + hidden_irreps: Official ``hidden_irreps`` string, e.g. ``"32x0e+32x1o"``. + mlp_irreps: Official ``MLP_irreps`` string, e.g. ``"8x0e"``. + drop: Config keys to leave out, for the missing-key cases. + switches: Extra config entries — the architecture-switch keys + (``pair_repulsion`` / ``distance_transform``) the stock dump + carries. Omitted by default, which is the older-dump case. + + Returns: + ``(config_path, weights_path)``. + """ + config: dict[str, object] = { + "atomic_numbers": spec.atomic_numbers, + "atomic_energies": spec.atomic_energies, + "r_max": spec.r_max, + "num_bessel": spec.num_bessel, + "num_polynomial_cutoff": spec.num_polynomial_cutoff, + "max_ell": spec.l_max, + "correlation": spec.correlation, + "num_interactions": spec.num_interactions, + "hidden_irreps": hidden_irreps, + "MLP_irreps": mlp_irreps, + "radial_MLP": spec.radial_mlp, + "atomic_inter_scale": spec.scale, + "atomic_inter_shift": spec.shift, + } + config.update(switches or {}) + for key in drop: + del config[key] + config_path = directory / "matpes_config.json" + config_path.write_text(json.dumps(config)) + weights_path = directory / "matpes_cueq_state.pt" + torch.save(_official_state(model, MATPES_OFFICIAL_NAMES), weights_path) + return config_path, weights_path + + +class TestCheckpointRemap: + """Test the one remap that both foundation families now share.""" + + def test_rename_is_pure(self) -> None: + """``.rename`` copies, drops graph constants / masks, and renames.""" + weight = torch.zeros(2) + freqs = torch.ones(4) + official = { + "node_embedding.linear.weight": weight, + BESSEL_OFFICIAL_KEY: freqs, + "node_embedding.linear.f.m.graphs.0.graph.c0": torch.zeros(()), + "interactions.0.conv_tp.output_mask": torch.zeros(3), + "interactions.0.linear.weight": weight, + "r_max": torch.tensor(6.0), + } + snapshot = dict(official) + + renamed = MATPES_REMAP.rename(official) + + assert renamed is not official + assert list(official) == list(snapshot) + assert all(official[key] is value for key, value in snapshot.items()) + assert set(renamed) == { + "node_embedding.weight", + "bessel.freqs", + "interactions.0.linear.weight", + } + assert renamed["bessel.freqs"] is freqs + + def test_roundtrip_restores_every_parameter_matpes( + self, tiny_matpes_spec: MACEMatpesSpec + ) -> None: + """MatPES: official names load back bit-for-bit into a fresh model.""" + source = _filled(tiny_matpes_spec, seed=0) + target = _filled(tiny_matpes_spec, seed=1) + + MATPES_REMAP.load(target, _official_state(source, MATPES_OFFICIAL_NAMES)) + + for (name, want), (_, got) in zip( + source.named_parameters(), target.named_parameters(), strict=True + ): + assert torch.equal(want, got), name + + def test_roundtrip_restores_every_parameter_omol(self, tiny_omol_spec: MACEOMolSpec) -> None: + """OMol: same bijection on the family that shipped with no unit test.""" + source = _filled(tiny_omol_spec, seed=0) + target = _filled(tiny_omol_spec, seed=1) + + OMOL_REMAP.load(target, _official_state(source, OMOL_OFFICIAL_NAMES)) + + for (name, want), (_, got) in zip( + source.named_parameters(), target.named_parameters(), strict=True + ): + assert torch.equal(want, got), name + + def test_raise_policy_rejects_unexpected_key(self, tiny_matpes_spec: MACEMatpesSpec) -> None: + """MatPES policy: refuse before touching the model — no half load.""" + source = _filled(tiny_matpes_spec, seed=0) + target = _filled(tiny_matpes_spec, seed=1) + state = _official_state(source, MATPES_OFFICIAL_NAMES) + state[MYSTERY_KEY] = torch.zeros(3, dtype=torch.float64) + before = {name: parameter.detach().clone() for name, parameter in target.named_parameters()} + + with pytest.raises(RuntimeError, match="no home"): + MATPES_REMAP.load(target, state) + + for name, parameter in target.named_parameters(): + assert torch.equal(parameter, before[name]), name + + def test_return_policy_reports_unexpected_key(self, tiny_matpes_spec: MACEMatpesSpec) -> None: + """OMol policy: the same key is reported, not raised.""" + source = _filled(tiny_matpes_spec, seed=0) + target = _filled(tiny_matpes_spec, seed=1) + state = _official_state(source, MATPES_OFFICIAL_NAMES) + state[MYSTERY_KEY] = torch.zeros(3, dtype=torch.float64) + lenient = CheckpointRemap(MATPES_KEY_REMAP, on_unexpected="return") + + _, unexpected = lenient.load(target, state) + + assert unexpected == [MYSTERY_KEY] + + @pytest.mark.parametrize("policy", ["raise", "return"]) + def test_rejects_missing_parameter(self, tiny_matpes_spec: MACEMatpesSpec, policy: str) -> None: + """An unfilled parameter would keep its random init — never load it.""" + source = _filled(tiny_matpes_spec, seed=0) + target = _filled(tiny_matpes_spec, seed=1) + state = _official_state(source, MATPES_OFFICIAL_NAMES) + del state[COVERED_PARAMETER] + + with pytest.raises(RuntimeError, match="not covered") as raised: + CheckpointRemap(MATPES_KEY_REMAP, on_unexpected=policy).load(target, state) + + assert COVERED_PARAMETER in str(raised.value) + + @pytest.mark.parametrize("policy", ["raise", "return"]) + def test_rejects_missing_parameter_without_weight_suffix( + self, tiny_matpes_spec: MACEMatpesSpec, policy: str + ) -> None: + """``bessel.freqs`` is an ``nn.Parameter``; no name heuristic excuses it. + + Regression lock for ``mace_omol.py:437-439``: a ``.weight`` / ``.bias`` + suffix test let this one through, and the official weights differ from + the analytic init by ~2e-7 — straight into the reported parity residual. + """ + source = _filled(tiny_matpes_spec, seed=0) + target = _filled(tiny_matpes_spec, seed=1) + state = _official_state(source, MATPES_OFFICIAL_NAMES) + del state[BESSEL_OFFICIAL_KEY] + + with pytest.raises(RuntimeError, match="not covered") as raised: + CheckpointRemap(MATPES_KEY_REMAP, on_unexpected=policy).load(target, state) + + assert "bessel.freqs" in str(raised.value) + + @pytest.mark.parametrize("policy", ["raise", "return"]) + def test_rejects_shape_mismatch(self, tiny_matpes_spec: MACEMatpesSpec, policy: str) -> None: + """Wrong shapes mean the model was built with the wrong config.""" + source = _filled(tiny_matpes_spec, seed=0) + target = _filled(tiny_matpes_spec, seed=1) + state = _official_state(source, MATPES_OFFICIAL_NAMES) + state[COVERED_PARAMETER] = torch.zeros(1, 7, dtype=torch.float64) + + with pytest.raises(RuntimeError, match="shape mismatch"): + CheckpointRemap(MATPES_KEY_REMAP, on_unexpected=policy).load(target, state) + + @pytest.mark.parametrize("policy", ["raise", "return"]) + def test_accepts_rank_difference_on_frozen_scalars( + self, tiny_matpes_spec: MACEMatpesSpec, policy: str + ) -> None: + """MACE stores ``scale`` as ``(1,)``; molnex holds a 0-d buffer.""" + source = _filled(tiny_matpes_spec, seed=0) + target = _filled(tiny_matpes_spec, seed=1) + state = _official_state(source, MATPES_OFFICIAL_NAMES) + state["scale_shift.scale"] = torch.tensor([FROZEN_SCALAR_VALUE], dtype=torch.float64) + + CheckpointRemap(MATPES_KEY_REMAP, on_unexpected=policy).load(target, state) + + assert target.scale_shift.scale.shape == () + assert float(target.scale_shift.scale) == pytest.approx(FROZEN_SCALAR_VALUE) + + def test_reported_lists_are_sorted(self, tiny_matpes_spec: MACEMatpesSpec) -> None: + """Both returned lists are sorted — a deterministic normalisation. + + The flat OMol loader handed back torch's ``unexpected`` in checkpoint + order; the merged implementation sorts both lists so a diff of two runs + cannot move on dict ordering alone. + """ + source = _filled(tiny_matpes_spec, seed=0) + target = _filled(tiny_matpes_spec, seed=1) + state = _official_state(source, MATPES_OFFICIAL_NAMES) + state["interactions.9.zeta.weight"] = torch.zeros(3, dtype=torch.float64) + state["interactions.0.alpha.weight"] = torch.zeros(3, dtype=torch.float64) + lenient = CheckpointRemap(MATPES_KEY_REMAP, on_unexpected="return") + + missing, unexpected = lenient.load(target, state) + + assert unexpected == ["interactions.0.alpha.weight", "interactions.9.zeta.weight"] + # The graph constants cueq rebuilds are the natural multi-element case; + # torch reports them in module-registration order, which is not sorted. + assert len(missing) > 1 + assert missing == sorted(missing) + assert missing != [key for key in target.state_dict() if key in set(missing)] + + +class TestFromCheckpoint: + """Test ``MACEPotential.from_checkpoint`` — json + weights in one call.""" + + def test_builds_from_config_and_weights( + self, tmp_path: Path, tiny_matpes_spec: MACEMatpesSpec + ) -> None: + """An official-shaped config + checkpoint reproduces the source model.""" + source = _filled(tiny_matpes_spec, seed=0) + config_path, weights_path = _write_checkpoint( + tmp_path, + tiny_matpes_spec, + source, + hidden_irreps="16x0e+16x1o", + mlp_irreps="8x0e", + ) + + built = MACEPotential.from_checkpoint(config_path, weights_path, use_fallback=True) + + for (name, want), (_, got) in zip( + source.named_parameters(), built.named_parameters(), strict=True + ): + assert torch.equal(want, got), name + + def test_derives_dims_from_irreps(self, tmp_path: Path) -> None: + """``num_features`` / ``max_hidden_l`` / ``mlp_dim`` come from the irreps. + + ``run_nve.py:154-159`` hard-codes ``128 / 1 / 16``; a checkpoint with + any other width would then load into the wrong model — or, once the + shapes happen to fit, be quietly wrong. These dimensions (32 / 1 / 8) + share no value with that triple. + """ + spec = MACEMatpesSpec( + **{**TINY_MATPES_KWARGS, "num_features": 32, "max_hidden_l": 1, "mlp_dim": 8}, + atomic_numbers=ATOMIC_NUMBERS, + atomic_energies=ATOMIC_ENERGIES, + ) + config_path, weights_path = _write_checkpoint( + tmp_path, + spec, + _filled(spec, seed=0), + hidden_irreps="32x0e+32x1o", + mlp_irreps="8x0e", + ) + + built = MACEPotential.from_checkpoint(config_path, weights_path, use_fallback=True) + + derived = built._spec + assert isinstance(derived, MACEMatpesSpec) + assert (derived.num_features, derived.max_hidden_l, derived.mlp_dim) == (32, 1, 8) + # Corroborate on the built module, not only on the parsed spec. + assert built.node_embedding.weight.numel() == len(ATOMIC_NUMBERS) * 32 + assert built.readouts[-1].linear_1.weight.numel() == 32 * 8 + assert built.readouts[-1].linear_2.weight.numel() == 8 + + @pytest.mark.parametrize("absent", ["hidden_irreps", "MLP_irreps"]) + def test_rejects_config_missing_irreps( + self, tmp_path: Path, tiny_matpes_spec: MACEMatpesSpec, absent: str + ) -> None: + """No silent fallback to a default width — name the missing key.""" + config_path, weights_path = _write_checkpoint( + tmp_path, + tiny_matpes_spec, + _filled(tiny_matpes_spec, seed=0), + hidden_irreps="16x0e+16x1o", + mlp_irreps="8x0e", + drop=(absent,), + ) + + with pytest.raises((KeyError, ValueError), match=absent): + MACEPotential.from_checkpoint(config_path, weights_path, use_fallback=True) + + def test_tolerates_a_config_without_the_architecture_switches( + self, tmp_path: Path, tiny_matpes_spec: MACEMatpesSpec + ) -> None: + """A dump that names neither switch is loaded on the spec's defaults. + + Absence is not a contradiction: older dumps predate the two keys, and + the spec's ZBL + Agnesi are what a stock MatPES checkpoint was fitted + with anyway. This is the case the whole existing suite (and + :func:`_write_checkpoint`'s default config) runs on, pinned explicitly + so a switch guard cannot be written as "absent means off". + """ + config_path, weights_path = _write_checkpoint( + tmp_path, + tiny_matpes_spec, + _filled(tiny_matpes_spec, seed=0), + hidden_irreps="16x0e+16x1o", + mlp_irreps="8x0e", + ) + assert set(json.loads(config_path.read_text())).isdisjoint(OFFICIAL_SWITCHES) + + built = MACEPotential.from_checkpoint(config_path, weights_path, use_fallback=True) + + assert built._spec.pair_repulsion == "zbl" + assert built._spec.distance_transform == "agnesi" + + def test_accepts_the_switches_the_official_dump_spells( + self, tmp_path: Path, tiny_matpes_spec: MACEMatpesSpec + ) -> None: + """``pair_repulsion: true`` + ``distance_transform: "Agnesi"`` agree. + + The literal spelling of :data:`OFFICIAL_SWITCHES` is what the shipped + ``matpes_r2scan_config.json`` holds, so a case-sensitive comparison + against the spec's lower-case ``"agnesi"`` would reject the very + checkpoint this constructor exists for — and + ``test_official_checkpoint_bit_parity`` would only catch it where the + offline weights are installed. + """ + source = _filled(tiny_matpes_spec, seed=0) + config_path, weights_path = _write_checkpoint( + tmp_path, + tiny_matpes_spec, + source, + hidden_irreps="16x0e+16x1o", + mlp_irreps="8x0e", + switches=OFFICIAL_SWITCHES, + ) + + built = MACEPotential.from_checkpoint(config_path, weights_path, use_fallback=True) + + for (name, want), (_, got) in zip( + source.named_parameters(), built.named_parameters(), strict=True + ): + assert torch.equal(want, got), name + + @pytest.mark.parametrize( + ("switches", "offender", "hard_wired", "spellings"), + [ + ({"pair_repulsion": False}, "pair_repulsion", "zbl", ("False", "false")), + ({"distance_transform": None}, "distance_transform", "agnesi", ("None", "null")), + ( + {"distance_transform": "polynomial"}, + "distance_transform", + "agnesi", + ("polynomial",), + ), + ], + ids=["pair_repulsion_off", "distance_transform_null", "distance_transform_other"], + ) + def test_rejects_a_config_that_contradicts_a_hard_wired_switch( + self, + tmp_path: Path, + tiny_matpes_spec: MACEMatpesSpec, + switches: dict[str, object], + offender: str, + hard_wired: str, + spellings: tuple[str, ...], + ) -> None: + """A config from another MACE family must not build a MatPES model. + + :class:`~molzoo.mace.spec.MACEMatpesSpec` hard-wires ZBL pair repulsion + and the Agnesi distance transform, and the fitted constants of both are + **buffers**, not ``nn.Parameter`` (``trainable=False``). So the strict + remap cannot notice: the surplus block's constants land in the discarded + ``missing_buffers`` list, every learnable is filled, the load returns + happily — and the model computes a short-range repulsion (or transforms + every distance) the checkpoint was never fitted with. That is the + "runs, looks sane, quietly wrong" failure the module docstring is about, + one config key away. + + The message must name the offending key **and** both values, because + the fix is to build the spec by hand — the reader has to know which + switch disagreed and in which direction. + """ + config_path, weights_path = _write_checkpoint( + tmp_path, + tiny_matpes_spec, + _filled(tiny_matpes_spec, seed=0), + hidden_irreps="16x0e+16x1o", + mlp_irreps="8x0e", + switches={**OFFICIAL_SWITCHES, **switches}, + ) + + with pytest.raises(ValueError, match=offender) as raised: + MACEPotential.from_checkpoint(config_path, weights_path, use_fallback=True) + + message = str(raised.value) + assert hard_wired in message, message + assert any(spelling in message for spelling in spellings), message + + +class TestLoadMatpesStateDict: + """Test the compatibility alias ``variants.load_matpes_state_dict``. + + Migrated verbatim (criteria and messages) from + ``tests/test_molzoo/test_mace_matpes.py::TestLoadMatpesStateDict`` by + ``mace-subpackage-restructure-06-wire``. The function lives in + :mod:`molzoo.mace.variants`, but everything it promises — the key dialect + and the three refusals — is :mod:`molzoo.mace.checkpoint`'s doctrine, so + the cases sit next to :class:`TestCheckpointRemap` (spec §5). What they add + over that class is the *forwarding*: an alias that quietly called + ``load_state_dict`` instead would pass none of the three refusals. + """ + + def test_roundtrip_restores_every_parameter(self, matpes_variant: MACEMatpes) -> None: + """A checkpoint written in official names loads back bit-for-bit.""" + state = _official_state(matpes_variant, MATPES_OFFICIAL_NAMES) + torch.manual_seed(1) + fresh = MACEMatpes( + atomic_numbers=list(ATOMIC_NUMBERS), + atomic_energies=torch.tensor(ATOMIC_ENERGIES), + **{**TINY_MATPES_KWARGS, "use_fallback": False}, + ) + assert any( + not torch.equal(want, got) + for (_, want), (_, got) in zip( + matpes_variant.named_parameters(), fresh.named_parameters(), strict=True + ) + ), "the two models start out identical — the round trip would be vacuous" + + load_matpes_state_dict(fresh, state) + + for (name, want), (_, got) in zip( + matpes_variant.named_parameters(), fresh.named_parameters(), strict=True + ): + assert torch.equal(want, got), name + + def test_rejects_unknown_checkpoint_key(self, matpes_variant: MACEMatpes) -> None: + """A key with no home means the mapping is stale — do not load silently.""" + state = _official_state(matpes_variant, MATPES_OFFICIAL_NAMES) + state[MYSTERY_KEY] = torch.zeros(3) + with pytest.raises(RuntimeError, match="no home"): + load_matpes_state_dict(matpes_variant, state) + + def test_rejects_missing_parameter(self, matpes_variant: MACEMatpes) -> None: + """A parameter the checkpoint never fills would keep its random init.""" + state = _official_state(matpes_variant, MATPES_OFFICIAL_NAMES) + del state[COVERED_PARAMETER] + with pytest.raises(RuntimeError, match="not covered"): + load_matpes_state_dict(matpes_variant, state) + + def test_rejects_shape_mismatch(self, matpes_variant: MACEMatpes) -> None: + """Wrong shapes mean the model was built with the wrong config.""" + state = _official_state(matpes_variant, MATPES_OFFICIAL_NAMES) + state[COVERED_PARAMETER] = torch.zeros(1, 7, dtype=torch.float64) + with pytest.raises(RuntimeError, match="shape mismatch"): + load_matpes_state_dict(matpes_variant, state) + + def test_accepts_rank_difference_on_frozen_scalars(self, matpes_variant: MACEMatpes) -> None: + """MACE stores scale/shift as ``(1,)``; molnex holds a 0-d buffer.""" + state = _official_state(matpes_variant, MATPES_OFFICIAL_NAMES) + state["scale_shift.scale"] = torch.tensor([FROZEN_SCALAR_VALUE], dtype=torch.float64) + + load_matpes_state_dict(matpes_variant, state) + + assert float(matpes_variant.scale_shift.scale) == pytest.approx(FROZEN_SCALAR_VALUE) + + +@pytest.mark.skipif( + os.environ.get(WEIGHTS_DIR_ENV) is None, + reason=f"{WEIGHTS_DIR_ENV} unset — the converted official weights are an offline asset", +) +def test_official_checkpoint_bit_parity() -> None: + """The real MatPES checkpoint loads identically both ways (``science``). + + ``from_checkpoint`` against the hand-written "construct + ``MATPES_REMAP``" + path on ``MACE-matpes-r2scan-omat-ft`` (converted offline with + ``mace.cli.convert_e3nn_cueq``; that tool is never imported here). Bit + parity, not a tolerance: the two paths must place the very same tensors. + The oracle arm states the published ``128 / 1 / 16`` dimensions literally, + so ``from_checkpoint`` has to derive them from ``hidden_irreps`` / + ``MLP_irreps`` unaided. + """ + directory = Path(os.environ[WEIGHTS_DIR_ENV]) + config_path = directory / "matpes_r2scan_config.json" + weights_path = directory / "matpes_r2scan_cueq_state.pt" + official = json.loads(config_path.read_text()) + + manual = MACEPotential( + MACEMatpesSpec( + atomic_numbers=official["atomic_numbers"], + atomic_energies=official["atomic_energies"], + r_max=official["r_max"], + num_bessel=official["num_bessel"], + num_polynomial_cutoff=official["num_polynomial_cutoff"], + l_max=official["max_ell"], + num_features=OFFICIAL_NUM_FEATURES, + max_hidden_l=OFFICIAL_MAX_HIDDEN_L, + num_interactions=official["num_interactions"], + correlation=official["correlation"], + mlp_dim=OFFICIAL_MLP_DIM, + radial_mlp=official["radial_MLP"], + scale=official["atomic_inter_scale"], + shift=official["atomic_inter_shift"], + ), + use_fallback=True, + ) + MATPES_REMAP.load(manual, torch.load(weights_path, map_location="cpu", weights_only=True)) + + built = MACEPotential.from_checkpoint(config_path, weights_path, use_fallback=True) + + reference = manual.state_dict() + produced = built.state_dict() + assert set(produced) == set(reference) + for name, want in reference.items(): + assert torch.equal(produced[name], want), name + + +def _official_omol() -> tuple[MACEPotential, dict[str, torch.Tensor]]: + """The full-size OMol model with the official weights, and that state. + + Hyper-parameters are **not** guessed: the four that vary per checkpoint — + element table, ``E0`` table, ``scale``, ``shift`` — are read out of the + checkpoint itself, and every other field is the + :class:`~molzoo.mace.spec.MACEOMolSpec` default, which is exactly what the + deleted ``scripts/omol_port/verify_e2e.py`` did (commit ``b85d12f^``: + ``MACEOMol(atomic_numbers=ztab, atomic_energies=ae, scale=…, shift=…)``). + The checkpoint corroborates every default it can: ``r_max`` 6.0, + ``num_interactions`` 3, ``num_bessel`` 8 (``bessel_weights``), + ``num_polynomial_cutoff`` 5 (``cutoff_fn.p``), ``num_features`` 1024 and + ``charge_classes`` / ``spin_classes`` 201 / 101 (the joint-embedding + tables), ``mlp_dim`` 16 (``readouts.0.linear_1.weight`` = 1024 × 16). + + ``use_fallback=True`` because the fused cuEquivariance kernels need a GPU + and the ops wheel; :class:`~molzoo.mace.variants.MACEOMol` has no such + keyword (it hard-codes the fused path), so the spec is built directly. + + Returns: + The loaded ``MACEPotential`` in eval mode and the official state. + """ + weights_path = Path(os.environ[WEIGHTS_DIR_ENV]) / OMOL_WEIGHTS_FILE + state = torch.load(weights_path, map_location="cpu", weights_only=True) + spec = MACEOMolSpec( + atomic_numbers=state["atomic_numbers"].tolist(), + atomic_energies=state["atomic_energies_fn.atomic_energies"].flatten().tolist(), + scale=float(state["scale_shift.scale"]), + shift=float(state["scale_shift.shift"]), + use_fallback=True, + ) + model = MACEPotential(spec, use_fallback=True) + OMOL_REMAP.load(model, state) + return model.eval(), state + + +@pytest.mark.skipif( + _official_weights_absent(OMOL_WEIGHTS_FILE), + reason=( + f"{WEIGHTS_DIR_ENV} unset or {OMOL_WEIGHTS_FILE} absent — " + "the converted official OMol weights are an offline asset" + ), +) +class TestOfficialOMolWeights: + """The official OMol checkpoint, loaded and evaluated in tree (``science``). + + This class re-homes the role of the deleted ``scripts/omol_port/verify_*.py`` + oracles, which ``src/molzoo/specs/mace_omol.md`` §7.1 still cites for its + 7.0e-7 eV / 4.3e-6 eV/Å parity record (deleted in ``b85d12f``; §7.1's own + Appendix-A entry of 2026-08-09 calls the record unreproducible in-tree). + What it is **not** is a replacement for that record: nothing here compares + against ``mace-torch`` or ``e3nn``, and no claim about upstream parity can + be read out of a green run. It is a **stability lock** — the official + weights load strictly, and the resulting surface is the one measured on + this machine on the date below. A future refactor that shifts the OMol + energy by more than a microelectronvolt has to say so out loud. + + Provenance of the goldens (all captured 2026-08-09 on the machine this + repository is checked out on): + + * weights ``$MOLNEX_MACE_WEIGHTS_DIR/omol_cueq_state.pt``, + sha256 ``074c86154a1d709f…``, derived from ``OMOL-cueq.model`` + (sha256 ``8735a524e99fe80c…``, the ``mace.cli.convert_e3nn_cueq`` twin of + ``MACE-omol-0-extra-large-1024``) by the ``convert_omol_to_cueq_state.py`` + script stored beside it — see :data:`OMOL_WEIGHTS_FILE` for why a + re-dump was needed at all; + * ``torch 2.12.1+cpu``, ``cuequivariance 0.10.0``, CPU, fp64 (the autouse + ``fp64`` fixture), ``use_fallback=True``, ``OMP_NUM_THREADS=1``, no seed + anywhere — the model is fully determined by the checkpoint and the + geometry is a literal, but the thread count is *not* free + (:data:`OMOL_ENERGY_TOL`), and on a many-core runner the forward is two + orders slower than at ``OMP_NUM_THREADS=8`` (48-way oversubscription on + a five-atom graph: ~80 s against ~0.6 s); + * the extraction is corroborated at *value* level, not only by shapes: + ``radial_embedding.bessel_fn.bessel_weights`` sits 2.2120e-07 Å⁻¹ from + the analytic ``nπ/r_max`` init, the fingerprint recorded independently in + :mod:`molzoo.mace.checkpoint`'s docstring and §7.1. + """ + + def test_official_omol_load_fills_every_learnable(self) -> None: + """Every ``nn.Parameter`` is covered; the unhoused keys are the snapshot. + + :meth:`~molzoo.mace.checkpoint.CheckpointRemap.load` already raises on + an unfilled learnable, so the load itself is half the assertion; the + explicit subset check states the doctrine where a reader can see it, + and the count keeps it from holding vacuously on a shrunken model. + """ + model, state = _official_omol() + + renamed = OMOL_REMAP.rename(state) + parameters = {name for name, _ in model.named_parameters()} + + assert len(parameters) == OMOL_PARAMETER_COUNT + assert parameters <= set(renamed) + assert sorted(set(renamed) - set(model.state_dict())) == OMOL_UNEXPECTED_KEYS + + def test_official_omol_energy_and_forces_match_the_goldens( + self, omol_cluster: TensorDict + ) -> None: + """The loaded surface reproduces the captured energy and forces. + + Five atoms (``O H H C H``, all inside OMol's 83-element table) at the + package's literal :data:`~tests.test_molzoo.test_mace.conftest.CLUSTER_POS` + coordinates, neutral closed-shell singlet, every ordered intra-graph + pair as an edge. Goldens are this machine's own output, not an upstream + number — see the class docstring. + """ + model, _ = _official_omol() + + result = model(omol_cluster) + + energy = result["graphs", "energy"] + forces = result["atoms", "forces"] + assert energy.shape == (1,) + assert forces.shape == (len(omol_cluster["atoms", "Z"]), 3) + assert energy.item() == pytest.approx(OMOL_GOLDEN_ENERGY, abs=OMOL_ENERGY_TOL) + assert forces.abs().max().item() == pytest.approx(OMOL_GOLDEN_MAX_FORCE, abs=OMOL_FORCE_TOL) + assert forces[0].tolist() == pytest.approx( + list(OMOL_GOLDEN_FIRST_FORCE), abs=OMOL_FORCE_TOL + ) diff --git a/tests/test_molzoo/test_mace/test_encoder.py b/tests/test_molzoo/test_mace/test_encoder.py new file mode 100644 index 0000000..1371243 --- /dev/null +++ b/tests/test_molzoo/test_mace/test_encoder.py @@ -0,0 +1,236 @@ +"""Tests for molzoo.mace.encoder — the spec-driven MACE backbone. + +The chain gate is ``state_dict`` parity: :class:`~molzoo.mace.encoder.MACEEncoder` +must register exactly the modules the keyword-constructed +:class:`molzoo.mace.variants.MACEMatpes` / :class:`~molzoo.mace.variants.MACEOMol` +register, under exactly the same names, so an official checkpoint keeps loading +without a key rewrite (05-checkpoint) and the variant classes can simply inherit +the backbone. + +Everything else here pins one primitive at a time: the encoder exposes +``validate_elements`` / ``node_attrs`` / ``initial_node_features`` / +``angular_features`` / ``radial_features`` / ``conditioning`` / +``layer_features`` and the caller composes them. +""" + +from __future__ import annotations + +import ast +import inspect +import math +from pathlib import Path +from typing import Any + +import pytest +import torch +from tensordict import TensorDict + +from molzoo.mace.encoder import MACEEncoder +from tests.conftest import rotate_graph, translate_graph +from tests.test_molzoo.test_mace.conftest import ( + CLUSTER_Z, + TINY_MATPES_KWARGS, + TINY_OMOL_KWARGS, +) + +#: Rigid translation applied by the invariance test (Å). +TRANSLATION = [0.31, -0.72, 1.13] + +#: Rotation angles (rad) of the ``Rz(γ) @ Rx(β)`` test rotation. +ROTATION_Z = 0.7 +ROTATION_X = 0.4 + +#: Scalar features are dimensionless activations; fp64 invariance holds to ~1e-13. +INVARIANCE_ATOL = 1e-10 + +#: ``Z = [8, 1, 6, 1]`` against the ascending table ``[1, 6, 8]``. +ONE_HOT_Z = [8, 1, 6, 1] +EXPECTED_ONE_HOT = [ + [0.0, 0.0, 1.0], + [1.0, 0.0, 0.0], + [0.0, 1.0, 0.0], + [1.0, 0.0, 0.0], +] + + +def _rotation_matrix() -> torch.Tensor: + """``Rz(0.7) @ Rx(0.4)`` in fp64 — a fixed, general SO(3) element.""" + cz, sz = math.cos(ROTATION_Z), math.sin(ROTATION_Z) + cx, sx = math.cos(ROTATION_X), math.sin(ROTATION_X) + rz = torch.tensor([[cz, -sz, 0.0], [sz, cz, 0.0], [0.0, 0.0, 1.0]], dtype=torch.float64) + rx = torch.tensor([[1.0, 0.0, 0.0], [0.0, cx, -sx], [0.0, sx, cx]], dtype=torch.float64) + return rz @ rx + + +def _layer_inputs(encoder: MACEEncoder, batch: TensorDict) -> dict[str, Any]: + """Compose the encoder primitives into ``layer_features`` keyword arguments.""" + Z = batch["atoms", "Z"] + edge_index = batch["edges", "edge_index"] + vectors = batch["edges", "edge_diff"] + lengths = batch["edges", "edge_dist"] + + node_attrs = encoder.node_attrs(Z, vectors.dtype) + edge_feats, cutoff = encoder.radial_features(lengths, Z, edge_index) + return { + "node_feats": encoder.initial_node_features(node_attrs), + "node_attrs": node_attrs, + "edge_attrs": encoder.angular_features(vectors), + "edge_feats": edge_feats, + "edge_index": edge_index, + "cutoff": cutoff, + } + + +def _spec_reading_methods(cls: type) -> list[str]: + """Names of ``cls`` methods (other than ``__init__``) that read ``self._spec``.""" + source = Path(inspect.getsourcefile(cls) or "").read_text(encoding="utf-8") + class_def = next( + node + for node in ast.walk(ast.parse(source)) + if isinstance(node, ast.ClassDef) and node.name == cls.__name__ + ) + offenders: list[str] = [] + for node in class_def.body: + if not isinstance(node, ast.FunctionDef) or node.name == "__init__": + continue + for sub in ast.walk(node): + if ( + isinstance(sub, ast.Attribute) + and sub.attr == "_spec" + and isinstance(sub.value, ast.Name) + and sub.value.id == "self" + ): + offenders.append(node.name) + break + return offenders + + +@pytest.fixture +def matpes_encoder(tiny_matpes_spec) -> MACEEncoder: + """``MACEEncoder`` on the tiny MatPES spec.""" + return MACEEncoder(tiny_matpes_spec).eval() + + +@pytest.fixture +def omol_encoder(tiny_omol_spec) -> MACEEncoder: + """``MACEEncoder`` on the tiny OMOL spec.""" + return MACEEncoder(tiny_omol_spec).eval() + + +class TestMACEEncoder: + """Test the configuration-driven MACE backbone.""" + + # -- state_dict parity: the chain gate ----------------------------------- + + def test_state_dict_keys_match_the_matpes_variant(self, matpes_encoder, matpes_variant): + """Same module names as ``MACEMatpes`` — checkpoints load without a rewrite.""" + assert set(matpes_encoder.state_dict()) == set(matpes_variant.state_dict()) + + def test_state_dict_shapes_and_dtypes_match_the_matpes_variant( + self, matpes_encoder, matpes_variant + ): + """Every shared MatPES entry agrees on ``shape`` and ``dtype``.""" + oracle = matpes_variant.state_dict() + own = matpes_encoder.state_dict() + mine = {k: (tuple(v.shape), v.dtype) for k, v in own.items() if k in oracle} + theirs = {k: (tuple(v.shape), v.dtype) for k, v in oracle.items() if k in own} + assert mine == theirs + + def test_state_dict_keys_match_the_omol_variant(self, omol_encoder, omol_variant): + """Same module names as ``MACEOMol`` — the second variant of one backbone.""" + assert set(omol_encoder.state_dict()) == set(omol_variant.state_dict()) + + def test_state_dict_shapes_and_dtypes_match_the_omol_variant(self, omol_encoder, omol_variant): + """Every shared OMOL entry agrees on ``shape`` and ``dtype``.""" + oracle = omol_variant.state_dict() + own = omol_encoder.state_dict() + mine = {k: (tuple(v.shape), v.dtype) for k, v in own.items() if k in oracle} + theirs = {k: (tuple(v.shape), v.dtype) for k, v in oracle.items() if k in own} + assert mine == theirs + + # -- element table primitives -------------------------------------------- + + def test_node_attrs_is_the_hard_coded_one_hot(self, matpes_encoder): + """``Z`` is looked up in the ascending table and one-hot encoded.""" + attrs = matpes_encoder.node_attrs(torch.tensor(ONE_HOT_Z, dtype=torch.long), torch.float64) + expected = torch.tensor(EXPECTED_ONE_HOT, dtype=torch.float64) + assert torch.equal(attrs, expected) + + def test_node_attrs_honours_the_requested_dtype(self, matpes_encoder): + """The dtype argument is the caller's, not the module's.""" + attrs = matpes_encoder.node_attrs(torch.tensor(ONE_HOT_Z, dtype=torch.long), torch.float32) + assert attrs.dtype == torch.float32 + + def test_validate_elements_accepts_the_table(self, matpes_encoder): + """Atomic numbers inside the table pass silently.""" + assert matpes_encoder.validate_elements(torch.tensor([1, 6, 8, 1])) is None + + def test_validate_elements_lists_the_out_of_table_numbers(self, matpes_encoder): + """An unknown element must name itself — ``searchsorted`` would snap it.""" + with pytest.raises(ValueError, match=r"\[2, 7\]"): + matpes_encoder.validate_elements(torch.tensor([1, 7, 2, 6])) + + # -- conditioning --------------------------------------------------------- + + def test_conditioning_rejects_an_unconditioned_variant(self, matpes_encoder, cluster): + """MatPES has no charge/spin embedding — asking for one is an error.""" + with pytest.raises(ValueError): + matpes_encoder.conditioning( + cluster["atoms", "batch"], + total_spin=torch.tensor([1], dtype=torch.long), + total_charge=torch.tensor([0], dtype=torch.long), + ) + + def test_conditioning_returns_one_row_per_atom(self, omol_encoder, cluster): + """The OMOL joint embedding broadcasts per-graph scalars onto atoms.""" + conditioned = omol_encoder.conditioning( + cluster["atoms", "batch"], + total_spin=torch.tensor([1], dtype=torch.long), + total_charge=torch.tensor([0], dtype=torch.long), + ) + assert conditioned.shape == (len(CLUSTER_Z), TINY_OMOL_KWARGS["num_features"]) + + # -- layer_features ------------------------------------------------------- + + def test_layer_features_returns_one_tensor_per_interaction(self, matpes_encoder, cluster): + """A list (not a stacked tensor): MatPES layer widths differ.""" + feats = matpes_encoder.layer_features(**_layer_inputs(matpes_encoder, cluster)) + assert isinstance(feats, list) + assert len(feats) == TINY_MATPES_KWARGS["num_interactions"] + + def test_layer_features_rows_match_the_atom_count(self, matpes_encoder, cluster): + """Every layer carries one row per atom, at the configured precision.""" + feats = matpes_encoder.layer_features(**_layer_inputs(matpes_encoder, cluster)) + assert all(f.shape[0] == len(CLUSTER_Z) for f in feats) + assert all(f.dtype == torch.float64 for f in feats) + + def test_last_layer_features_are_translation_invariant(self, matpes_encoder, cluster): + """The last layer is pure-scalar irreps — a rigid shift cannot move it.""" + moved = translate_graph(cluster, torch.tensor(TRANSLATION, dtype=torch.float64)) + before = matpes_encoder.layer_features(**_layer_inputs(matpes_encoder, cluster))[-1] + after = matpes_encoder.layer_features(**_layer_inputs(matpes_encoder, moved))[-1] + assert torch.allclose(before, after, atol=INVARIANCE_ATOL, rtol=0.0) + + def test_last_layer_features_are_rotation_invariant_with_shifts( + self, matpes_encoder, periodic_cluster + ): + """``S = n · h`` rotates with the box, so scalar features stay put.""" + turned = rotate_graph(periodic_cluster, _rotation_matrix()) + before = matpes_encoder.layer_features(**_layer_inputs(matpes_encoder, periodic_cluster))[ + -1 + ] + after = matpes_encoder.layer_features(**_layer_inputs(matpes_encoder, turned))[-1] + assert torch.allclose(before, after, atol=INVARIANCE_ATOL, rtol=0.0) + + # -- hot-path discipline --------------------------------------------------- + + def test_primitives_run_without_the_spec_attribute(self, matpes_encoder, cluster): + """``_spec`` is provenance only; deleting it must not break the model.""" + del matpes_encoder._spec + feats = matpes_encoder.layer_features(**_layer_inputs(matpes_encoder, cluster)) + assert len(feats) == TINY_MATPES_KWARGS["num_interactions"] + + def test_primitive_methods_never_read_the_spec(self): + """Only ``__init__`` may touch ``self._spec`` — pydantic attribute reads + on a hot path are a dynamo graph break (cf. mace_matpes.py:118).""" + assert _spec_reading_methods(MACEEncoder) == [] diff --git a/tests/test_molzoo/test_mace/test_geometry.py b/tests/test_molzoo/test_mace/test_geometry.py new file mode 100644 index 0000000..a09b8d8 --- /dev/null +++ b/tests/test_molzoo/test_mace/test_geometry.py @@ -0,0 +1,150 @@ +"""Tests for molzoo.mace.geometry — MACE's edge displacement / length owner. + +MACE's edge vector is ``r_ij = pos[target] - pos[source] + S_ij`` with an +*additive* PBC shift ``S_ij = n_ij · h``. This is a different mathematical +object from PiNet's minimum-image straight-through ``edge_bond_diff`` +(``src/molzoo/pinet/geometry.py:28-49``), which is why MACE owns its own +geometry module — see that module's docstring for the rationale. + +All expectations are hard-coded from three literal coordinates; nothing here +is derived by re-running the function under test. +""" + +from __future__ import annotations + +import pytest +import torch + +from molzoo.mace.geometry import edge_lengths, edge_vectors + +#: Three atoms at literal coordinates (Å). +POS = [[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 2.0, 0.0]] + +#: Four directed edges, ``[:, 0]`` = source, ``[:, 1]`` = target. +EDGE_INDEX = [[0, 1], [1, 0], [0, 2], [2, 1]] + +#: ``pos[target] - pos[source]``, by hand. +EXPECTED_VECTORS = [[1.0, 0.0, 0.0], [-1.0, 0.0, 0.0], [0.0, 2.0, 0.0], [1.0, -2.0, 0.0]] + +#: ``‖r_ij‖`` for the four edges above; ``sqrt(5) = 2.23606797749979``. +EXPECTED_LENGTHS = [1.0, 1.0, 2.0, 2.23606797749979] + +#: Integer images ``n_ij`` for a 3.0 Å cubic box, i.e. ``S = n · 3.0``. +UNIT_SHIFTS = [[1, 0, 0], [-1, 0, 0], [0, 0, 0], [0, 1, 0]] +BOX_LENGTH = 3.0 + +#: ``pos[target] - pos[source] + S_ij``, by hand. +EXPECTED_SHIFTED_VECTORS = [ + [4.0, 0.0, 0.0], + [-4.0, 0.0, 0.0], + [0.0, 2.0, 0.0], + [1.0, 1.0, 0.0], +] + +#: ``‖r_ij + S_ij‖``; ``sqrt(2) = 1.4142135623730951``. +EXPECTED_SHIFTED_LENGTHS = [4.0, 4.0, 2.0, 1.4142135623730951] + +#: ``d(Σ r_ij)/d pos`` = (in-degree − out-degree) per atom, per component. +#: Atom 0 is a target once and a source twice, atom 1 twice/once, atom 2 once/once. +EXPECTED_VECTOR_SUM_GRAD = [[-1.0, -1.0, -1.0], [1.0, 1.0, 1.0], [0.0, 0.0, 0.0]] + +#: ``d(Σ ‖r_ij‖)/d pos`` = Σ ±r̂; ``1/sqrt(5) = 0.4472135954999579``. +EXPECTED_LENGTH_SUM_GRAD = [ + [-2.0, -1.0, 0.0], + [2.4472135954999579, -0.8944271909999159, 0.0], + [-0.4472135954999579, 1.8944271909999159, 0.0], +] + +#: Positions are exact to 1e-12; these are fp64 sums of a handful of terms. +POSITION_ATOL = 1e-12 + + +@pytest.fixture +def pos() -> torch.Tensor: + """``(3, 3)`` fp64 positions in Å.""" + return torch.tensor(POS, dtype=torch.float64) + + +@pytest.fixture +def edge_index() -> torch.Tensor: + """``(4, 2)`` source/target index pairs.""" + return torch.tensor(EDGE_INDEX, dtype=torch.long) + + +@pytest.fixture +def shifts() -> torch.Tensor: + """``(4, 3)`` fp64 PBC shift vectors for a 3.0 Å cubic box.""" + return torch.tensor(UNIT_SHIFTS, dtype=torch.float64) * BOX_LENGTH + + +class TestEdgeVectors: + """Test ``edge_vectors(pos, edge_index, shifts=None)``.""" + + def test_matches_hand_computed_displacements(self, pos, edge_index): + """``r_ij = pos[target] - pos[source]`` — the repo-wide edge convention.""" + vectors = edge_vectors(pos, edge_index) + expected = torch.tensor(EXPECTED_VECTORS, dtype=torch.float64) + assert torch.allclose(vectors, expected, atol=POSITION_ATOL, rtol=0.0) + + def test_output_shape_is_one_vector_per_edge(self, pos, edge_index): + """``(E, 3)`` regardless of the atom count.""" + assert edge_vectors(pos, edge_index).shape == (len(EDGE_INDEX), 3) + + def test_shifts_are_added_to_the_displacement(self, pos, edge_index, shifts): + """``S_ij`` enters additively — not as a minimum-image wrap.""" + vectors = edge_vectors(pos, edge_index, shifts=shifts) + expected = torch.tensor(EXPECTED_SHIFTED_VECTORS, dtype=torch.float64) + assert torch.allclose(vectors, expected, atol=POSITION_ATOL, rtol=0.0) + + def test_shifted_result_equals_shift_free_plus_shifts(self, pos, edge_index, shifts): + """The additive law, stated directly against the shift-free call.""" + shifted = edge_vectors(pos, edge_index, shifts=shifts) + assert torch.allclose( + shifted, edge_vectors(pos, edge_index) + shifts, atol=POSITION_ATOL, rtol=0.0 + ) + + def test_position_gradient_matches_the_edge_incidence(self, pos, edge_index): + """``∂r/∂pos`` is the signed incidence matrix of the directed graph.""" + leaf = pos.clone().requires_grad_(True) + (grad,) = torch.autograd.grad(edge_vectors(leaf, edge_index).sum(), leaf) + expected = torch.tensor(EXPECTED_VECTOR_SUM_GRAD, dtype=torch.float64) + assert torch.allclose(grad, expected, atol=POSITION_ATOL, rtol=0.0) + + def test_position_gradient_is_unchanged_by_shifts(self, pos, edge_index, shifts): + """``S_ij`` is constant w.r.t. ``pos``, so it drops out of the gradient.""" + leaf = pos.clone().requires_grad_(True) + (grad,) = torch.autograd.grad(edge_vectors(leaf, edge_index, shifts=shifts).sum(), leaf) + expected = torch.tensor(EXPECTED_VECTOR_SUM_GRAD, dtype=torch.float64) + assert torch.allclose(grad, expected, atol=POSITION_ATOL, rtol=0.0) + + +class TestEdgeLengths: + """Test ``edge_lengths(vectors, *, keepdim=False)``.""" + + def test_matches_hand_computed_distances(self, pos, edge_index): + """``d_ij = ‖r_ij‖`` for the four literal edges.""" + lengths = edge_lengths(edge_vectors(pos, edge_index)) + expected = torch.tensor(EXPECTED_LENGTHS, dtype=torch.float64) + assert torch.allclose(lengths, expected, atol=POSITION_ATOL, rtol=0.0) + + def test_default_shape_is_one_scalar_per_edge(self, pos, edge_index): + """``keepdim=False`` (the default) gives ``(E,)`` — MatPES's shape.""" + assert edge_lengths(edge_vectors(pos, edge_index)).shape == (len(EDGE_INDEX),) + + def test_keepdim_shape_is_a_trailing_singleton(self, pos, edge_index): + """``keepdim=True`` gives ``(E, 1)`` — OMOL's shape.""" + lengths = edge_lengths(edge_vectors(pos, edge_index), keepdim=True) + assert lengths.shape == (len(EDGE_INDEX), 1) + + def test_shifted_lengths_are_the_minimum_image_distances(self, pos, edge_index, shifts): + """With ``S_ij`` folded in, the length is the imaged distance.""" + lengths = edge_lengths(edge_vectors(pos, edge_index, shifts=shifts)) + expected = torch.tensor(EXPECTED_SHIFTED_LENGTHS, dtype=torch.float64) + assert torch.allclose(lengths, expected, atol=POSITION_ATOL, rtol=0.0) + + def test_position_gradient_is_the_analytic_unit_vectors(self, pos, edge_index): + """``∂‖r‖/∂pos`` accumulates ``+r̂`` on the target and ``−r̂`` on the source.""" + leaf = pos.clone().requires_grad_(True) + (grad,) = torch.autograd.grad(edge_lengths(edge_vectors(leaf, edge_index)).sum(), leaf) + expected = torch.tensor(EXPECTED_LENGTH_SUM_GRAD, dtype=torch.float64) + assert torch.allclose(grad, expected, atol=POSITION_ATOL, rtol=0.0) diff --git a/tests/test_molzoo/test_mace/test_potential.py b/tests/test_molzoo/test_mace/test_potential.py new file mode 100644 index 0000000..8be614f --- /dev/null +++ b/tests/test_molzoo/test_mace/test_potential.py @@ -0,0 +1,747 @@ +"""Tests for molzoo.mace.potential — the unified MACE energy/force potential. + +:class:`~molzoo.mace.potential.MACEPotential` merges the two pre-cutover flat +foundation models — now the keyword aliases +:class:`molzoo.mace.variants.MACEMatpes` / +:class:`~molzoo.mace.variants.MACEOMol` — into one spec-driven class with two +seams: + +* ``energy_core(...) -> (B,)`` — flat, public, compilable (1 dynamo graph); +* ``_write_energy(batch) -> batch`` — the ``molpot.derivation.protocol`` + hook that ``call_energy`` dispatches to. + +Every branch (forces on/off, charge/spin conditioning) is resolved in +``__init__`` into ``self._pipeline``; ``forward`` is a one-line dispatch and no +public signature carries ``compute_forces``. + +The two keyword variants are imported here as **parity oracles**: the numbers +must not move when the code moves house. Since the cutover +(``mace-subpackage-restructure-06-wire``) they are subclasses of the class under +test, so the parity is no longer "old module vs new module" but the two seams a +caller actually has — ``MACEPotential.forward`` (batch schema, shared force +kernel, own position leaf policy) against ``MACEMatpes.energy_forces`` / +``MACEOMol.energy_forces`` (raw tensors, the variants' own leaf). Those are two +genuinely different code paths on one energy, and +``regressions/mace-subpackage-restructure-04-potential.py`` freezes the same +numbers as hard-coded literals with no import of either. + +Every model is tiny (2 layers, 16 channels, ``l_max=1``), fp64 (autouse +``fp64`` fixture in ``conftest.py``), CPU, seeded — see +``.claude/specs/mace-subpackage-restructure-04-potential.md`` §Testing strategy. +""" + +from __future__ import annotations + +import ast +import inspect +import textwrap +from collections.abc import Callable +from pathlib import Path + +import pytest +import torch +from tensordict import TensorDict + +from molpot.derivation.protocol import call_energy +from molzoo.mace.encoder import MACEEncoder +from molzoo.mace.potential import MACEPotential +from molzoo.mace.spec import MACEMatpesSpec, MACEOMolSpec +from molzoo.mace.variants import MACEMatpes, MACEOMol +from tests.conftest import make_graph_batch, translate_graph +from tests.test_molzoo.test_mace.conftest import ( + CLUSTER_POS, + CLUSTER_Z, + PAIR_Z, + SINGLE_GRAPH, + full_edge_index, + raw_tensors, +) + +#: Iron (26) is outside the ``[1, 6, 8]`` table — ``searchsorted`` would snap it. +OFF_TABLE_Z = [8, 1, 1, 6, 26] + +#: Rigid translation of the whole system (Å); the energy must not notice. +TRANSLATION = [1.0, 0.0, 0.0] + +#: Re-homing the code must not move a bit: fp64 parity between the batch path +#: and the variants' raw-tensor ``energy_forces`` (spec §Domain basis — +#: measured bitwise equal at the same operator order). +ENERGY_PARITY_ATOL = 1e-12 # eV +FORCE_PARITY_ATOL = 1e-12 # eV/Å + +#: Newton's third law on an isolated cluster (measured 1.4e-17 on this model). +NET_FORCE_ATOL = 1e-10 # eV/Å + +#: Central-difference step and tolerance (measured worst error 2.8e-7). +FD_STEP = 1e-4 # Å +FD_ATOL = 1e-6 # eV/Å + +#: Atom/axis pairs probed by the finite-difference test. +FD_PROBES = ((0, 0), (2, 1), (4, 2)) + +#: Hard-coded golden: the parameter names of the pre-cutover flat energy core +#: (``MACEMatpes._compute_energy``, ``mace_matpes.py:229-237`` at 0e05959). +#: ``energy_core`` must keep them so 07's call-site re-point is a pure rename. +FLAT_ENERGY_CORE_PARAMETERS = ("self", "positions", "Z", "edge_index", "batch", "num_graphs") + +#: The two per-graph conditioning tensors only the OMOL variant consumes. +CONDITIONING_PARAMETERS = ("shifts", "total_charge", "total_spin") + +#: A third force path in this file is forbidden (CLAUDE.md "Force derivation"). +FORBIDDEN_GRAD_CALLS = frozenset( + {"torch.autograd.grad", "autograd.grad", "torch.func.grad", "func.grad", "grad"} +) + + +# -------------------------------------------------------------------------- +# Precondition guard (spec §Testing strategy): the package must not be shadowed +# by a same-named module, or the whole directory is silently skipped. +# -------------------------------------------------------------------------- + +_TEST_PACKAGE = Path(__file__).resolve().parent + + +def test_mace_test_directory_is_a_package() -> None: + """``tests/test_molzoo/test_mace/`` must carry an ``__init__.py``.""" + assert (_TEST_PACKAGE / "__init__.py").is_file() + + +def test_no_module_shadows_the_mace_test_package() -> None: + """A leftover ``test_mace.py`` next to ``test_mace/`` hides one of them.""" + assert not (_TEST_PACKAGE.parent / "test_mace.py").exists() + + +# -------------------------------------------------------------------------- +# Helpers +# -------------------------------------------------------------------------- + + +def _transfer(potential: MACEPotential, variant: torch.nn.Module) -> MACEPotential: + """Move the oracle's weights into the potential by a strict key match.""" + potential.load_state_dict(variant.state_dict(), strict=True) + return potential + + +def _dotted_name(node: ast.expr) -> str: + """``ast.Attribute``/``ast.Name`` chain as ``"torch.autograd.grad"``.""" + parts: list[str] = [] + while isinstance(node, ast.Attribute): + parts.append(node.attr) + node = node.value + if isinstance(node, ast.Name): + parts.append(node.id) + return ".".join(reversed(parts)) + + +def _called_names(tree: ast.AST) -> set[str]: + """Every dotted callee name appearing in ``tree``.""" + return {_dotted_name(node.func) for node in ast.walk(tree) if isinstance(node, ast.Call)} + + +def _potential_module_tree() -> ast.Module: + """Parse ``src/molzoo/mace/potential.py``.""" + source = Path(inspect.getsourcefile(MACEPotential) or "").read_text(encoding="utf-8") + return ast.parse(source) + + +def _function_body(function: Callable[..., object]) -> list[ast.stmt]: + """Statements of ``function``, with a leading docstring dropped.""" + tree = ast.parse(textwrap.dedent(inspect.getsource(function))) + definition = tree.body[0] + assert isinstance(definition, ast.FunctionDef) + body = definition.body + if body and isinstance(body[0], ast.Expr) and isinstance(body[0].value, ast.Constant): + body = body[1:] + return body + + +# -------------------------------------------------------------------------- +# Fixtures +# -------------------------------------------------------------------------- + + +@pytest.fixture +def leaf_cluster() -> TensorDict: + """The cluster whose ``atoms.pos`` is already a live ``requires_grad`` leaf.""" + pos = torch.tensor(CLUSTER_POS, dtype=torch.float64).requires_grad_(True) + return make_graph_batch( + pos=pos, + Z=torch.tensor(CLUSTER_Z, dtype=torch.long), + edge_index=full_edge_index(SINGLE_GRAPH), + batch=torch.zeros(len(CLUSTER_Z), dtype=torch.long), + ) + + +@pytest.fixture +def branch_cluster() -> TensorDict: + """The cluster whose ``atoms.pos`` requires grad but is **not** a leaf.""" + pos = torch.tensor(CLUSTER_POS, dtype=torch.float64).requires_grad_(True) * 1.0 + return make_graph_batch( + pos=pos, + Z=torch.tensor(CLUSTER_Z, dtype=torch.long), + edge_index=full_edge_index(SINGLE_GRAPH), + batch=torch.zeros(len(CLUSTER_Z), dtype=torch.long), + ) + + +@pytest.fixture +def matpes_potential(tiny_matpes_spec: MACEMatpesSpec) -> MACEPotential: + """Energy + force MatPES potential on the tiny spec.""" + torch.manual_seed(0) + return MACEPotential(tiny_matpes_spec).eval() + + +@pytest.fixture +def energy_only_potential(tiny_matpes_spec: MACEMatpesSpec) -> MACEPotential: + """Energy-only MatPES potential — ``compute_forces=False`` is a *construction*.""" + torch.manual_seed(0) + return MACEPotential(tiny_matpes_spec, compute_forces=False).eval() + + +@pytest.fixture +def omol_potential(tiny_omol_spec: MACEOMolSpec) -> MACEPotential: + """Energy + force OMOL potential on the tiny spec.""" + torch.manual_seed(0) + return MACEPotential(tiny_omol_spec).eval() + + +@pytest.fixture +def matpes_parity( + tiny_matpes_spec: MACEMatpesSpec, matpes_variant: MACEMatpes +) -> tuple[MACEPotential, MACEMatpes]: + """MatPES potential holding the oracle variant's weights, plus the oracle.""" + return _transfer(MACEPotential(tiny_matpes_spec).eval(), matpes_variant), matpes_variant + + +@pytest.fixture +def omol_parity( + tiny_omol_spec: MACEOMolSpec, omol_variant: MACEOMol +) -> tuple[MACEPotential, MACEOMol]: + """OMOL potential holding the oracle variant's weights, plus the oracle.""" + return _transfer(MACEPotential(tiny_omol_spec).eval(), omol_variant), omol_variant + + +class TestMACEPotential: + """Test the unified MACE energy/force potential.""" + + # -- construction --------------------------------------------------------- + + def test_inherits_the_encoder_so_state_dict_keys_stay_flat(self) -> None: + """Holding an encoder would prefix every key with ``encoder.`` (ac-007).""" + assert issubclass(MACEPotential, MACEEncoder) + + def test_constructor_takes_a_spec_and_two_keyword_switches(self) -> None: + """``MACEPotential(spec, *, compute_forces=True, use_fallback=False)``.""" + parameters = inspect.signature(MACEPotential.__init__).parameters + assert list(parameters) == ["self", "spec", "compute_forces", "use_fallback"] + assert parameters["spec"].kind is inspect.Parameter.POSITIONAL_OR_KEYWORD + assert parameters["compute_forces"].kind is inspect.Parameter.KEYWORD_ONLY + assert parameters["use_fallback"].kind is inspect.Parameter.KEYWORD_ONLY + + def test_forces_are_computed_by_default(self) -> None: + """MD / evaluation is the main use — unlike ``PiNetPotential``.""" + parameters = inspect.signature(MACEPotential.__init__).parameters + assert parameters["compute_forces"].default is True + assert parameters["use_fallback"].default is False + + def test_the_spec_decides_the_cuequivariance_path( + self, matpes_potential: MACEPotential + ) -> None: + """``spec.use_fallback=True`` (the tiny CPU spec) builds pure-torch blocks.""" + assert matpes_potential.interactions[0].use_fallback is True + assert matpes_potential.products[0].use_fallback is True + + def test_the_constructor_can_escalate_a_fused_spec_to_the_fallback( + self, tiny_omol_spec: MACEOMolSpec + ) -> None: + """``use_fallback=True`` at the call site wins over a fused spec. + + The knob is escalation-only: the ``False`` default never downgrades a + spec that asked for the pure-torch path (a ``bool`` cannot express + "not given"), while a CPU / test call site can always force it on. + """ + assert tiny_omol_spec.use_fallback is False + potential = MACEPotential(tiny_omol_spec, use_fallback=True) + assert potential.interactions[0].use_fallback is True + assert potential.products[0].use_fallback is True + + # -- state_dict parity: the chain gate ------------------------------------ + + def test_state_dict_keys_match_the_matpes_variant( + self, matpes_potential: MACEPotential, matpes_variant: MACEMatpes + ) -> None: + """An official MatPES checkpoint must keep loading without a key rewrite.""" + assert set(matpes_potential.state_dict()) == set(matpes_variant.state_dict()) + + def test_state_dict_shapes_match_the_matpes_variant( + self, matpes_potential: MACEPotential, matpes_variant: MACEMatpes + ) -> None: + """Same names *and* same shapes, or ``strict=True`` transfer is a lie.""" + oracle = {k: tuple(v.shape) for k, v in matpes_variant.state_dict().items()} + own = {k: tuple(v.shape) for k, v in matpes_potential.state_dict().items()} + assert own == oracle + + def test_state_dict_keys_match_the_omol_variant( + self, omol_potential: MACEPotential, omol_variant: MACEOMol + ) -> None: + """The OMOL arm of the same gate.""" + assert set(omol_potential.state_dict()) == set(omol_variant.state_dict()) + + def test_state_dict_shapes_match_the_omol_variant( + self, omol_potential: MACEPotential, omol_variant: MACEOMol + ) -> None: + """The OMOL arm of the same gate, shapes.""" + oracle = {k: tuple(v.shape) for k, v in omol_variant.state_dict().items()} + own = {k: tuple(v.shape) for k, v in omol_potential.state_dict().items()} + assert own == oracle + + def test_loads_the_matpes_variant_state_dict_strictly( + self, matpes_potential: MACEPotential, matpes_variant: MACEMatpes + ) -> None: + """No missing, no unexpected: a direct hand-over, not a fuzzy remap.""" + report = matpes_potential.load_state_dict(matpes_variant.state_dict(), strict=True) + assert (list(report.missing_keys), list(report.unexpected_keys)) == ([], []) + + def test_loads_the_omol_variant_state_dict_strictly( + self, omol_potential: MACEPotential, omol_variant: MACEOMol + ) -> None: + """The OMOL arm of the strict hand-over.""" + report = omol_potential.load_state_dict(omol_variant.state_dict(), strict=True) + assert (list(report.missing_keys), list(report.unexpected_keys)) == ([], []) + + # -- forward: the molix.md.forcefield contract ---------------------------- + + def test_forward_returns_the_same_batch_object( + self, matpes_potential: MACEPotential, cluster: TensorDict + ) -> None: + """Writes are in place — ``PotentialForceField`` keeps its own handle.""" + assert matpes_potential(cluster) is cluster + + def test_forward_writes_one_energy_per_graph( + self, matpes_potential: MACEPotential, pair_batch: TensorDict + ) -> None: + """``graphs.energy`` is ``(B,)`` in eV.""" + out = matpes_potential(pair_batch) + assert out["graphs", "energy"].shape == (2,) + + def test_forward_writes_one_force_vector_per_atom( + self, matpes_potential: MACEPotential, pair_batch: TensorDict + ) -> None: + """``atoms.forces`` is ``(N, 3)`` in eV/Å.""" + out = matpes_potential(pair_batch) + assert out["atoms", "forces"].shape == (len(PAIR_Z), 3) + + def test_forward_builds_the_graphs_namespace_with_the_batch_size( + self, matpes_potential: MACEPotential, pair_batch: TensorDict + ) -> None: + """A missing ``graphs`` must come back as ``batch_size=[B]``, never ``[]``. + + Since 07-cleanup ``write_energy`` sizes the namespace itself, but the + potential still creates it *before* calling ``write_energy`` so the + schema holds even for callers that bypass the protocol helpers. + """ + del pair_batch["graphs"] + out = matpes_potential(pair_batch) + assert out["graphs"].batch_size == torch.Size([2]) + + def test_energy_only_forward_also_builds_graphs_with_the_batch_size( + self, energy_only_potential: MACEPotential, pair_batch: TensorDict + ) -> None: + """The energy-only pipeline is ``_write_energy`` — same schema duty.""" + del pair_batch["graphs"] + out = energy_only_potential(pair_batch) + assert out["graphs"].batch_size == torch.Size([2]) + + def test_forward_consumes_the_periodic_shifts( + self, matpes_potential: MACEPotential, periodic_cluster: TensorDict + ) -> None: + """``edges.shifts`` must reach the edge vectors, not be dropped.""" + pos, Z, edge_index, batch, num_graphs = raw_tensors(periodic_cluster) + shifts = periodic_cluster["edges", "shifts"] + with torch.no_grad(): + expected = matpes_potential.energy_core( + pos, Z, edge_index, batch, num_graphs, shifts=shifts + ) + out = matpes_potential(periodic_cluster) + assert torch.allclose(out["graphs", "energy"], expected, atol=ENERGY_PARITY_ATOL, rtol=0.0) + + def test_energy_only_instance_writes_no_forces( + self, energy_only_potential: MACEPotential, cluster: TensorDict + ) -> None: + """ "Energy only" is a constructed object, and it must not pay for forces.""" + out = energy_only_potential(cluster) + assert ("atoms", "forces") not in out.keys(include_nested=True) + + def test_energy_only_instance_still_writes_the_energy( + self, energy_only_potential: MACEPotential, cluster: TensorDict + ) -> None: + """The energy-only pipeline is still a full energy pipeline.""" + out = energy_only_potential(cluster) + assert out["graphs", "energy"].shape == (1,) + + # -- needs_leaf policy ---------------------------------------------------- + + def test_detached_positions_yield_a_detached_energy( + self, matpes_potential: MACEPotential, cluster: TensorDict + ) -> None: + """The potential owns the leaf, so nobody downstream can lose a graph.""" + out = matpes_potential(cluster) + assert out["graphs", "energy"].requires_grad is False + + def test_non_leaf_positions_yield_a_detached_energy( + self, matpes_potential: MACEPotential, branch_cluster: TensorDict + ) -> None: + """``needs_leaf = not (requires_grad and is_leaf)`` — a branch is not a leaf.""" + out = matpes_potential(branch_cluster) + assert out["graphs", "energy"].requires_grad is False + + def test_grad_leaf_positions_keep_the_energy_attached( + self, matpes_potential: MACEPotential, leaf_cluster: TensorDict + ) -> None: + """Detaching here would cut the training path at the first potential.""" + out = matpes_potential(leaf_cluster) + assert out["graphs", "energy"].requires_grad is True + + def test_grad_leaf_energy_loss_reaches_the_parameters( + self, matpes_potential: MACEPotential, leaf_cluster: TensorDict + ) -> None: + """The whole point of not detaching: ``loss.backward()`` trains the model.""" + out = matpes_potential(leaf_cluster) + out["graphs", "energy"].sum().backward() + weight = matpes_potential.readouts[0].linear.weight + assert weight.grad is not None and torch.any(weight.grad != 0) + + # -- energy_core: the flat compile seam ----------------------------------- + + def test_energy_core_keeps_the_flat_parameter_names(self) -> None: + """07's call-site re-point must be a pure rename of ``_compute_energy``.""" + parameters = tuple(inspect.signature(MACEPotential.energy_core).parameters) + assert parameters[: len(FLAT_ENERGY_CORE_PARAMETERS)] == FLAT_ENERGY_CORE_PARAMETERS + + def test_energy_core_takes_the_conditioning_tensors_last(self) -> None: + """``shifts`` then the two per-graph OMOL tensors, all optional.""" + parameters = inspect.signature(MACEPotential.energy_core).parameters + assert tuple(parameters)[len(FLAT_ENERGY_CORE_PARAMETERS) :] == CONDITIONING_PARAMETERS + assert all(parameters[name].default is None for name in CONDITIONING_PARAMETERS) + + def test_energy_core_first_six_parameters_are_positional(self) -> None: + """``run_nve.py`` calls it positionally on the MD hot path.""" + parameters = inspect.signature(MACEPotential.energy_core).parameters + assert all( + parameters[name].kind is inspect.Parameter.POSITIONAL_OR_KEYWORD + for name in FLAT_ENERGY_CORE_PARAMETERS[1:] + ) + + def test_energy_core_returns_one_energy_per_graph( + self, matpes_potential: MACEPotential, pair_batch: TensorDict + ) -> None: + """Flat tensors in, ``(B,)`` eV out — no TensorDict on this seam.""" + with torch.no_grad(): + energy = matpes_potential.energy_core(*raw_tensors(pair_batch)) + assert energy.shape == (2,) + + def test_energy_core_agrees_with_the_forward_energy( + self, matpes_potential: MACEPotential, cluster: TensorDict + ) -> None: + """One energy, two seams: the TensorDict path must add nothing.""" + with torch.no_grad(): + direct = matpes_potential.energy_core(*raw_tensors(cluster)) + out = matpes_potential(cluster) + assert torch.allclose(out["graphs", "energy"], direct, atol=ENERGY_PARITY_ATOL, rtol=0.0) + + def test_matpes_energy_core_rejects_a_total_charge( + self, matpes_potential: MACEPotential, cluster: TensorDict + ) -> None: + """Silently ignoring the conditioning would return a wrong energy.""" + with pytest.raises(ValueError): + matpes_potential.energy_core( + *raw_tensors(cluster), total_charge=torch.zeros(1, dtype=torch.long) + ) + + def test_matpes_energy_core_rejects_a_total_spin( + self, matpes_potential: MACEPotential, cluster: TensorDict + ) -> None: + """Same for the spin channel — this variant has no joint embedding.""" + with pytest.raises(ValueError): + matpes_potential.energy_core( + *raw_tensors(cluster), total_spin=torch.ones(1, dtype=torch.long) + ) + + def test_omol_energy_core_consumes_the_total_charge( + self, omol_parity: tuple[MACEPotential, MACEOMol], omol_cluster: TensorDict + ) -> None: + """A different charge must reach the joint embedding and move the energy.""" + potential, _ = omol_parity + pos, Z, edge_index, batch, num_graphs = raw_tensors(omol_cluster) + spin = torch.ones(1, dtype=torch.long) + with torch.no_grad(): + neutral = potential.energy_core( + pos, + Z, + edge_index, + batch, + num_graphs, + total_charge=torch.zeros(1, dtype=torch.long), + total_spin=spin, + ) + cation = potential.energy_core( + pos, + Z, + edge_index, + batch, + num_graphs, + total_charge=torch.ones(1, dtype=torch.long), + total_spin=spin, + ) + assert float((cation - neutral).abs().max()) > 1e-6 + + # -- _write_energy: the protocol hook ------------------------------------- + + def test_call_energy_dispatches_to_the_energy_core( + self, matpes_potential: MACEPotential, cluster: TensorDict + ) -> None: + """``EnergyReadout`` must reach the core, not the force pipeline.""" + out = call_energy(matpes_potential, cluster) + assert ("graphs", "energy") in out.keys(include_nested=True) + assert ("atoms", "forces") not in out.keys(include_nested=True) + + def test_call_energy_keeps_the_energy_on_the_callers_leaf( + self, matpes_potential: MACEPotential, leaf_cluster: TensorDict + ) -> None: + """``_write_energy`` must not detach — ``ForceReadout`` differentiates it.""" + out = call_energy(matpes_potential, leaf_cluster) + assert out["graphs", "energy"].requires_grad is True + + # -- monomorphism --------------------------------------------------------- + + @pytest.mark.parametrize("method", ["forward", "energy_core", "_write_energy"]) + def test_no_compute_forces_on_the_public_surface(self, method: str) -> None: + """ "Energy only" is a construction, never a per-call flag.""" + parameters = inspect.signature(getattr(MACEPotential, method)).parameters + assert "compute_forces" not in parameters + + def test_forward_is_a_single_pipeline_dispatch(self) -> None: + """``return self._pipeline(batch)`` — every branch resolved in ``__init__``.""" + body = _function_body(MACEPotential.forward) + assert len(body) == 1 + statement = body[0] + assert isinstance(statement, ast.Return) + assert isinstance(statement.value, ast.Call) + assert _dotted_name(statement.value.func) == "self._pipeline" + + def test_no_hand_rolled_force_pass_in_this_module(self) -> None: + """Forces come from the shared kernel; a third force path is forbidden.""" + tree = _potential_module_tree() + called = _called_names(tree) + assert not called & FORBIDDEN_GRAD_CALLS + assert not {name for name in called if name.split(".")[-1] == "backward"} + + # -- numerical parity with the keyword variants (fp64, hard tolerance) ---- + + def test_matpes_energy_matches_the_variant_energy_forces( + self, matpes_parity: tuple[MACEPotential, MACEMatpes], cluster: TensorDict + ) -> None: + """Moving house must not move a bit (1e-12 eV on a −3.1 keV total).""" + potential, variant = matpes_parity + pos, Z, edge_index, batch, num_graphs = raw_tensors(cluster) + reference = variant.energy_forces(pos, Z, edge_index, batch, num_graphs=num_graphs) + energy = potential(cluster)["graphs", "energy"] + assert torch.allclose(energy, reference["energy"], atol=ENERGY_PARITY_ATOL, rtol=0.0), ( + f"bitwise equal: {torch.equal(energy, reference['energy'])}" + ) + + def test_matpes_forces_match_the_variant_energy_forces( + self, matpes_parity: tuple[MACEPotential, MACEMatpes], cluster: TensorDict + ) -> None: + """Same for ``F = -dE/dx`` (1e-12 eV/Å).""" + potential, variant = matpes_parity + pos, Z, edge_index, batch, num_graphs = raw_tensors(cluster) + reference = variant.energy_forces(pos, Z, edge_index, batch, num_graphs=num_graphs) + forces = potential(cluster)["atoms", "forces"] + assert torch.allclose(forces, reference["forces"], atol=FORCE_PARITY_ATOL, rtol=0.0), ( + f"bitwise equal: {torch.equal(forces, reference['forces'])}" + ) + + def test_matpes_periodic_energy_matches_the_variant_energy_forces( + self, matpes_parity: tuple[MACEPotential, MACEMatpes], periodic_cluster: TensorDict + ) -> None: + """The ``edges.shifts`` path is the MD path — it gets its own parity check.""" + potential, variant = matpes_parity + pos, Z, edge_index, batch, num_graphs = raw_tensors(periodic_cluster) + shifts = periodic_cluster["edges", "shifts"] + reference = variant.energy_forces( + pos, Z, edge_index, batch, num_graphs=num_graphs, shifts=shifts + ) + energy = potential(periodic_cluster)["graphs", "energy"] + assert torch.allclose(energy, reference["energy"], atol=ENERGY_PARITY_ATOL, rtol=0.0), ( + f"bitwise equal: {torch.equal(energy, reference['energy'])}" + ) + + def test_matpes_periodic_forces_match_the_variant_energy_forces( + self, matpes_parity: tuple[MACEPotential, MACEMatpes], periodic_cluster: TensorDict + ) -> None: + """``S_ij`` is constant w.r.t. ``pos``, so the forces stay exact.""" + potential, variant = matpes_parity + pos, Z, edge_index, batch, num_graphs = raw_tensors(periodic_cluster) + shifts = periodic_cluster["edges", "shifts"] + reference = variant.energy_forces( + pos, Z, edge_index, batch, num_graphs=num_graphs, shifts=shifts + ) + forces = potential(periodic_cluster)["atoms", "forces"] + assert torch.allclose(forces, reference["forces"], atol=FORCE_PARITY_ATOL, rtol=0.0), ( + f"bitwise equal: {torch.equal(forces, reference['forces'])}" + ) + + def test_omol_energy_matches_the_variant_energy_forces( + self, omol_parity: tuple[MACEPotential, MACEOMol], omol_cluster: TensorDict + ) -> None: + """The conditioned arm: charge 0 / spin 1, the flat model's own defaults.""" + potential, variant = omol_parity + pos, Z, edge_index, batch, _ = raw_tensors(omol_cluster) + reference = variant.energy_forces( + pos, + Z, + edge_index, + batch, + omol_cluster["graphs", "total_charge"], + omol_cluster["graphs", "total_spin"], + ) + energy = potential(omol_cluster)["graphs", "energy"] + assert torch.allclose(energy, reference["energy"], atol=ENERGY_PARITY_ATOL, rtol=0.0), ( + f"bitwise equal: {torch.equal(energy, reference['energy'])}" + ) + + def test_omol_forces_match_the_variant_energy_forces( + self, omol_parity: tuple[MACEPotential, MACEOMol], omol_cluster: TensorDict + ) -> None: + """Non-vacuous by construction — see :func:`.conftest.wake_zero_init_readout`.""" + potential, variant = omol_parity + pos, Z, edge_index, batch, _ = raw_tensors(omol_cluster) + reference = variant.energy_forces( + pos, + Z, + edge_index, + batch, + omol_cluster["graphs", "total_charge"], + omol_cluster["graphs", "total_spin"], + ) + largest = float(reference["forces"].detach().abs().max()) + assert largest > 0.0, "vacuous oracle: every reference force is exactly zero" + forces = potential(omol_cluster)["atoms", "forces"] + assert torch.allclose(forces, reference["forces"], atol=FORCE_PARITY_ATOL, rtol=0.0), ( + f"bitwise equal: {torch.equal(forces, reference['forces'])}" + ) + + def test_omol_defaults_to_a_neutral_closed_shell_singlet( + self, omol_parity: tuple[MACEPotential, MACEOMol], omol_cluster: TensorDict + ) -> None: + """Unconditioned batch → charge 0 / spin 1, as the flat model defaults. + + Spin 0 would index an untrained embedding row and quietly return + garbage, so the default must stay ``1`` (``mace_omol.py:330-335``). + """ + potential, _ = omol_parity + # Build the unconditioned twin first: ``forward`` swaps ``atoms.pos`` + # for its own leaf, and reading it back afterwards would compare a + # different tensor object. + bare = make_graph_batch( + pos=omol_cluster["atoms", "pos"].clone(), + Z=omol_cluster["atoms", "Z"], + edge_index=omol_cluster["edges", "edge_index"], + batch=omol_cluster["atoms", "batch"], + ) + conditioned = potential(omol_cluster)["graphs", "energy"].clone() + assert torch.allclose( + potential(bare)["graphs", "energy"], conditioned, atol=ENERGY_PARITY_ATOL, rtol=0.0 + ) + + # -- physics -------------------------------------------------------------- + + def test_net_force_vanishes(self, matpes_potential: MACEPotential, cluster: TensorDict) -> None: + """Newton's third law: the energy depends on ``pos`` only through ``r_ij``.""" + forces = matpes_potential(cluster)["atoms", "forces"] + assert float(forces.detach().sum(0).abs().max()) <= NET_FORCE_ATOL + + def test_energy_is_translation_invariant( + self, matpes_potential: MACEPotential, cluster: TensorDict + ) -> None: + """A rigid 1 Å shift cannot change a physical energy.""" + moved = translate_graph(cluster, torch.tensor(TRANSLATION, dtype=torch.float64)) + before = matpes_potential(cluster)["graphs", "energy"].clone() + after = matpes_potential(moved)["graphs", "energy"] + assert float((after - before).abs().max()) <= ENERGY_PARITY_ATOL + + def test_forces_match_central_differences( + self, matpes_potential: MACEPotential, cluster: TensorDict + ) -> None: + """``F = -dE/dx`` against a numerical derivative of ``energy_core``.""" + forces = matpes_potential(cluster)["atoms", "forces"] + pos, Z, edge_index, batch, num_graphs = raw_tensors(cluster) + pos = pos.detach() + for atom, axis in FD_PROBES: + shifted = pos.clone() + shifted[atom, axis] += FD_STEP + with torch.no_grad(): + plus = float( + matpes_potential.energy_core(shifted, Z, edge_index, batch, num_graphs).sum() + ) + shifted[atom, axis] -= 2 * FD_STEP + with torch.no_grad(): + minus = float( + matpes_potential.energy_core(shifted, Z, edge_index, batch, num_graphs).sum() + ) + numerical = -(plus - minus) / (2 * FD_STEP) + assert float(forces[atom, axis]) == pytest.approx(numerical, abs=FD_ATOL) + + # -- element table gate --------------------------------------------------- + + def test_rejects_atomic_numbers_outside_the_table( + self, matpes_potential: MACEPotential + ) -> None: + """``searchsorted`` would snap Fe onto a neighbour and be quietly wrong.""" + off_table = make_graph_batch( + pos=torch.tensor(CLUSTER_POS, dtype=torch.float64), + Z=torch.tensor(OFF_TABLE_Z, dtype=torch.long), + edge_index=full_edge_index([0] * len(OFF_TABLE_Z)), + batch=torch.zeros(len(OFF_TABLE_Z), dtype=torch.long), + ) + with pytest.raises(ValueError, match="outside this model"): + matpes_potential(off_table) + + def test_the_element_gate_runs_once_per_instance( + self, + matpes_potential: MACEPotential, + cluster: TensorDict, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + """Re-validating per call costs a host sync and a dynamo graph break. + + ``Z`` is constant over a trajectory and a wrongly wired model/dataset + pair fails on the very first batch, so the gate is spent after one + forward. Callers binding a *new* system to an existing instance call + ``validate_elements`` themselves (``mace_matpes.py:206-222`` semantics, + preserved deliberately). + """ + checks: list[torch.Tensor] = [] + validate = matpes_potential.validate_elements + + def counting(Z: torch.Tensor) -> None: + checks.append(Z) + validate(Z) + + monkeypatch.setattr(matpes_potential, "validate_elements", counting) + matpes_potential(cluster) + matpes_potential(cluster) + assert len(checks) == 1 + + # -- compile smoke -------------------------------------------------------- + + @pytest.mark.slow + def test_energy_core_compiles_to_one_graph( + self, matpes_potential: MACEPotential, cluster: TensorDict + ) -> None: + """The flat seam exists to be compiled: 1 dynamo graph, 0 breaks.""" + torch._dynamo.reset() + explanation = torch._dynamo.explain(matpes_potential.energy_core)(*raw_tensors(cluster)) + assert explanation.graph_count == 1 + assert explanation.graph_break_count == 0 diff --git a/tests/test_molzoo/test_mace/test_research.py b/tests/test_molzoo/test_mace/test_research.py new file mode 100644 index 0000000..7cbb8a7 --- /dev/null +++ b/tests/test_molzoo/test_mace/test_research.py @@ -0,0 +1,192 @@ +"""Tests for the research MACE encoder (``molzoo.mace.research``). + +Migration note (mace-subpackage-restructure-02-core): the deleted +``tests/test_molzoo/test_mace.py`` contained **no** research-encoder cases — +it only exercised ``EmbeddingBlock`` / ``InteractionBlock`` / ``ProductHead``, +which 01 promoted to ``molrep`` and whose cases now live in the +``tests/test_molrep`` mirrors. The ``molzoo.mace`` re-export surface that file +imported is already guarded by +``tests/test_molrep/test_reexport_compat.py::TestMolzooMaceShim``. + +What is left un-tested is the hot-path fix this step owes the research encoder: +``forward`` reads ``self.config.num_interactions`` on every call +(``src/molzoo/mace.py:254``), i.e. a pydantic attribute lookup inside the loop +head — the same dynamo graph-break hazard ``mace_matpes.py:118`` already +avoids with a plain ``self.num_interactions``. Those two cases are below. + +Migration note (mace-subpackage-restructure-06-wire): the one case of the +deleted ``tests/test_molzoo/test_mace_encoder.py`` — the encoder-output +contract — lands here and not in ``test_encoder.py``. That file mirrors +``molzoo/mace/encoder.py`` (:class:`~molzoo.mace.encoder.MACEEncoder`, the +foundation backbone); the symbol the case pins is +:class:`molzoo.mace.research.MACE`, whose mirror is this module, and whose +``TestMACE`` class already exists here. Its ``torch.randn`` geometry and +``torch.randint`` species were replaced by the package's fixed cluster, so the +case no longer depends on the ambient RNG. + +The import stays ``from molzoo.mace import MACE``, the stable package surface. + +fp64 support is pinned by :meth:`TestMACE.test_builds_and_runs_at_fp64`. +Until that case is green the research encoder **cannot be built at fp64**: +under ``config.set_precision("fp64")`` it comes out mixed-precision (17 fp64 / +14 fp32 parameters) and the first forward dies with ``mat1 and mat2 must have +the same dtype``. Two ``molrep`` layer families ignore the ``config["ftype"]`` +singleton that ``molrep.embedding.mlp`` (``:57``) honours: + +* ``molrep.embedding.node`` — ``nn.Embedding`` / ``nn.Linear`` / ``cuet.Linear`` + built without ``dtype=`` (``src/molrep/embedding/node.py:143,147,156,264, + 269-276``), which is where the two ``embedding.node_embedding.*`` fp32 + parameters come from; +* ``molrep.interaction.radial.RadialWeightMLP`` — ``nn.Linear`` without + ``dtype=`` (``src/molrep/interaction/radial.py:79,82``), which is where the + twelve ``interactions.*.radial_mlp.*`` fp32 parameters come from. + +The remaining fp32-only cases below (:func:`fp32`) are **not** a weakening — +they are the coverage the deleted flat file had, kept green across the fix so +the default precision cannot regress while fp64 is being enabled. Once fp64 is +green, whoever lands the fix should re-read this note: the site lists above are +the only part of it that goes stale. +""" + +from __future__ import annotations + +import ast +import inspect +from pathlib import Path + +import pytest +import torch +from tensordict import TensorDict + +from molix import config +from molrep.embedding.node import DiscreteEmbeddingSpec +from molzoo.mace import MACE +from tests.conftest import make_graph_batch +from tests.test_molzoo.test_mace.conftest import CLUSTER_POS, SINGLE_GRAPH, full_edge_index + +NUM_ELEMENTS = 5 +NUM_FEATURES = 8 +NUM_INTERACTIONS = 2 + +#: Species labels of :data:`~tests.test_molzoo.test_mace.conftest.CLUSTER_POS` +#: for the research encoder. Unlike the foundation variants it has no z-table: +#: ``Z`` is a plain 0-based class index into a ``num_classes`` embedding, so the +#: labels must stay below :data:`NUM_ELEMENTS`. +CLUSTER_SPECIES = [0, 1, 1, 2, 1] + + +def _config_reads(cls: type, method: str) -> list[str]: + """Attribute names read off ``self.config`` inside ``cls.method``.""" + source = Path(inspect.getsourcefile(cls) or "").read_text(encoding="utf-8") + class_def = next( + node + for node in ast.walk(ast.parse(source)) + if isinstance(node, ast.ClassDef) and node.name == cls.__name__ + ) + func = next( + node for node in class_def.body if isinstance(node, ast.FunctionDef) and node.name == method + ) + return [ + node.attr + for node in ast.walk(func) + if isinstance(node, ast.Attribute) + and isinstance(node.value, ast.Attribute) + and node.value.attr == "config" + and isinstance(node.value.value, ast.Name) + and node.value.value.id == "self" + ] + + +@pytest.fixture +def fp32(fp64) -> None: + """Override the package's autouse fp64 for the research encoder. + + Depends on ``fp64`` so it is torn down first and hands the precision back + in the state that fixture expects. See the module docstring for why the + research encoder cannot be built at fp64. + """ + config.set_precision("fp32") + yield + config.set_precision("fp64") + + +def tiny_mace() -> MACE: + """Build a tiny research MACE at the ambient precision. + + Two layers, eight scalar channels. Called by the fp32 :func:`encoder` + fixture and by the fp64 case, so both exercise the same configuration and + only the ambient ``config["ftype"]`` differs. + """ + return MACE( + node_attr_specs=[ + DiscreteEmbeddingSpec(input_key="Z", num_classes=NUM_ELEMENTS, emb_dim=NUM_FEATURES) + ], + num_elements=NUM_ELEMENTS, + num_features=NUM_FEATURES, + r_max=5.0, + num_bessel=4, + l_max=1, + num_interactions=NUM_INTERACTIONS, + correlation=2, + ) + + +def tiny_cluster() -> TensorDict: + """The package's five-atom cluster at the ambient precision.""" + return make_graph_batch( + pos=torch.tensor(CLUSTER_POS, dtype=config.ftype), + Z=torch.tensor(CLUSTER_SPECIES, dtype=torch.long), + edge_index=full_edge_index(SINGLE_GRAPH), + batch=torch.zeros(len(CLUSTER_SPECIES), dtype=torch.long), + ) + + +@pytest.fixture +def encoder(fp32) -> MACE: + """A tiny research MACE — two layers, eight scalar channels, fp32.""" + return tiny_mace() + + +@pytest.fixture +def species_cluster(fp32) -> TensorDict: + """The package's five-atom cluster relabelled with research species indices.""" + return tiny_cluster() + + +class TestMACE: + """Test the research MACE encoder.""" + + def test_num_interactions_is_a_plain_attribute(self, encoder): + """The layer count is frozen into a plain ``int`` at construction.""" + assert encoder.num_interactions == NUM_INTERACTIONS + assert isinstance(encoder.num_interactions, int) + + def test_forward_does_not_read_the_pydantic_config(self): + """``forward`` must not touch ``self.config`` — it is a graph break.""" + assert _config_reads(MACE, "forward") == [] + + # -- migrated from tests/test_molzoo/test_mace_encoder.py::TestMACE by + # mace-subpackage-restructure-06-wire (assertion verbatim) ------------ + + def test_forward_writes_per_layer_node_features(self, encoder, species_cluster): + """The encoder contract: ``atoms.node_features`` is ``(N, layers, features)``.""" + node_features = encoder(species_cluster)["atoms", "node_features"] + assert isinstance(node_features, torch.Tensor) + assert node_features.shape == (len(CLUSTER_SPECIES), NUM_INTERACTIONS, NUM_FEATURES) + + def test_builds_and_runs_at_fp64(self): + """The encoder is buildable and runnable at the package's ambient fp64. + + Deliberately does **not** request the :func:`fp32` override: it runs + under the package-wide autouse ``fp64`` fixture, which is the precision + the foundation-weight path uses. Every parameter must come out fp64 — + a single fp32 layer both loses the precision the caller asked for and + breaks the forward on a dtype-mismatched matmul. + """ + encoder = tiny_mace() + + assert {p.dtype for p in encoder.parameters()} == {torch.float64} + + node_features = encoder(tiny_cluster())["atoms", "node_features"] + assert node_features.dtype == torch.float64 + assert node_features.shape == (len(CLUSTER_SPECIES), NUM_INTERACTIONS, NUM_FEATURES) diff --git a/tests/test_molzoo/test_mace/test_spec.py b/tests/test_molzoo/test_mace/test_spec.py new file mode 100644 index 0000000..82b1c6d --- /dev/null +++ b/tests/test_molzoo/test_mace/test_spec.py @@ -0,0 +1,273 @@ +"""Tests for molzoo.mace.spec — the torch-free MACE configuration family. + +``spec.py`` is the one module in the MACE sub-package that must stay importable +without paying for the cuEquivariance stack, so half of this file is about what +it may *not* do: no heavy imports, no tensors in ``model_dump()``, and no +transitive load of ``molzoo.mace.encoder`` when the shim hands out a spec class. + +The expected defaults are hard-coded from the flat constructors they replace +(``src/molzoo/mace_matpes.py`` ``MACEMatpes.__init__`` and +``src/molzoo/mace_omol.py`` ``MACEOMol.__init__``) — that parity is the whole +point of the field table. +""" + +from __future__ import annotations + +import ast +import inspect +import json +import os +import subprocess +import sys +from pathlib import Path +from typing import Any + +import pytest +from pydantic import ValidationError + +from molzoo.mace.spec import MACEMatpesSpec, MACEOMolSpec, MACESpec +from tests.test_molzoo.test_mace.conftest import ATOMIC_ENERGIES, ATOMIC_NUMBERS + +#: The five variant switches are required on the base and defaulted on the +#: subclasses, so every direct ``MACESpec`` construction has to spell them out. +BASE_SWITCHES: dict[str, str] = { + "interaction": "density", + "readout": "per_layer", + "distance_transform": "none", + "pair_repulsion": "none", + "conditioning": "none", +} + +#: ``MACEMatpes.__init__`` defaults, read from src/molzoo/mace_matpes.py:89-107. +#: ``radial_mlp`` is ``None`` there and materialises as ``[64, 64, 64]`` at +#: mace_matpes.py:112; the spec carries the materialised list. +MATPES_DEFAULTS: dict[str, Any] = { + "r_max": 6.0, + "num_bessel": 10, + "num_polynomial_cutoff": 5, + "l_max": 3, + "num_features": 128, + "max_hidden_l": 1, + "num_interactions": 2, + "correlation": 3, + "mlp_dim": 16, + "radial_mlp": [64, 64, 64], + "scale": 1.0, + "shift": 0.0, + "use_fallback": False, + "interaction": "density", + "readout": "per_layer", + "distance_transform": "agnesi", + "pair_repulsion": "zbl", + "conditioning": "none", +} + +#: ``MACEOMol.__init__`` defaults, read from src/molzoo/mace_omol.py:65-85. +#: ``use_fallback`` is not a constructor argument there — the flat model +#: hard-codes the fused cuEq path (mace_omol.py:172, 182), i.e. ``False``. +OMOL_DEFAULTS: dict[str, Any] = { + "r_max": 6.0, + "num_bessel": 8, + "num_polynomial_cutoff": 5, + "l_max": 3, + "num_features": 1024, + "num_interactions": 3, + "correlation": 2, + "mlp_dim": 16, + "edge_channels": 128, + "charge_classes": 201, + "charge_offset": 100, + "spin_classes": 101, + "spin_offset": 0, + "scale": 1.0, + "shift": 0.0, + "use_fallback": False, + "interaction": "residual", + "readout": "final", + "distance_transform": "none", + "pair_repulsion": "none", + "conditioning": "charge_spin", +} + +#: Import roots that would defeat the point of a torch-free config module. +FORBIDDEN_IMPORT_ROOTS = frozenset( + { + "torch", + "cuequivariance", + "cuequivariance_torch", + "tensordict", + "molrep", + "molix", + "molpot", + "molzoo", + } +) + +#: Probe run in a fresh interpreter: reaching a spec class through the package +#: shim must not drag in the encoder (and with it the whole cuEq stack). +_LAZY_PROBE = ( + "import sys\n" + "import molzoo.mace\n" + "assert molzoo.mace.MACEMatpesSpec is not None\n" + "print('molzoo.mace.encoder' in sys.modules)\n" +) + + +def _base_spec(**overrides: Any) -> MACESpec: + """A valid base ``MACESpec`` with the shared table, plus ``overrides``.""" + return MACESpec( + atomic_numbers=list(ATOMIC_NUMBERS), + atomic_energies=list(ATOMIC_ENERGIES), + **{**BASE_SWITCHES, **overrides}, + ) + + +def _imported_roots(module: Any) -> set[str]: + """Top-level package names imported by ``module``'s source, via ``ast``.""" + source = Path(inspect.getsourcefile(module) or "").read_text(encoding="utf-8") + roots: set[str] = set() + for node in ast.walk(ast.parse(source)): + if isinstance(node, ast.Import): + roots.update(alias.name.split(".")[0] for alias in node.names) + elif isinstance(node, ast.ImportFrom): + if node.level: # relative import — resolves inside molzoo + roots.add("molzoo") + elif node.module: + roots.add(node.module.split(".")[0]) + return roots + + +class TestMACESpec: + """Test the shared MACE configuration base class.""" + + def test_variant_switches_are_required(self): + """The five variant switches have no base default — subclasses set them.""" + with pytest.raises(ValidationError): + MACESpec( + atomic_numbers=list(ATOMIC_NUMBERS), + atomic_energies=list(ATOMIC_ENERGIES), + ) + + def test_atomic_energies_length_must_match_the_table(self): + """One ``E0`` per element — a short list would silently mis-index.""" + with pytest.raises(ValueError, match="atomic_energies"): + MACESpec( + atomic_numbers=[1, 6, 8], + atomic_energies=[-13.6, -1029.0], + **BASE_SWITCHES, + ) + + def test_non_ascending_atomic_numbers_are_rejected(self): + """``torch.searchsorted`` needs an ordered table; unordered = wrong energy.""" + with pytest.raises(ValueError, match="atomic_numbers"): + MACESpec( + atomic_numbers=[8, 1, 6], + atomic_energies=list(ATOMIC_ENERGIES), + **BASE_SWITCHES, + ) + + def test_duplicated_atomic_numbers_are_rejected(self): + """A duplicate row makes the one-hot ambiguous — strictly ascending only.""" + with pytest.raises(ValueError, match="atomic_numbers"): + MACESpec( + atomic_numbers=[1, 6, 6], + atomic_energies=list(ATOMIC_ENERGIES), + **BASE_SWITCHES, + ) + + def test_zero_num_bessel_is_rejected(self): + """``num_bessel`` is a positive count (``Field(gt=0)``).""" + with pytest.raises(ValidationError): + _base_spec(num_bessel=0) + + def test_zero_r_max_is_rejected(self): + """A zero cutoff has no neighbours (``Field(gt=0.0)``); units are Å.""" + with pytest.raises(ValidationError): + _base_spec(r_max=0.0) + + def test_atomic_energies_stay_a_plain_list(self): + """Tensor conversion belongs to ``MACEEncoder.__init__``, not the spec.""" + assert _base_spec().atomic_energies == list(ATOMIC_ENERGIES) + + def test_model_dump_is_json_native(self): + """``model_dump()`` must round-trip through ``json`` — no tensors.""" + json.dumps(_base_spec().model_dump()) + + def test_spec_module_imports_no_heavy_dependencies(self): + """spec.py stays torch-free: only ``typing`` / ``pydantic`` and friends.""" + import molzoo.mace.spec as spec_module + + assert not (_imported_roots(spec_module) & FORBIDDEN_IMPORT_ROOTS) + + def test_reaching_a_spec_class_does_not_import_the_encoder(self): + """``molzoo.mace.MACEMatpesSpec`` must not trigger the encoder module.""" + molzoo_root = Path(inspect.getsourcefile(sys.modules["molzoo"]) or "").parents[1] + env = dict(os.environ) + env["PYTHONPATH"] = os.pathsep.join([str(molzoo_root), env.get("PYTHONPATH", "")]).rstrip( + os.pathsep + ) + + result = subprocess.run( + [sys.executable, "-c", _LAZY_PROBE], + capture_output=True, + text=True, + env=env, + check=True, + ) + assert result.stdout.strip() == "False" + + +class TestMACEMatpesSpec: + """Test the MACE-MatPES configuration.""" + + @pytest.mark.parametrize(("field", "expected"), sorted(MATPES_DEFAULTS.items())) + def test_field_default_matches_the_flat_constructor(self, field: str, expected: Any): + """Every default equals the ``MACEMatpes.__init__`` default it replaces.""" + spec = MACEMatpesSpec( + atomic_numbers=list(ATOMIC_NUMBERS), + atomic_energies=list(ATOMIC_ENERGIES), + ) + assert getattr(spec, field) == expected + + def test_single_interaction_is_rejected(self): + """MatPES needs a residual second layer (mace_matpes.py:109-110).""" + with pytest.raises(ValueError, match="num_interactions"): + MACEMatpesSpec( + atomic_numbers=list(ATOMIC_NUMBERS), + atomic_energies=list(ATOMIC_ENERGIES), + num_interactions=1, + ) + + def test_tiny_spec_keeps_the_overridden_values(self, tiny_matpes_spec): + """The shared tiny fixture round-trips its overrides (parity baseline).""" + assert tiny_matpes_spec.num_features == 16 + assert tiny_matpes_spec.radial_mlp == [8] + assert tiny_matpes_spec.use_fallback is True + + +class TestMACEOMolSpec: + """Test the MACE-OMOL configuration.""" + + @pytest.mark.parametrize(("field", "expected"), sorted(OMOL_DEFAULTS.items())) + def test_field_default_matches_the_flat_constructor(self, field: str, expected: Any): + """Every default equals the ``MACEOMol.__init__`` default it replaces.""" + spec = MACEOMolSpec( + atomic_numbers=list(ATOMIC_NUMBERS), + atomic_energies=list(ATOMIC_ENERGIES), + ) + assert getattr(spec, field) == expected + + def test_scalar_only_angular_order_is_rejected(self): + """OMOL's mid-layer edge irreps use ``range(l_max)`` (mace_omol.py:151-153).""" + with pytest.raises(ValueError, match="l_max"): + MACEOMolSpec( + atomic_numbers=list(ATOMIC_NUMBERS), + atomic_energies=list(ATOMIC_ENERGIES), + l_max=0, + ) + + def test_tiny_spec_keeps_the_overridden_values(self, tiny_omol_spec): + """The shared tiny fixture round-trips its overrides (parity baseline).""" + assert tiny_omol_spec.num_features == 16 + assert tiny_omol_spec.edge_channels == 8 + assert tiny_omol_spec.conditioning == "charge_spin" diff --git a/tests/test_molzoo/test_mace/test_variants.py b/tests/test_molzoo/test_mace/test_variants.py new file mode 100644 index 0000000..8b69130 --- /dev/null +++ b/tests/test_molzoo/test_mace/test_variants.py @@ -0,0 +1,727 @@ +"""Tests for molzoo.mace.variants — the two named foundation models. + +:class:`~molzoo.mace.variants.MACEMatpes` and +:class:`~molzoo.mace.variants.MACEOMol` are **existing public API**: they are +constructed by keyword in ``scripts/matpes_port/run_nve.py:148-164`` and +``benchmarks/bench_mace_matpes.py:41-54``, which the +``mace-subpackage-restructure-06-wire`` cutover must not touch. After the +cutover they are thin adapters — old keyword signature in, 04/05's spec presets +and :class:`~molzoo.mace.potential.MACEPotential` out — so what this file pins +is the surface those callers bind, plus everything that is genuinely the +variants' own: :meth:`MACEMatpes.energy_forces` / +:meth:`MACEOMol.energy_forces` exist **only** here (``MACEPotential`` has +``forward`` and ``energy_core``), so the raw-tensor physics belongs in this +module and not in ``test_potential.py``. + +The load-bearing oddity is ``_compute_energy``: a **private** method consumed +across packages (``run_nve.py:118`` binds it as the compiled energy core, +``bench_mace_matpes.py``'s compiled arm compiles it). Re-pointing that consumer at the +public :meth:`~molzoo.mace.potential.MACEPotential.energy_core` belongs to +``mace-subpackage-restructure-07-cleanup``; until then the positional signature +``(pos, Z, edge_index, batch, num_graphs, shifts)`` is a contract, and the case +below is what keeps the NVE script alive on cutover day. + +Migration note (``mace-subpackage-restructure-06-wire``): the cases below the +``-- migrated from tests/test_molzoo/test_mace_{matpes,omol}.py --`` markers +come from the two pre-cutover flat test files, deleted by that step. Their +assertion criteria and tolerances are preserved verbatim; what changed is the +import (``molzoo.mace.variants`` instead of the flat modules), the geometry +(the package's fixed :data:`~tests.test_molzoo.test_mace.conftest.CLUSTER_POS` +instead of a seeded ``randn`` cloud) and, for OMOL, the precision — the flat +OMOL file built at fp32 and called ``.double()``, which leaves cuEquivariance +contracting in fp32; the package's autouse ``fp64`` fixture builds at fp64, so +the same tolerances are strictly harder to meet. + +Every model here is tiny (2 layers, 16 channels, ``l_max=1``), fp64 (autouse +``fp64`` fixture in ``conftest.py``), CPU and seeded. +""" + +from __future__ import annotations + +import math + +import pytest +import torch +from tensordict import TensorDict + +from molzoo.mace.potential import MACEPotential +from molzoo.mace.spec import MACEMatpesSpec, MACEOMolSpec +from molzoo.mace.variants import ( + MACEMatpes, + MACEOMol, + load_matpes_state_dict, + load_omol_state_dict, +) +from tests.conftest import make_graph_batch +from tests.test_molzoo.test_mace.conftest import ( + ATOMIC_ENERGIES, + ATOMIC_NUMBERS, + CLUSTER_POS, + CLUSTER_Z, + PAIR_Z, + SINGLE_GRAPH, + TINY_MATPES_KWARGS, + TINY_OMOL_KWARGS, + full_edge_index, + raw_tensors, +) + +#: Iron (26) is outside the ``[1, 6, 8]`` table — ``searchsorted`` would snap it +#: onto a neighbouring row and return a plausible, wrong energy. +OFF_TABLE_Z = [26] + +#: The same probe inside a whole five-atom system, for the ``forward`` gate. +OFF_TABLE_CLUSTER_Z = [1, 6, 8, 1, 26] + +#: ``scale`` / ``shift`` are the two ``run_nve.py`` kwargs the tiny preset omits; +#: pinned to non-default values so the construction case really passes all 15. +SCALE = 1.5 +SHIFT = -0.25 + +#: OMOL's charge/spin conditioning sizes, the four kwargs outside the tiny +#: preset. Small on purpose: the embedding tables are ``num_classes``-sized. +CHARGE_CLASSES = 7 +CHARGE_OFFSET = 3 +SPIN_CLASSES = 5 +SPIN_OFFSET = 0 + +#: A checkpoint key that exists in no MACE dialect — the loaders must refuse it +#: rather than quietly leaving the model at its initialisation. +ALIEN_CHECKPOINT: dict[str, torch.Tensor] = {"not.a.mace.key": torch.zeros(1)} + +#: Same class, same weights, same operator order: the kwargs form and the spec +#: form are one code path, so their energies must agree *bitwise*, not closely. +EXACT = 0.0 + +#: ``forward(batch)`` (batch schema, shared force kernel) against +#: ``energy_forces(...)`` (raw tensors, own position leaf) — two code paths on +#: one energy. Criteria carried over verbatim from the flat test files. +FORWARD_ENERGY_ATOL = 1e-9 # eV +FORWARD_FORCE_ATOL = 1e-8 # eV/Å + +#: Newton's third law on an isolated system, as the two flat files asserted it. +MATPES_NET_FORCE_ATOL = 1e-8 # eV/Å +OMOL_NET_FORCE_ATOL = 1e-7 # eV/Å + +#: Rotation angles (rad) of the two rigid test rotations: ``Rz`` for the energy +#: invariance, ``Rx`` for the force equivariance (flat file's own choices). +ROTATION_Z = 0.7 +ROTATION_X = -0.4 + +#: Central-difference step and tolerance of the finite-difference force check. +FD_STEP = 1e-6 # Å +FD_ATOL = 1e-5 # eV/Å + +#: Atom/axis pairs probed by the finite-difference case. +FD_PROBES = ((0, 0), (2, 1), (4, 2)) + +#: Cubic cell edge (Å) and neighbour cutoff (Å) of the dead-edge padding cases. +PBC_CELL = 7.0 +PBC_CUTOFF = 3.4 + +#: Capacity factors of the un-padded and heavily padded neighbour lists. +PBC_CAPACITY_FACTORS = (1.0, 3.0) + +#: A four-atom water-like cluster inside :data:`PBC_CELL`, one atom wrapped +#: across the boundary so the minimum-image path is exercised. +PBC_POS = [ + [0.5, 0.5, 0.5], + [1.6, 0.6, 0.4], + [0.4, 1.7, 0.6], + [6.8, 6.9, 0.2], +] +PBC_Z = [8, 1, 1, 1] + +#: The rebuild case drops the wrapped atom — three atoms, no periodic image. +REBUILD_CAPACITY_FACTOR = 2.0 + +#: Per-graph conditioning of the batched OMOL case: a neutral and an anionic +#: molecule with different spin rows. +PAIR_TOTAL_CHARGE = [0, -1] +PAIR_TOTAL_SPIN = [0, 1] + +#: Conditioning of the ``charged_omol_cluster``: deliberately **not** the +#: ``(0, 1)`` neutral singlet ``forward`` falls back to, so a ``forward`` that +#: ignored ``graphs`` could not pass the consistency cases. +CHARGED_TOTAL_CHARGE = 1 +CHARGED_TOTAL_SPIN = 0 + +#: Parameter perturbation of the force-loss cases: a fresh MACE readout is +#: zero-initialised, so an unperturbed model has identically zero forces and no +#: force loss can reach anything (a vacuous pass). +PERTURBATION = 0.05 +PERTURBATION_SEED = 7 + + +def _rotation_z(angle: float) -> torch.Tensor: + """Right-handed rotation about ``z`` by ``angle`` rad, fp64.""" + c, s = math.cos(angle), math.sin(angle) + return torch.tensor([[c, -s, 0.0], [s, c, 0.0], [0.0, 0.0, 1.0]], dtype=torch.float64) + + +def _rotation_x(angle: float) -> torch.Tensor: + """Right-handed rotation about ``x`` by ``angle`` rad, fp64.""" + c, s = math.cos(angle), math.sin(angle) + return torch.tensor([[1.0, 0.0, 0.0], [0.0, c, -s], [0.0, s, c]], dtype=torch.float64) + + +def _perturb(model: torch.nn.Module) -> None: + """Move every parameter off its initialisation, deterministically.""" + torch.manual_seed(PERTURBATION_SEED) + with torch.no_grad(): + for parameter in model.parameters(): + parameter.add_(torch.randn_like(parameter) * PERTURBATION) + + +def _fraction_with_gradient(model: torch.nn.Module) -> tuple[int, int]: + """``(parameters carrying a non-zero grad, total parameters)``.""" + parameters = list(model.parameters()) + with_grad = sum(int(p.grad is not None and float(p.grad.abs().sum()) > 0.0) for p in parameters) + return with_grad, len(parameters) + + +@pytest.fixture +def charged_omol_cluster() -> TensorDict: + """The cluster conditioned on a non-default charge/spin pair.""" + return make_graph_batch( + pos=torch.tensor(CLUSTER_POS, dtype=torch.float64), + Z=torch.tensor(CLUSTER_Z, dtype=torch.long), + edge_index=full_edge_index(SINGLE_GRAPH), + batch=torch.zeros(len(CLUSTER_Z), dtype=torch.long), + graphs={ + "total_charge": torch.tensor([CHARGED_TOTAL_CHARGE], dtype=torch.long), + "total_spin": torch.tensor([CHARGED_TOTAL_SPIN], dtype=torch.long), + }, + ) + + +@pytest.fixture +def omol_pair_batch(pair_batch: TensorDict) -> TensorDict: + """Two molecules in one batch, each with its own charge and spin.""" + pair_batch["graphs", "total_charge"] = torch.tensor(PAIR_TOTAL_CHARGE, dtype=torch.long) + pair_batch["graphs", "total_spin"] = torch.tensor(PAIR_TOTAL_SPIN, dtype=torch.long) + return pair_batch + + +class TestMACEMatpes: + """The keyword surface ``run_nve.py`` / ``bench_mace_matpes.py`` bind.""" + + def test_constructs_from_the_nve_script_keywords(self) -> None: + """All 15 ``run_nve.py:148-164`` keywords, ``scale`` / ``shift`` included.""" + torch.manual_seed(0) + model = MACEMatpes( + atomic_numbers=list(ATOMIC_NUMBERS), + atomic_energies=torch.tensor(ATOMIC_ENERGIES), + scale=SCALE, + shift=SHIFT, + **TINY_MATPES_KWARGS, + ) + assert isinstance(model, torch.nn.Module) + + def test_exposes_the_consumed_method_surface(self, matpes_variant: MACEMatpes) -> None: + """``forward`` / ``energy_forces`` / ``validate_elements`` all survive.""" + assert all( + callable(getattr(matpes_variant, name, None)) + for name in ("forward", "energy_forces", "validate_elements") + ) + + def test_registers_the_element_table_as_a_buffer(self, matpes_variant: MACEMatpes) -> None: + """``z_table`` is a persistent buffer, so it moves with ``.to(device)``.""" + buffers = dict(matpes_variant.named_buffers()) + assert torch.equal(buffers["z_table"], torch.tensor(ATOMIC_NUMBERS, dtype=torch.long)) + + def test_compute_energy_takes_six_positional_arguments( + self, matpes_variant: MACEMatpes, cluster: TensorDict + ) -> None: + """``run_nve.py:118`` binds this private core and calls it positionally.""" + edge_index = cluster["edges", "edge_index"] + shifts = torch.zeros(edge_index.shape[0], 3, dtype=torch.float64) + with torch.no_grad(): + energy = matpes_variant._compute_energy( + cluster["atoms", "pos"], + cluster["atoms", "Z"], + edge_index, + cluster["atoms", "batch"], + 1, + shifts, + ) + assert energy.shape == (1,) + + def test_compute_energy_agrees_with_the_forward_energy( + self, matpes_variant: MACEMatpes, cluster: TensorDict + ) -> None: + """The compiled MD core and the batch path are one energy, not two.""" + with torch.no_grad(): + direct = matpes_variant._compute_energy( + cluster["atoms", "pos"], + cluster["atoms", "Z"], + cluster["edges", "edge_index"], + cluster["atoms", "batch"], + 1, + None, + ) + energy = matpes_variant(cluster)["graphs", "energy"] + assert float((energy - direct).abs().max()) == EXACT + + def test_energy_forces_returns_energy_and_forces( + self, matpes_variant: MACEMatpes, cluster: TensorDict + ) -> None: + """``bench_mace_matpes.py``'s eager arm calls it with ``num_graphs=`` / ``shifts=``.""" + out = matpes_variant.energy_forces( + cluster["atoms", "pos"], + cluster["atoms", "Z"], + cluster["edges", "edge_index"], + cluster["atoms", "batch"], + num_graphs=1, + shifts=None, + ) + assert out["energy"].shape == (1,) and out["forces"].shape == (len(CLUSTER_Z), 3) + + def test_forward_writes_the_energy_and_forces_on_the_batch( + self, matpes_variant: MACEMatpes, cluster: TensorDict + ) -> None: + """The molnex pipeline entry point: ``graphs.energy`` + ``atoms.forces``.""" + out = matpes_variant(cluster) + assert out["graphs", "energy"].shape == (1,) + assert out["atoms", "forces"].shape == (len(CLUSTER_Z), 3) + + def test_kwargs_form_matches_the_spec_form_exactly( + self, matpes_variant: MACEMatpes, tiny_matpes_spec: MACEMatpesSpec, cluster: TensorDict + ) -> None: + """The alias must *be* the spec form, not merely approximate it.""" + spec_form = MACEPotential(tiny_matpes_spec).eval() + spec_form.load_state_dict(matpes_variant.state_dict(), strict=True) + alias_energy = matpes_variant(cluster.clone())["graphs", "energy"] + spec_energy = spec_form(cluster.clone())["graphs", "energy"] + assert float((alias_energy - spec_energy).abs().max()) == EXACT + + def test_rejects_atomic_numbers_outside_the_table(self, matpes_variant: MACEMatpes) -> None: + """An off-table element would be snapped onto a neighbouring row.""" + with pytest.raises(ValueError): + matpes_variant.validate_elements(torch.tensor(OFF_TABLE_Z, dtype=torch.long)) + + def test_load_matpes_state_dict_refuses_an_alien_checkpoint( + self, matpes_variant: MACEMatpes + ) -> None: + """The alias forwards to the strict ``CheckpointRemap``, not to ``load_state_dict``.""" + with pytest.raises(RuntimeError): + load_matpes_state_dict(matpes_variant, ALIEN_CHECKPOINT) + + # -- migrated from tests/test_molzoo/test_mace_matpes.py::TestMACEMatpes by + # mace-subpackage-restructure-06-wire (criteria verbatim) -------------- + + def test_forward_matches_energy_forces( + self, matpes_variant: MACEMatpes, cluster: TensorDict + ) -> None: + """forward(td) energy/forces == raw energy_forces() to machine precision.""" + pos, Z, edge_index, batch, _ = raw_tensors(cluster) + reference = matpes_variant.energy_forces(pos, Z, edge_index, batch, num_graphs=1) + out = matpes_variant.forward(cluster) + assert torch.allclose( + out["graphs", "energy"], reference["energy"], atol=FORWARD_ENERGY_ATOL, rtol=0 + ) + assert torch.allclose( + out["atoms", "forces"], reference["forces"], atol=FORWARD_FORCE_ATOL, rtol=0 + ) + + def test_forward_returns_the_same_batch_with_new_keys( + self, matpes_variant: MACEMatpes, cluster: TensorDict + ) -> None: + """forward mutates in place and returns the same TensorDict object.""" + out = matpes_variant.forward(cluster) + assert out is cluster + nested = out.keys(include_nested=True) + assert ("graphs", "energy") in nested + assert ("atoms", "forces") in nested + + def test_net_force_vanishes_on_an_isolated_cluster( + self, matpes_variant: MACEMatpes, cluster: TensorDict + ) -> None: + """Net force through the variant's own raw-tensor seam is ~zero. + + ``forward``-path net force is covered by + ``test_potential.py::TestMACEPotential::test_net_force_vanishes``; this + case drives ``energy_forces`` so the variant-only leaf path carries its + own Newton's-third-law lock. + """ + pos, Z, edge_index, batch, _ = raw_tensors(cluster) + result = matpes_variant.energy_forces(pos, Z, edge_index, batch, num_graphs=1) + assert float(result["forces"].detach().sum(0).abs().max()) < MATPES_NET_FORCE_ATOL + + def test_energy_is_rotation_invariant( + self, matpes_variant: MACEMatpes, cluster: TensorDict + ) -> None: + """A rigid rotation leaves the energy unchanged (O(3) invariance).""" + pos, Z, edge_index, batch, _ = raw_tensors(cluster) + rotation = _rotation_z(ROTATION_Z) + base = matpes_variant.energy_forces(pos, Z, edge_index, batch, num_graphs=1) + turned = matpes_variant.energy_forces(pos @ rotation.T, Z, edge_index, batch, num_graphs=1) + assert torch.allclose(base["energy"], turned["energy"], atol=FORWARD_ENERGY_ATOL, rtol=0) + + def test_forces_rotate_with_the_system( + self, matpes_variant: MACEMatpes, cluster: TensorDict + ) -> None: + """Forces are equivariant: F(Rx) = R F(x).""" + pos, Z, edge_index, batch, _ = raw_tensors(cluster) + rotation = _rotation_x(ROTATION_X) + base = matpes_variant.energy_forces(pos, Z, edge_index, batch, num_graphs=1)["forces"] + turned = matpes_variant.energy_forces(pos @ rotation.T, Z, edge_index, batch, num_graphs=1)[ + "forces" + ] + assert torch.allclose(turned, base @ rotation.T, atol=FORWARD_FORCE_ATOL, rtol=0) + + def test_forces_match_finite_differences( + self, matpes_variant: MACEMatpes, cluster: TensorDict + ) -> None: + """``F = -dE/dpos`` — the autograd force must match a numeric gradient.""" + pos, Z, edge_index, batch, _ = raw_tensors(cluster) + forces = matpes_variant.energy_forces(pos, Z, edge_index, batch, num_graphs=1)["forces"] + + for atom, axis in FD_PROBES: + shifted = pos.clone() + shifted[atom, axis] += FD_STEP + plus = float( + matpes_variant.energy_forces(shifted, Z, edge_index, batch, num_graphs=1)[ + "energy" + ].sum() + ) + shifted[atom, axis] -= 2 * FD_STEP + minus = float( + matpes_variant.energy_forces(shifted, Z, edge_index, batch, num_graphs=1)[ + "energy" + ].sum() + ) + assert float(forces[atom, axis]) == pytest.approx( + -(plus - minus) / (2 * FD_STEP), abs=FD_ATOL + ) + + def test_batched_graphs(self, matpes_variant: MACEMatpes, pair_batch: TensorDict) -> None: + """Two clusters in one batch: per-graph energies, correct atom routing.""" + pos, Z, edge_index, batch, num_graphs = raw_tensors(pair_batch) + reference = matpes_variant.energy_forces(pos, Z, edge_index, batch, num_graphs=num_graphs) + out = matpes_variant.forward(pair_batch) + assert out["graphs", "energy"].shape == (2,) + assert torch.allclose( + out["graphs", "energy"], reference["energy"], atol=FORWARD_ENERGY_ATOL, rtol=0 + ) + + def test_energy_is_extensive_over_separated_graphs( + self, matpes_variant: MACEMatpes, cluster: TensorDict + ) -> None: + """Two non-interacting copies carry twice one copy's energy.""" + pos, Z, edge_index, batch, _ = raw_tensors(cluster) + n_atoms = len(CLUSTER_Z) + one = matpes_variant.energy_forces(pos, Z, edge_index, batch, num_graphs=1)["energy"] + pair = matpes_variant.energy_forces( + torch.cat([pos, pos]), + torch.cat([Z, Z]), + torch.cat([edge_index, edge_index + n_atoms], dim=0), + torch.cat([batch, batch + 1]), + num_graphs=2, + )["energy"] + assert torch.allclose(pair, one.repeat(2), atol=FORWARD_ENERGY_ATOL, rtol=0) + + def test_forward_rejects_an_element_outside_the_table( + self, matpes_variant: MACEMatpes, cluster: TensorDict + ) -> None: + """An out-of-table Z would be snapped onto a neighbour and silently wrong.""" + cluster["atoms", "Z"] = torch.tensor(OFF_TABLE_CLUSTER_Z, dtype=torch.long) + with pytest.raises(ValueError, match="outside this model"): + matpes_variant.forward(cluster) + + def test_rejects_single_interaction(self) -> None: + """The readout schedule needs a first and a last layer. + + The spec owns the rule (``test_spec.py``); this pins that the keyword + constructor propagates it instead of silently building a one-layer model. + """ + with pytest.raises(ValueError, match="num_interactions"): + MACEMatpes( + atomic_numbers=list(ATOMIC_NUMBERS), + atomic_energies=torch.tensor(ATOMIC_ENERGIES), + num_interactions=1, + ) + + def test_first_product_has_no_skip_connection(self, matpes_variant: MACEMatpes) -> None: + """MACE only carries a residual from the second layer on.""" + assert matpes_variant.products[0].use_sc is False + assert matpes_variant.products[1].use_sc is True + + def test_last_layer_keeps_scalars_only(self, matpes_variant: MACEMatpes) -> None: + """The final hidden state is scalar; earlier layers keep l>0.""" + assert matpes_variant.readouts[0].__class__.__name__ == "LinearReadout" + assert matpes_variant.readouts[1].__class__.__name__ == "NonLinearReadout" + + # -- migrated from tests/test_molzoo/test_mace_matpes.py::TestDeadEdgePadding + # (same step; the subject is the variant, so the two cases fold into this + # class rather than keeping a second class for one production unit) ---- + # + # A padded (dead) edge must contribute exactly nothing. This is the + # load-bearing property of :class:`molix.md.NeighborList`: its + # fixed-capacity buffers pad with self-loops on atom 0 displaced beyond the + # cutoff, and the whole rebuild-under-CUDA-graphs design assumes such edges + # are invisible to the model. Random weights make the check stronger — the + # zeros must come from the cutoff structure, not from trained smallness. + + def test_energy_and_forces_ignore_dead_edges(self, matpes_variant: MACEMatpes) -> None: + """Same system, 0% vs heavy padding: identical energy and forces.""" + from molix.md import NeighborList + + cell = torch.eye(3, dtype=torch.float64) * PBC_CELL + pos = torch.tensor(PBC_POS, dtype=torch.float64) + Z = torch.tensor(PBC_Z, dtype=torch.long) + batch = torch.zeros(len(PBC_Z), dtype=torch.long) + + outs = [] + for factor in PBC_CAPACITY_FACTORS: + nl = NeighborList(cell=cell, cutoff=PBC_CUTOFF, positions=pos, capacity_factor=factor) + out = matpes_variant.energy_forces( + pos, Z, nl.edge_index, batch, num_graphs=1, shifts=nl.shifts + ) + outs.append((nl, out)) + + (nl_a, a), (nl_b, b) = outs + assert nl_b.capacity > nl_a.capacity # the padded arm really is padded + assert nl_a.num_edges == nl_b.num_edges + assert torch.equal(a["energy"], b["energy"]) + assert torch.equal(a["forces"], b["forces"]) + + def test_rebuild_preserves_energy_at_identical_positions( + self, matpes_variant: MACEMatpes + ) -> None: + """rebuild() at the same positions must not change the physics.""" + from molix.md import NeighborList + + cell = torch.eye(3, dtype=torch.float64) * PBC_CELL + pos = torch.tensor(PBC_POS[:3], dtype=torch.float64) + Z = torch.tensor(PBC_Z[:3], dtype=torch.long) + batch = torch.zeros(3, dtype=torch.long) + nl = NeighborList( + cell=cell, + cutoff=PBC_CUTOFF, + positions=pos, + capacity_factor=REBUILD_CAPACITY_FACTOR, + ) + before = matpes_variant.energy_forces( + pos, Z, nl.edge_index, batch, num_graphs=1, shifts=nl.shifts + )["energy"].clone() + nl.rebuild(pos) + after = matpes_variant.energy_forces( + pos, Z, nl.edge_index, batch, num_graphs=1, shifts=nl.shifts + )["energy"] + assert torch.equal(before, after) + + +class TestMACEOMol: + """The keyword surface of the OMOL foundation variant.""" + + def test_constructs_from_the_seventeen_keywords(self) -> None: + """The full consumed signature, charge/spin conditioning sizes included.""" + torch.manual_seed(0) + model = MACEOMol( + atomic_numbers=list(ATOMIC_NUMBERS), + atomic_energies=torch.tensor(ATOMIC_ENERGIES), + charge_classes=CHARGE_CLASSES, + charge_offset=CHARGE_OFFSET, + spin_classes=SPIN_CLASSES, + spin_offset=SPIN_OFFSET, + scale=SCALE, + shift=SHIFT, + **TINY_OMOL_KWARGS, + ) + assert isinstance(model, torch.nn.Module) + + def test_exposes_the_consumed_method_surface(self, omol_variant: MACEOMol) -> None: + """``forward`` / ``energy_forces`` / ``validate_elements`` all survive.""" + assert all( + callable(getattr(omol_variant, name, None)) + for name in ("forward", "energy_forces", "validate_elements") + ) + + def test_registers_the_element_table_as_a_buffer(self, omol_variant: MACEOMol) -> None: + """``z_table`` is a persistent buffer, so it moves with ``.to(device)``.""" + buffers = dict(omol_variant.named_buffers()) + assert torch.equal(buffers["z_table"], torch.tensor(ATOMIC_NUMBERS, dtype=torch.long)) + + def test_energy_forces_returns_energy_and_forces( + self, omol_variant: MACEOMol, omol_cluster: TensorDict + ) -> None: + """OMOL's own raw-tensor entry point takes the two conditioning tensors.""" + out = omol_variant.energy_forces( + omol_cluster["atoms", "pos"], + omol_cluster["atoms", "Z"], + omol_cluster["edges", "edge_index"], + omol_cluster["atoms", "batch"], + omol_cluster["graphs", "total_charge"], + omol_cluster["graphs", "total_spin"], + ) + assert out["energy"].shape == (1,) and out["forces"].shape == (len(CLUSTER_Z), 3) + + def test_forward_writes_the_energy_and_forces_on_the_batch( + self, omol_variant: MACEOMol, omol_cluster: TensorDict + ) -> None: + """The molnex pipeline entry point: ``graphs.energy`` + ``atoms.forces``.""" + out = omol_variant(omol_cluster) + assert out["graphs", "energy"].shape == (1,) + assert out["atoms", "forces"].shape == (len(CLUSTER_Z), 3) + + def test_kwargs_form_matches_the_spec_form_exactly( + self, omol_variant: MACEOMol, tiny_omol_spec: MACEOMolSpec, omol_cluster: TensorDict + ) -> None: + """Non-vacuous by construction — see ``conftest.wake_zero_init_readout``.""" + spec_form = MACEPotential(tiny_omol_spec).eval() + spec_form.load_state_dict(omol_variant.state_dict(), strict=True) + alias_energy = omol_variant(omol_cluster.clone())["graphs", "energy"] + spec_energy = spec_form(omol_cluster.clone())["graphs", "energy"] + assert float((alias_energy - spec_energy).abs().max()) == EXACT + + def test_rejects_atomic_numbers_outside_the_table(self, omol_variant: MACEOMol) -> None: + """The pre-cutover flat ``MACEOMol`` had no such gate; the shared encoder brings one.""" + with pytest.raises(ValueError): + omol_variant.validate_elements(torch.tensor(OFF_TABLE_Z, dtype=torch.long)) + + def test_load_omol_state_dict_refuses_an_alien_checkpoint(self, omol_variant: MACEOMol) -> None: + """OMOL returns unhoused keys, but still refuses to leave parameters unfilled.""" + with pytest.raises(RuntimeError): + load_omol_state_dict(omol_variant, ALIEN_CHECKPOINT) + + # -- migrated from tests/test_molzoo/test_mace_omol.py (module-level cases) + # by mace-subpackage-restructure-06-wire (criteria verbatim) ----------- + + def test_forward_matches_energy_forces( + self, omol_variant: MACEOMol, charged_omol_cluster: TensorDict + ) -> None: + """forward(td) energy/forces == raw energy_forces() to machine precision. + + Conditioned on a *non-default* charge/spin pair, so a ``forward`` that + ignored ``graphs`` and fell back to the neutral singlet would fail here. + """ + pos, Z, edge_index, batch, _ = raw_tensors(charged_omol_cluster) + total_charge = charged_omol_cluster["graphs", "total_charge"] + total_spin = charged_omol_cluster["graphs", "total_spin"] + reference = omol_variant.energy_forces(pos, Z, edge_index, batch, total_charge, total_spin) + out = omol_variant.forward(charged_omol_cluster) + assert torch.allclose( + out["graphs", "energy"], reference["energy"], atol=FORWARD_ENERGY_ATOL, rtol=0 + ) + assert torch.allclose( + out["atoms", "forces"], reference["forces"], atol=FORWARD_FORCE_ATOL, rtol=0 + ) + + def test_forward_returns_the_same_batch_with_new_keys( + self, omol_variant: MACEOMol, charged_omol_cluster: TensorDict + ) -> None: + """forward mutates in place and returns the same TensorDict object.""" + out = omol_variant.forward(charged_omol_cluster) + assert out is charged_omol_cluster + assert ("graphs", "energy") in out.keys(include_nested=True) + assert ("atoms", "forces") in out.keys(include_nested=True) + + def test_net_force_vanishes_on_an_isolated_molecule( + self, omol_variant: MACEOMol, charged_omol_cluster: TensorDict + ) -> None: + """Net force on an isolated molecule is ~zero (translation invariance).""" + out = omol_variant.forward(charged_omol_cluster) + assert float(out["atoms", "forces"].detach().sum(0).abs().max()) < OMOL_NET_FORCE_ATOL + + def test_missing_charge_spin_defaults_to_neutral( + self, omol_variant: MACEOMol, cluster: TensorDict + ) -> None: + """Absent graphs.total_charge/total_spin → neutral singlet, still runs.""" + pos, Z, edge_index, batch, _ = raw_tensors(cluster) + # forward defaults to the OMOL neutral closed-shell singlet: charge=0, + # spin=1 (spin index 1 is a trained row; spin 0 hits an untrained + # embedding). + charge = torch.zeros(1, dtype=torch.long) + spin = torch.ones(1, dtype=torch.long) + reference = omol_variant.energy_forces(pos, Z, edge_index, batch, charge, spin) + out = omol_variant.forward(cluster) + assert torch.allclose( + out["graphs", "energy"], reference["energy"], atol=FORWARD_ENERGY_ATOL, rtol=0 + ) + + def test_force_loss_reaches_parameters_in_eval_mode( + self, omol_variant: MACEOMol, charged_omol_cluster: TensorDict + ) -> None: + """Force-supervised training must backprop to the parameters. + + Regression for the pre-cutover OMOL model gating ``create_graph`` on + ``self.training``: with the model in its default eval mode the autograd + force was detached from the parameter graph, so a force loss produced + zero gradient for every parameter (and ``loss.backward()`` raised). The + robust form keeps the force in the graph whenever grad is enabled, so a + strict majority of parameters receive a gradient. See spec + ``cuet-force-doublebackward`` Findings (run 2). + + The readout's output layers are zero-initialised (MACE starts at the E0 + baseline), which makes a fresh model's energy position-independent and + its forces identically zero — a degenerate state in which no force loss + can reach any parameter. Perturb the parameters first so the model + produces real forces, then assert the loss reaches them. + """ + pos, Z, edge_index, batch, _ = raw_tensors(charged_omol_cluster) + total_charge = charged_omol_cluster["graphs", "total_charge"] + total_spin = charged_omol_cluster["graphs", "total_spin"] + assert not omol_variant.training # default eval mode — the regression condition + + _perturb(omol_variant) + omol_variant.zero_grad(set_to_none=True) + forces = omol_variant.energy_forces(pos, Z, edge_index, batch, total_charge, total_spin)[ + "forces" + ] + assert forces.abs().max() > 0.0, "perturbed model still produces zero forces" + + (forces**2).mean().backward() + + with_grad, total = _fraction_with_gradient(omol_variant) + assert with_grad > total // 2, ( + f"force loss reached only {with_grad}/{total} parameters; " + "the force is detached from the parameter graph" + ) + + def test_force_loss_through_forward_reaches_parameters( + self, omol_variant: MACEOMol, charged_omol_cluster: TensorDict + ) -> None: + """The molnex-pipeline ``forward`` also unlocks force-supervised training. + + ``forward`` derives forces through the shared batch-level force kernel + (``molpot.derivation.kernels``) rather than the variant's own + ``energy_forces`` leaf, so it is a second path that must stay connected + to the parameters. Before the port it raised ``setup_context`` because + ``torch.func.grad`` cannot trace cuEquivariance's fused custom ops; see + spec ``cuet-force-doublebackward`` Resolution (run 3). + """ + _perturb(omol_variant) + + out = omol_variant.forward(charged_omol_cluster) + forces = out["atoms", "forces"] + assert forces.abs().max() > 0.0 + + omol_variant.zero_grad(set_to_none=True) + (forces**2).mean().backward() + + with_grad, total = _fraction_with_gradient(omol_variant) + assert with_grad > total // 2, ( + f"force loss through forward reached only {with_grad}/{total} parameters" + ) + + def test_batched_graphs(self, omol_variant: MACEOMol, omol_pair_batch: TensorDict) -> None: + """Two molecules in one batch: per-graph energies, correct atom routing.""" + pos, Z, edge_index, batch, _ = raw_tensors(omol_pair_batch) + reference = omol_variant.energy_forces( + pos, + Z, + edge_index, + batch, + omol_pair_batch["graphs", "total_charge"], + omol_pair_batch["graphs", "total_spin"], + ) + out = omol_variant.forward(omol_pair_batch) + assert out["graphs", "energy"].shape == (2,) + assert torch.allclose( + out["graphs", "energy"], reference["energy"], atol=FORWARD_ENERGY_ATOL, rtol=0 + ) + assert torch.allclose( + out["atoms", "forces"], reference["forces"], atol=FORWARD_FORCE_ATOL, rtol=0 + ) + assert out["atoms", "forces"].shape == (len(PAIR_Z), 3) diff --git a/tests/test_molzoo/test_mace_encoder.py b/tests/test_molzoo/test_mace_encoder.py deleted file mode 100644 index 97d39b1..0000000 --- a/tests/test_molzoo/test_mace_encoder.py +++ /dev/null @@ -1,75 +0,0 @@ -"""Tests for encoder-only MACE API.""" - -from __future__ import annotations - -import pytest -import torch -from tensordict import TensorDict - -from molrep.embedding.node import DiscreteEmbeddingSpec -from molzoo import MACE - - -@pytest.fixture -def graph_data(): - n_nodes = 5 - edge_index = torch.tensor( - [ - [0, 1], - [1, 0], - [1, 2], - [2, 1], - [2, 3], - [3, 2], - [3, 4], - [4, 3], - ], - dtype=torch.long, - ) - pos = torch.randn(n_nodes, 3) - edge_diff = pos[edge_index[:, 1]] - pos[edge_index[:, 0]] - edge_dist = edge_diff.norm(dim=-1).clamp(min=1e-4) - n_edges = edge_index.shape[0] - - atoms = TensorDict( - Z=torch.randint(0, 6, (n_nodes,)), - pos=pos, - batch=torch.zeros(n_nodes, dtype=torch.long), - batch_size=[n_nodes], - ) - edges = TensorDict( - edge_index=edge_index, - edge_diff=edge_diff, - edge_dist=edge_dist, - batch_size=[n_edges], - ) - return TensorDict(atoms=atoms, edges=edges, batch_size=[]) - - -def _build_encoder() -> MACE: - return MACE( - node_attr_specs=[ - DiscreteEmbeddingSpec( - input_key="Z", - num_classes=6, - emb_dim=16, - ) - ], - num_elements=6, - num_features=16, - r_max=5.0, - num_interactions=2, - l_max=2, - ) - - -class TestMACE: - """Full MACE encoder contract and compile compatibility.""" - - def test_forward_encoder_contract(self, graph_data): - encoder = _build_encoder() - output = encoder(graph_data) - node_features = output["atoms", "node_features"] - n_nodes = graph_data["atoms", "Z"].shape[0] - assert isinstance(node_features, torch.Tensor) - assert node_features.shape == (n_nodes, 2, 16) diff --git a/tests/test_molzoo/test_mace_omol.py b/tests/test_molzoo/test_mace_omol.py deleted file mode 100644 index cc7bba0..0000000 --- a/tests/test_molzoo/test_mace_omol.py +++ /dev/null @@ -1,234 +0,0 @@ -"""Integration tests for the MACEOMol molnex-pipeline adapter. - -Covers ``mace-omol-port-02-pipeline-integration`` ac-001: ``MACEOMol.forward`` -consumes the post-collate ``atoms / edges / graphs`` TensorDict, routes forces -through ``molpot.derivation.ForceDerivation``, and agrees with the raw-tensor -``MACEOMol.energy_forces`` entry point to machine precision. - -The full-OMOL energy/force faithfulness vs the official model is covered by -``mace-omol-port-01`` (scripts/omol_port/verify_e2e.py); here we only assert the -two molnex entry points are consistent, on a small l_max=1 instance. -""" - -from __future__ import annotations - -import pytest -import torch -from tensordict import TensorDict - -from molzoo.mace_omol import MACEOMol - - -def _full_edges(batch: torch.Tensor) -> torch.Tensor: - """All intra-graph ordered pairs as ``(E, 2)`` ``[source, target]`` edges.""" - src, dst = [], [] - n = batch.shape[0] - for i in range(n): - for j in range(n): - if i != j and batch[i] == batch[j]: - src.append(i) - dst.append(j) - return torch.tensor([src, dst], dtype=torch.long).t().contiguous() - - -def _make_batch(pos, Z, batch, edge_index_e2, total_charge, total_spin) -> TensorDict: - """Assemble a post-collate-style TensorDict (atoms / edges / graphs).""" - src = edge_index_e2[:, 0] - dst = edge_index_e2[:, 1] - edge_diff = pos[dst] - pos[src] - edge_dist = torch.linalg.norm(edge_diff, dim=-1) - n = Z.shape[0] - e = edge_index_e2.shape[0] - b = total_charge.shape[0] - return TensorDict( - { - "atoms": TensorDict({"Z": Z, "pos": pos, "batch": batch}, batch_size=[n]), - "edges": TensorDict( - { - "edge_index": edge_index_e2, - "edge_diff": edge_diff, - "edge_dist": edge_dist, - }, - batch_size=[e], - ), - "graphs": TensorDict( - {"total_charge": total_charge, "total_spin": total_spin}, - batch_size=[b], - ), - }, - batch_size=[], - ) - - -@pytest.fixture -def model(): - """Small fp64 MACEOMol (l_max=1) — non-OMOL dims, fast on CPU.""" - torch.manual_seed(0) - ae = torch.tensor([-13.6, -1029.0, -2041.0]) - m = MACEOMol( - atomic_numbers=[1, 6, 8], - atomic_energies=ae, - r_max=5.0, - num_bessel=8, - l_max=1, - num_features=64, - num_interactions=2, - correlation=2, - ).double() - return m.eval() - - -@pytest.fixture -def single_graph(): - torch.manual_seed(1) - pos = torch.randn(5, 3, dtype=torch.float64) * 1.5 - Z = torch.tensor([1, 6, 8, 1, 1]) - batch = torch.zeros(5, dtype=torch.long) - total_charge = torch.tensor([1], dtype=torch.long) # charged molecule - total_spin = torch.tensor([0], dtype=torch.long) - return pos, Z, batch, total_charge, total_spin - - -def test_forward_matches_energy_forces(model, single_graph): - """forward(td) energy/forces == raw energy_forces() to machine precision.""" - pos, Z, batch, tc, ts = single_graph - edge_e2 = _full_edges(batch) - - ref = model.energy_forces(pos, Z, edge_e2.t().contiguous(), batch, tc, ts) - - td = _make_batch(pos.clone(), Z, batch, edge_e2, tc, ts) - out = model.forward(td) - - assert torch.allclose(out["graphs", "energy"], ref["energy"], atol=1e-9, rtol=0) - assert torch.allclose(out["atoms", "forces"], ref["forces"], atol=1e-8, rtol=0) - - -def test_forward_returns_same_td_with_new_keys(model, single_graph): - """forward mutates in place and returns the same TensorDict object.""" - pos, Z, batch, tc, ts = single_graph - td = _make_batch(pos, Z, batch, _full_edges(batch), tc, ts) - out = model.forward(td) - assert out is td - assert ("graphs", "energy") in out.keys(include_nested=True) - assert ("atoms", "forces") in out.keys(include_nested=True) - - -def test_forces_translation_invariant(model, single_graph): - """Net force on an isolated molecule is ~zero (translation invariance).""" - pos, Z, batch, tc, ts = single_graph - td = _make_batch(pos, Z, batch, _full_edges(batch), tc, ts) - out = model.forward(td) - net = out["atoms", "forces"].sum(0) - assert net.abs().max() < 1e-7 - - -def test_missing_charge_spin_defaults_to_neutral(model, single_graph): - """Absent graphs.total_charge/total_spin → neutral singlet, still runs.""" - pos, Z, batch, _, _ = single_graph - edge_e2 = _full_edges(batch) - td = TensorDict( - { - "atoms": TensorDict({"Z": Z, "pos": pos, "batch": batch}, batch_size=[5]), - "edges": TensorDict({"edge_index": edge_e2}, batch_size=[edge_e2.shape[0]]), - }, - batch_size=[], - ) - out = model.forward(td) - # forward defaults to the OMOL neutral closed-shell singlet: charge=0, spin=1 - # (spin index 1 is a trained row; spin 0 hits an untrained embedding). - charge = torch.zeros(1, dtype=torch.long) - spin = torch.ones(1, dtype=torch.long) - ref = model.energy_forces(pos, Z, edge_e2.t().contiguous(), batch, charge, spin) - assert torch.allclose(out["graphs", "energy"], ref["energy"], atol=1e-9, rtol=0) - - -def test_force_loss_reaches_parameters_in_eval_mode(model, single_graph): - """Force-supervised training must backprop to the parameters. - - Regression for ``mace_omol.py`` ``energy_forces`` gating ``create_graph`` on - ``self.training``: with the model in its default eval mode the autograd - force was detached from the parameter graph, so a force loss produced zero - gradient for every parameter (and ``loss.backward()`` raised). The robust - form keeps the force in the graph whenever grad is enabled, so a strict - majority of parameters receive a gradient. See spec - ``cuet-force-doublebackward`` Findings (run 2). - - The readout's output layers are zero-initialised (MACE starts at the E0 - baseline), which makes a fresh model's energy position-independent and its - forces identically zero — a degenerate state in which no force loss can - reach any parameter. Perturb the parameters first so the model produces real - forces, then assert the loss reaches them. - """ - pos, Z, batch, tc, ts = single_graph - edge_index = _full_edges(batch).t().contiguous() - assert not model.training # default eval mode — the regression condition - - torch.manual_seed(7) - with torch.no_grad(): - for p in model.parameters(): - p.add_(torch.randn_like(p) * 0.05) - - model.zero_grad(set_to_none=True) - out = model.energy_forces(pos, Z, edge_index, batch, tc, ts, compute_forces=True) - forces = out["forces"] - assert forces.abs().max() > 0.0, "perturbed model still produces zero forces" - - loss = (forces**2).mean() - loss.backward() - - params = list(model.parameters()) - n_with_grad = sum(int(p.grad is not None and float(p.grad.abs().sum()) > 0.0) for p in params) - assert n_with_grad > len(params) // 2, ( - f"force loss reached only {n_with_grad}/{len(params)} parameters; " - "the force is detached from the parameter graph" - ) - - -def test_functorch_force_loss_trains_via_forward(model, single_graph): - """The molnex-pipeline ``forward`` (functorch ``ForceDerivation``) unlocks - force-supervised training. - - ``forward`` derives forces with ``torch.func.grad`` (compile-friendly, single - backward). That transform cannot trace cuEquivariance's fused custom ops, so - before the functorch port it raised ``setup_context``; now the spherical - harmonics are pure-torch and the equivariant tensor products run on their - ``use_fallback`` path, so a force loss backprops to the parameters. See spec - ``cuet-force-doublebackward`` Resolution (run 3, functorch port). - """ - pos, Z, batch, tc, ts = single_graph - - torch.manual_seed(7) - with torch.no_grad(): - for p in model.parameters(): - p.add_(torch.randn_like(p) * 0.05) - - td = _make_batch(pos, Z, batch, _full_edges(batch), tc, ts) - out = model.forward(td) - forces = out["atoms", "forces"] - assert forces.abs().max() > 0.0 - - model.zero_grad(set_to_none=True) - (forces**2).mean().backward() - - params = list(model.parameters()) - n_with_grad = sum(int(p.grad is not None and float(p.grad.abs().sum()) > 0.0) for p in params) - assert n_with_grad > len(params) // 2, ( - f"functorch force loss reached only {n_with_grad}/{len(params)} parameters" - ) - - -def test_batched_graphs(model): - """Two molecules in one batch: per-graph energies, correct atom routing.""" - torch.manual_seed(2) - pos = torch.randn(7, 3, dtype=torch.float64) * 1.5 - Z = torch.tensor([1, 6, 8, 1, 8, 6, 1]) - batch = torch.tensor([0, 0, 0, 0, 1, 1, 1]) - tc = torch.tensor([0, -1], dtype=torch.long) - ts = torch.tensor([0, 1], dtype=torch.long) - edge_e2 = _full_edges(batch) - td = _make_batch(pos, Z, batch, edge_e2, tc, ts) - out = model.forward(td) - ref = model.energy_forces(pos, Z, edge_e2.t().contiguous(), batch, tc, ts) - assert out["graphs", "energy"].shape == (2,) - assert torch.allclose(out["graphs", "energy"], ref["energy"], atol=1e-9, rtol=0) - assert torch.allclose(out["atoms", "forces"], ref["forces"], atol=1e-8, rtol=0) diff --git a/tests/test_molzoo/test_pinet/test_encoder.py b/tests/test_molzoo/test_pinet/test_encoder.py index c216a10..d429619 100644 --- a/tests/test_molzoo/test_pinet/test_encoder.py +++ b/tests/test_molzoo/test_pinet/test_encoder.py @@ -71,6 +71,43 @@ def test_rank_output_shapes(rank): assert out["edges", "i5_features"].shape == (8, 2, 5, 8) +@pytest.mark.parametrize("rank", [1, 3, 5]) +def test_emit_property_features_off_drops_only_property_tracks(rank): + """Opting out skips the property-head tracks; scalar outputs are unchanged.""" + torch.manual_seed(0) + lean = PiNet( + atom_types=[1, 6, 7, 8], + r_max=4.0, + n_basis=3, + pp_nodes=[8, 8], + pi_nodes=[8, 8], + ii_nodes=[8, 8], + depth=2, + rank=rank, + emit_property_features=False, + ) + lean.eval() + full = _encoder(rank=rank) # same seed/architecture, emission on + + out_lean = lean(_graph()) + out_full = full(_graph()) + + # The contract keys survive and are bit-identical to the emitting encoder. + for key in ("node_features", "p1_block_outputs"): + assert torch.equal(out_lean["atoms", key], out_full["atoms", key]) + + # Only the property-head tracks are gone. + assert "i1_features" not in out_lean["edges"].keys() + for key in ("p3_features", "p5_features"): + assert key not in out_lean["atoms"].keys() + for key in ("i3_features", "i5_features"): + assert key not in out_lean["edges"].keys() + + +def test_emit_property_features_defaults_on(): + assert PiNet(atom_types=[1, 6]).emit_property_features is True + + def test_translation_invariance(): enc = _encoder(rank=5) g = _graph() diff --git a/tests/test_molzoo/test_pinet/test_potential.py b/tests/test_molzoo/test_pinet/test_potential.py index 4086385..7f7d5e3 100644 --- a/tests/test_molzoo/test_pinet/test_potential.py +++ b/tests/test_molzoo/test_pinet/test_potential.py @@ -44,7 +44,7 @@ def _batch( ) -def _model() -> PiNetPotential: +def _model(*, compute_forces: bool = True) -> PiNetPotential: torch.manual_seed(0) return PiNetPotential( atom_types=[1, 6, 7, 8], @@ -56,6 +56,7 @@ def _model() -> PiNetPotential: depth=3, rank=3, hidden_dim=16, + compute_forces=compute_forces, ).to(DEVICE) @@ -90,9 +91,9 @@ def test_energy_and_force_shapes(): graphs=TensorDict(num_atoms=torch.tensor([n]), batch_size=[1]), batch_size=[], ) - out = model(batch, compute_forces=True) - assert out["energy"].shape == (1,) - assert out["forces"].shape == (n, 3) + batch = model(batch) + assert batch["graphs", "energy"].shape == (1,) + assert batch["atoms", "forces"].shape == (n, 3) def test_encoder_kwarg_composition(): @@ -110,57 +111,63 @@ def test_encoder_kwarg_composition(): assert model.encoder is enc -def test_functorch_forces_match_autograd_reference(): - model = _model() +def test_func_forces_match_autograd_reference(): + model = _model(compute_forces=True) model.eval() batch = _batch() - f_functorch = model(batch.clone(), compute_forces=True)["forces"].detach() + b1 = model(batch.clone()) + f_func = b1["atoms", "forces"].detach() + + # Energy-only model for the reference path (same architecture / weights). + ref_model = _model(compute_forces=False) + ref_model.load_state_dict( + {k: v for k, v in model.state_dict().items()}, + strict=False, + ) + # Copy only shared params (force pipeline has no extra params). + ref_model.load_state_dict(model.state_dict()) + ref_model.eval() ref = batch.clone() pos = ref["atoms", "pos"].detach().clone().requires_grad_(True) ref["atoms", "pos"] = pos with torch.enable_grad(): - energy = model(ref, compute_forces=False)["energy"].sum() + ref = ref_model(ref) + energy = ref["graphs", "energy"].sum() f_autograd = -torch.autograd.grad(energy, pos)[0].detach() - assert f_functorch.shape == f_autograd.shape - assert torch.allclose(f_functorch, f_autograd, atol=1e-5, rtol=1e-5), ( - f"max abs diff {(f_functorch - f_autograd).abs().max().item():.3e}" + assert f_func.shape == f_autograd.shape + assert torch.allclose(f_func, f_autograd, atol=1e-5, rtol=1e-5), ( + f"max abs diff {(f_func - f_autograd).abs().max().item():.3e}" ) -def test_eval_single_pass_matches_train_two_pass(): - model = _model() +def test_eval_and_train_forward_match(): + model = _model(compute_forces=True) torch.manual_seed(1) batch = _batch() model.train() - out_train = model(batch.clone(), compute_forces=True) + out_train = model(batch.clone()) model.eval() - out_eval = model(batch.clone(), compute_forces=True) + out_eval = model(batch.clone()) - assert set(out_eval) == set(out_train) - for key in ("energy", "atomic_energy", "forces"): - assert torch.allclose(out_eval[key], out_train[key].detach(), atol=1e-6, rtol=1e-6), ( - f"{key} diverges between eval single-pass and train two-pass" - ) + for key in (("graphs", "energy"), ("atoms", "energy"), ("atoms", "forces")): + assert torch.allclose(out_eval[key], out_train[key].detach(), atol=1e-6, rtol=1e-6), key def test_force_loss_single_backward_populates_param_grads(): - model = _model() + model = _model(compute_forces=True) model.train() - batch = _batch() - - out = model(batch.clone(), compute_forces=True) - loss = (out["forces"] - batch["atoms", "forces"]).pow(2).mean() - loss.backward() + batch = model(_batch().clone()) + batch["atoms", "forces"].pow(2).mean().backward() n_grad = sum(1 for p in model.parameters() if p.grad is not None and p.grad.abs().sum() > 0) assert n_grad > 0, "no parameter received a gradient from the force loss" def test_no_lazy_linear_in_potential(): - model = _model() + model = _model(compute_forces=True) for m in model.modules(): assert type(m).__name__ != "LazyLinear" diff --git a/tests/test_molzoo/test_pinet/test_spec.py b/tests/test_molzoo/test_pinet/test_spec.py index 006816f..b5819ba 100644 --- a/tests/test_molzoo/test_pinet/test_spec.py +++ b/tests/test_molzoo/test_pinet/test_spec.py @@ -17,4 +17,4 @@ def test_defaults_build_encoder(): def test_rejects_bad_rank(): with pytest.raises(ValidationError): - PiNetSpec(rank=2) # type: ignore[arg-type] + PiNetSpec(rank=2) diff --git a/tests/test_scripts/test_mm_param_learning/__init__.py b/tests/test_scripts/test_mm_param_learning/__init__.py new file mode 100644 index 0000000..c8dff75 --- /dev/null +++ b/tests/test_scripts/test_mm_param_learning/__init__.py @@ -0,0 +1 @@ +"""Tests for scripts/mm_param_learning workspace scaffolding.""" diff --git a/tests/test_scripts/test_mm_param_learning/test_constants.py b/tests/test_scripts/test_mm_param_learning/test_constants.py new file mode 100644 index 0000000..899c8c6 --- /dev/null +++ b/tests/test_scripts/test_mm_param_learning/test_constants.py @@ -0,0 +1,36 @@ +"""Unit tests for mm_param_learning constants.""" + +from __future__ import annotations + +from scripts.mm_param_learning.constants import ( + EXPERIMENT_SLUGS, + MOLHUB_COORDINATES, + VALIDATION_STAGES, +) + + +class TestConstants: + def test_four_experiment_slugs(self): + assert set(EXPERIMENT_SLUGS) == { + "potential-parity", + "zinc-typing-recovery", + "phalkethoh-mm-energy", + "latent-analysis", + } + assert set(MOLHUB_COORDINATES) == set(EXPERIMENT_SLUGS) + + def test_coordinates_are_dataset_prefixed(self): + for slug, coord in MOLHUB_COORDINATES.items(): + assert coord.startswith("dataset:"), (slug, coord) + + def test_six_validation_stages(self): + assert len(VALIDATION_STAGES) == 6 + nums = [n for n, _ in VALIDATION_STAGES] + assert nums == ["1", "2", "3", "4", "5", "6"] + titles = " ".join(t for _, t in VALIDATION_STAGES).lower() + assert "inventory" in titles + assert "b0" in titles or "parity" in titles + assert "typing" in titles + assert "energy" in titles or "b1" in titles + assert "latent" in titles + assert "gate" in titles or "synthesis" in titles diff --git a/tests/test_scripts/test_mm_param_learning/test_materialize_workspace.py b/tests/test_scripts/test_mm_param_learning/test_materialize_workspace.py new file mode 100644 index 0000000..dae7605 --- /dev/null +++ b/tests/test_scripts/test_mm_param_learning/test_materialize_workspace.py @@ -0,0 +1,74 @@ +"""Unit tests for MmParamLearningWorkspace.materialize (temp root only).""" + +from __future__ import annotations + +from pathlib import Path + +import pytest + +pytest.importorskip("molexp") + +from scripts.mm_param_learning.constants import ( + DEFAULT_WORKSPACE_ROOT, + EXPERIMENT_SLUGS, + KNOWLEDGE_NOTE_NAME, + MOLHUB_COORDINATES, + PROGRAM, + PROJECT_SLUG, +) +from scripts.mm_param_learning.materialize_workspace import MmParamLearningWorkspace + + +class TestMmParamLearningWorkspace: + def test_materialize_temp_root(self, tmp_path: Path): + root = tmp_path / "ws" + summary = MmParamLearningWorkspace(root=root).materialize() + assert (root / "workspace.json").is_file() + assert summary["project"] == PROJECT_SLUG + assert set(summary["experiments"]) == set(EXPERIMENT_SLUGS) + # Project + four experiments + project_dir = root / "projects" / PROJECT_SLUG + assert project_dir.is_dir() + exp_dir = project_dir / "experiments" + assert set(p.name for p in exp_dir.iterdir() if p.is_dir()) == set(EXPERIMENT_SLUGS) + # Seed run params (molexp may prefix run ids with ``run-``). + for slug in EXPERIMENT_SLUGS: + runs_dir = exp_dir / slug / "runs" + run_dirs = [p for p in runs_dir.iterdir() if p.is_dir() and "seed" in p.name] + assert len(run_dirs) == 1, run_dirs + run_json = run_dirs[0] / "run.json" + assert run_json.is_file() + text = run_json.read_text() + assert "dataset_coordinate" in text + assert PROGRAM in text or "mm-param-learning-baseline" in text + assert MOLHUB_COORDINATES[slug] in text or slug in text + + def test_idempotent(self, tmp_path: Path): + root = tmp_path / "ws" + MmParamLearningWorkspace(root=root).materialize() + MmParamLearningWorkspace(root=root).materialize() + exp_dir = root / "projects" / PROJECT_SLUG / "experiments" + assert len([p for p in exp_dir.iterdir() if p.is_dir()]) == 4 + # Knowledge note not duplicated (single slug dir) + notes = list(root.rglob(KNOWLEDGE_NOTE_NAME)) + # at most one concept directory with that name + note_dirs = [p for p in notes if p.is_dir()] + assert len(note_dirs) == 1 + + def test_knowledge_note_content(self, tmp_path: Path): + root = tmp_path / "ws" + MmParamLearningWorkspace(root=root).materialize() + bodies = list(root.rglob("index.md")) + assert bodies, "expected Knowledge Note index.md" + body = "\n".join(p.read_text() for p in bodies) + for num in ("1", "2", "3", "4", "5", "6"): + assert f"Stage {num}" in body + for coord in MOLHUB_COORDINATES.values(): + assert coord in body + + def test_explicit_root_not_default(self, tmp_path: Path): + root = tmp_path / "only-here" + MmParamLearningWorkspace(root=root).materialize() + assert (root / "workspace.json").is_file() + # Must not touch operator default from this call. + assert root.resolve() != DEFAULT_WORKSPACE_ROOT.resolve() diff --git a/tests/test_scripts/test_mm_param_learning/test_workflows.py b/tests/test_scripts/test_mm_param_learning/test_workflows.py new file mode 100644 index 0000000..c71cd76 --- /dev/null +++ b/tests/test_scripts/test_mm_param_learning/test_workflows.py @@ -0,0 +1,25 @@ +"""Unit tests for MmParamWorkflows stubs.""" + +from __future__ import annotations + +import pytest + +pytest.importorskip("molexp") + +from scripts.mm_param_learning.constants import EXPERIMENT_SLUGS +from scripts.mm_param_learning.workflows import MmParamWorkflows + + +class TestMmParamWorkflows: + def test_three_tasks_per_slug(self): + wf = MmParamWorkflows() + for slug in EXPERIMENT_SLUGS: + names = wf.task_names(slug) + assert names == MmParamWorkflows.TASK_NAMES + assert "resolve_molhub_coordinates" in names + assert "record_run_params" in names + assert "write_placeholder_metrics" in names + + def test_unknown_slug_raises(self): + with pytest.raises(ValueError, match="Unknown experiment slug"): + MmParamWorkflows().build("not-a-real-experiment") diff --git a/zensical.toml b/zensical.toml index 0ff25f1..e74cc10 100644 --- a/zensical.toml +++ b/zensical.toml @@ -1,7 +1,7 @@ [project] site_name = "MolNex Documentation" site_description = "Dict-first molecular ML framework" -site_url = "https://molcrafts.github.io/molnex/" +site_url = "https://docs.molcrafts.org/molnex/" repo_url = "https://github.com/molcrafts/molnex" repo_name = "molcrafts/molnex" edit_uri = "edit/main/docs" @@ -26,11 +26,14 @@ nav = [ { "Hooks" = "molix/user-guide/hooks.md" }, { "Data Pipeline" = "molix/user-guide/data.md" }, { "Data Loading" = "molix/user-guide/data-loading.md" }, - { "Data Modules" = "molix/user-guide/data-modules.md" } + { "Data Modules" = "molix/user-guide/data-modules.md" }, + { "Molecular Dynamics" = "molix/user-guide/md.md" }, + { "Profiling" = "molix/user-guide/profiling.md" } ]}, { "Explanation" = [ { "Execution Model" = "molix/explanation/execution-model.md" }, - { "Batch Schema" = "molix/explanation/batch-schema.md" } + { "Batch Schema" = "molix/explanation/batch-schema.md" }, + { "Training Throughput & torch.compile" = "molix/explanation/throughput-and-compilation.md" } ]} ]}, { "MolRep" = [ @@ -79,84 +82,30 @@ nav = [ ]} ] +# Shared MolCrafts theme: brand palette, light/dark schemes, navigation +# features, and fonts. Do not re-list features/palette/fonts here. [project.theme] -name = "material" -features = [ - "navigation.sections", - "navigation.indexes", - "content.code.copy", - "search.highlight", -] - -[project.theme.palette] -scheme = "default" -primary = "indigo" -accent = "blue" - -[project.theme.font] -text = "Inter" -code = "JetBrains Mono" - -[project.markdown_extensions.abbr] - -[project.markdown_extensions.admonition] - -[project.markdown_extensions.attr_list] +name = "molcrafts" +language = "en" -[project.markdown_extensions.def_list] +[project.extra.molcrafts] +product = "molnex" +accent = "#db2777" # primary — molnex signature colour +accent_soft = "rgba(219, 39, 119, 0.14)" # secondary — soft fill behind it -[project.markdown_extensions.footnotes] - -[project.markdown_extensions.md_in_html] - -[project.markdown_extensions.tables] - -[project.markdown_extensions.toc] -permalink = true +[[project.extra.social]] +icon = "fontawesome/brands/github" +link = "https://github.com/MolCrafts/molnex" +# Markdown: keep arithmatex for MathJax; other defaults come from Zensical. [project.markdown_extensions."pymdownx.arithmatex"] generic = true -[project.markdown_extensions."pymdownx.betterem"] - -[project.markdown_extensions."pymdownx.caret"] - -[project.markdown_extensions."pymdownx.details"] - -[project.markdown_extensions."pymdownx.highlight"] -anchor_linenums = true -line_spans = "__span" -pygments_lang_class = true -use_pygments = false - -[project.markdown_extensions."pymdownx.inlinehilite"] - -[project.markdown_extensions."pymdownx.keys"] - -[project.markdown_extensions."pymdownx.magiclink"] - -[project.markdown_extensions."pymdownx.mark"] - -[project.markdown_extensions."pymdownx.smartsymbols"] - -[project.markdown_extensions."pymdownx.superfences"] - -[project.markdown_extensions."pymdownx.tabbed"] -alternate_style = true -combine_header_slug = true - -[project.markdown_extensions."pymdownx.tasklist"] -custom_checkbox = true - -[project.markdown_extensions."pymdownx.tilde"] - [project.plugins.search] enabled = true [project.plugins.mkdocstrings] default_handler = "python" -inventory_project = "molnex" -inventory_version = "0.1.0" [project.plugins.mkdocstrings.handlers.python] paths = ["src"]