feat: MM parameter learning baseline (molnex)

.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`.)*
 
 <!-- mol:map:managed end -->

.claude/notes/learnable-classical-ff.md

@@ -0,0 +1,136 @@
+<!-- mol:note:topic:learnable-classical-ff -->
+# 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).