matterlens 0.0.1__tar.gz

matterlens-0.0.1/.github/workflows/ci.yml

@@ -0,0 +1,16 @@
+name: CI
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+
+jobs:
+  check:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: astral-sh/setup-uv@v6
+      - run: uv build
+      - run: uvx ruff check src
+      - run: uv run --isolated --no-project --python 3.12 --with dist/*.whl python -c "import matterlens"

matterlens-0.0.1/.github/workflows/publish.yml

@@ -0,0 +1,41 @@
+name: Publish to PyPI
+
+# Builds with uv and publishes via PyPI trusted publishing (OIDC), no API token.
+# One-time setup on pypi.org: Publishing -> add a pending publisher for project
+# "matterlens" with owner Garden-AI, repository matterlens, workflow
+# publish.yml, environment pypi. Then push a tag like v0.0.1.
+
+on:
+  push:
+    tags: ["v*"]
+  workflow_dispatch:
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: astral-sh/setup-uv@v6
+      - name: Build sdist and wheel
+        run: uv build
+      - name: Smoke-test the wheel in a clean environment
+        run: uv run --isolated --no-project --python 3.12 --with dist/*.whl python -c "import matterlens; print(matterlens.__version__)"
+      - uses: actions/upload-artifact@v4
+        with:
+          name: dist
+          path: dist/
+
+  publish:
+    needs: build
+    runs-on: ubuntu-latest
+    environment: pypi
+    permissions:
+      id-token: write
+    steps:
+      - uses: actions/download-artifact@v4
+        with:
+          name: dist
+          path: dist/
+      - uses: astral-sh/setup-uv@v6
+      - name: Publish
+        run: uv publish

matterlens-0.0.1/.gitignore

@@ -0,0 +1,13 @@
+__pycache__/
+*.pyc
+.venv/
+dist/
+build/
+*.egg-info/
+.pytest_cache/
+.ruff_cache/
+*.npz
+*.pt
+scratch/
+docs/draft.tex
+.obsidian/

matterlens-0.0.1/PKG-INFO

@@ -0,0 +1,66 @@
+Metadata-Version: 2.5
+Name: matterlens
+Version: 0.0.1
+Summary: Hooks, caches, and lenses for machine-learning interatomic potentials and other geometric GNNs (the TransformerLens idea, for MLIPs).
+Project-URL: Homepage, https://github.com/Garden-AI/matterlens
+Project-URL: Repository, https://github.com/Garden-AI/matterlens
+Author-email: Owen Price Skelly <owenpriceskelly@uchicago.edu>
+License: MIT
+Keywords: equivariant,gnn,interatomic-potential,interpretability,mace,mlip,uma
+Classifier: Development Status :: 1 - Planning
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Topic :: Scientific/Engineering :: Chemistry
+Classifier: Topic :: Scientific/Engineering :: Physics
+Requires-Python: >=3.11
+Requires-Dist: ase>=3.23
+Requires-Dist: numpy
+Requires-Dist: torch>=2.4
+Provides-Extra: dev
+Requires-Dist: pytest; extra == 'dev'
+Requires-Dist: ruff; extra == 'dev'
+Provides-Extra: fairchem
+Requires-Dist: fairchem-core>=2.20; extra == 'fairchem'
+Provides-Extra: mace
+Requires-Dist: mace-torch>=0.3.15; extra == 'mace'
+Requires-Dist: torch<2.10; extra == 'mace'
+Provides-Extra: orb
+Requires-Dist: orb-models<0.8,>=0.6.2; extra == 'orb'
+Provides-Extra: sevenn
+Requires-Dist: sevenn>=0.10.0; extra == 'sevenn'
+Description-Content-Type: text/markdown
+
+# matterlens
+
+Hooks, caches, and lenses for machine-learning interatomic potentials and
+other geometric GNNs. The TransformerLens idea, built for models whose
+activations are per-atom irreps rather than a token residual stream.
+
+Status (2026-09-10): design stage. Read `docs/v0-design.md` for what the first
+release is and why. `docs/architecture-survey.md` records what MACE, SevenNet,
+Orb, and UMA module trees actually look like; `docs/prior-art.md` records the
+landscape. `scripts/spike_hooked_mace.py` is a runnable feasibility spike.
+
+Intended v0 usage:
+
+```python
+from matterlens import load
+hm = load("mace", "mace-mp-0-medium")          # native mace-torch model underneath
+batch = hm.from_atoms(atoms)
+out, cache = hm.run_with_cache(batch)
+cache["blocks.1.node_feats"].shape             # [n_atoms, 256] as 128x0e+128x1o
+cache.blocks("blocks.1.node_feats")[1]         # [n_atoms, 128, 3] the vector channels
+with hm.hooks({"blocks.0.node_feats": lambda a, hp: torch.zeros_like(a)}):
+    ablated = hm.model(batch)
+```
+
+Part of the Amatterlens program (interpretability for scientific discovery),
+workstream 1. Companion to Rootstock, which supplies the isolated
+per-family environments each adapter installs into.
+
+Run the spike:
+
+```bash
+uv run scripts/spike_hooked_mace.py
+```

matterlens-0.0.1/README.md

@@ -0,0 +1,33 @@
+# matterlens
+
+Hooks, caches, and lenses for machine-learning interatomic potentials and
+other geometric GNNs. The TransformerLens idea, built for models whose
+activations are per-atom irreps rather than a token residual stream.
+
+Status (2026-09-10): design stage. Read `docs/v0-design.md` for what the first
+release is and why. `docs/architecture-survey.md` records what MACE, SevenNet,
+Orb, and UMA module trees actually look like; `docs/prior-art.md` records the
+landscape. `scripts/spike_hooked_mace.py` is a runnable feasibility spike.
+
+Intended v0 usage:
+
+```python
+from matterlens import load
+hm = load("mace", "mace-mp-0-medium")          # native mace-torch model underneath
+batch = hm.from_atoms(atoms)
+out, cache = hm.run_with_cache(batch)
+cache["blocks.1.node_feats"].shape             # [n_atoms, 256] as 128x0e+128x1o
+cache.blocks("blocks.1.node_feats")[1]         # [n_atoms, 128, 3] the vector channels
+with hm.hooks({"blocks.0.node_feats": lambda a, hp: torch.zeros_like(a)}):
+    ablated = hm.model(batch)
+```
+
+Part of the Amatterlens program (interpretability for scientific discovery),
+workstream 1. Companion to Rootstock, which supplies the isolated
+per-family environments each adapter installs into.
+
+Run the spike:
+
+```bash
+uv run scripts/spike_hooked_mace.py
+```

matterlens-0.0.1/docs/architecture-survey.md

@@ -0,0 +1,203 @@
+# Architecture survey: what the module trees actually look like
+
+Observed 2026-09-10 by loading each pretrained model on CPU (macOS arm64) and
+registering forward hooks on every submodule (depth 1 to 3 depending on family). Shapes are for
+the test structure named in each section. UMA was loaded with a
+HuggingFace token on the second pass (2026-09-10). This is the ground truth the adapter
+tables in `docs/v0-design.md` rest on; re-run the introspection scripts when
+bumping a package.
+
+Package versions: mace-torch 0.3.x (latest as of run), sevenn 0.10.x, orb-models
+0.7.x, fairchem-core 2.x (latest PyPI release as of run).
+
+## MACE (`mace-torch`), model `mace-mp-0 small`, H2O in a box (3 atoms, 6 edges)
+
+Class `mace.modules.models.ScaleShiftMACE`. Forward is one long method over a
+dict batch; blocks are three parallel `ModuleList`s indexed together.
+
+```
+node_embedding            LinearNodeEmbeddingBlock       out (3,128)      irreps 89x0e -> 128x0e
+radial_embedding          RadialEmbeddingBlock           out ((6,10), None)   tuple: (edge_feats, cutoff)
+spherical_harmonics       SphericalHarmonics             out (6,16)       1x0e+1x1o+1x2e+1x3o
+atomic_energies_fn        AtomicEnergiesBlock            out (3,1)
+interactions.{0,1}        RealAgnosticResidualInteractionBlock  out ((3,128,16), (3,128))  tuple: (message, sc)
+  .linear_up .conv_tp .conv_tp_weights .linear .skip_tp .reshape
+products.{0,1}            EquivariantProductBasisBlock   out (3,128)      node_feats, irreps 128x0e (L=0 model)
+  .symmetric_contractions .linear
+readouts.0                LinearReadoutBlock             out (3,1)        per-atom energy contribution, layer 0
+readouts.1                NonLinearReadoutBlock          out (3,1)        layer 1 (hidden 16x0e)
+scale_shift               ScaleShiftBlock                out (3,)
+```
+
+Facts that matter for hooking:
+
+- `interactions.i` output is a tuple; the message tensor is index 0 with
+  shape `[N, channels, (lmax+1)^2]` = `[3,128,16]`, i.e. **not** e3nn-flat
+  yet; `products.i` returns the flat node features. Hook `products.i` for
+  the residual-stream analogue.
+- `hidden_irreps` for the small model is `128x0e` (scalars only). Medium and
+  large are `128x0e+128x1o`; `hidden_irreps` is on each interaction block.
+- Per-layer `readouts` exist natively: a layer-resolved per-atom energy
+  decomposition for free. Observed values for H2O: layer 0 `[-0.186, -0.136, -0.136]`,
+  layer 1 `[0.157, 0.199, 0.199]` (eV, before scale/shift).
+- e3nn linears carry a `_compiled_main` TorchScript submodule: opaque to
+  hooks and to source inspection. Do not hook below block level.
+- dtype: the model is float64; building `AtomicData` without
+  `torch.set_default_dtype(torch.float64)` produces a float32 batch and a
+  TorchScript "both inputs should have same dtype" error deep inside the
+  tensor product. The calculator handles this for you; a raw `HookedModel`
+  must too.
+- Zero-ablating `products.0` output moved H2O energy from -14.0479 to -14.1266 eV
+  (float64), confirming interventions propagate.
+- Batch construction: `mace.data.config_from_atoms` ->
+  `AtomicData.from_config(cfg, z_table, cutoff=model.r_max)` -> geometric
+  DataLoader -> `.to_dict()`. `model(batch, compute_force=False)` returns a
+  dict with `energy`, `node_energy`, `forces`, ...
+
+## SevenNet (`sevenn`), model `7net-0`, Cu 2x2x2 conventional (32 atoms, 1344 edges)
+
+Class `sevenn.nn.sequential.AtomGraphSequential`: an `nn.Sequential` whose
+`forward` is literally `for module in self: data = module(data)` over an
+`AtomGraphData` dict. **Every module's output is the whole dict**, mutated
+in place; node features live under key `x`.
+
+```
+edge_embedding                 EdgeEmbedding          (.basis_function (1344,8), .spherical (1344,9) 1x0e+1x1e+1x2e)
+onehot_idx_to_onehot           OnehotEmbedding
+onehot_to_feature_x            IrrepsLinear           89x0e -> 128x0e            (.linear out (32,128))
+{t}_self_connection_intro      SelfConnectionLinearIntro   128x0e -> 224x0e+64x1e+32x2e
+{t}_self_interaction_1         IrrepsLinear
+{t}_convolution                IrrepsConvolution      (.convolution out (1344, 1152) edge-level TP; .weight_nn radial MLP)
+{t}_self_interaction_2         IrrepsLinear           -> 224x0e+64x1e+32x2e
+{t}_self_connection_outro      SelfConnectionOutro
+{t}_equivariant_gate           EquivariantGate        gate out (32,480) = 128x0e+64x1e+32x2e   <- node_feats after block t
+reduce_input_to_hidden         (.linear (32,64))
+reduce_hidden_to_energy        (.linear (32,1))
+rescale_atomic_energy, reduce_total_enegy, force_output
+```
+
+t = 0..4 (five interaction blocks). Block 4 collapses to `128x0e`.
+
+Facts that matter:
+
+- Hook points must select `data["x"]` from the dict output; write-back
+  interventions must assign `data["x"] = new` (in-place) rather than return
+  a value. Leaf `.linear` submodules return plain tensors if you prefer,
+  but the gate has no clean tensor-returning leaf for the block output.
+- Feature irreps `128x0e+64x1e+32x2e`; vector channels are labelled **even**
+  parity (`1e`), unlike MACE's `1o`. SevenNet's spherical encoding is
+  `1x1e -> 1x0e+1x1e+1x2e`: effectively SO(3), parity not tracked. A dipole
+  probe cannot use parity to select channels here; degree l only.
+- Layer names are `{t}_...` strings, not a `ModuleList`; the adapter should
+  discover `t` by regex over `named_children()`.
+- Module names `reduce_total_enegy` (sic) exist upstream; do not "fix" them.
+
+## Orb v3 (`orb-models`), `orb-v3-direct-20-omat`, Cu 2x2x2 (32 atoms, 640 edges)
+
+Class `orb_models.forcefield.models.direct_regressor.DirectForcefieldRegressor`
+wrapping `model: MoleculeGNS` plus a `heads` ModuleDict. Hooked to depth 3.
+
+```
+model.atom_emb           AtomEmbedding                 out (32,256)
+model.rbf_transform      BesselBasis                   out (640,8)
+model.angular_transform  SphericalHarmonics            out (640,16)   lmax 3, edge features only
+model._encoder           Encoder                       out ((32,256), (640,256))   (node, edge); ._node_fn / ._edge_fn are MLP + LayerNorm
+model.gnn_stacks.{0..4}  AttentionInteractionNetwork   out ((32,256), (640,256))   (node, edge)  <- node_feats = idx 0
+  ._send_attn            -> (640,1)   per-edge sigmoid gate on the sender side
+  ._receive_attn         -> (640,1)   per-edge sigmoid gate on the receiver side
+  ._edge_mlp             -> (640,256) updated edge features (the message)
+  ._node_mlp             -> (32,256)  node update from aggregated messages
+model._decoder           Decoder                       out (32,3)
+model                    -> {"node_features": (32,256), "edge_features": (640,256), "pred": (32,3)}
+heads.energy   mlp (1,1) + reference.linear (1,)      heads.forces (32,3)   heads.confidence (32,50)   heads.stress (1,6)
+pair_repulsion_fn        ZBLBasis                      -> {"energy","forces","stress"}
+```
+
+Facts that matter:
+
+- Plain `[N, 256]` node latents and `[E, 256]` edge latents, no irreps;
+  rotation handled by augmentation. Equivariance tests are
+  approximate-invariance tests.
+- Every block returns a `(node, edge)` tuple, same convention as the
+  encoder; `HookPoint.select = lambda o: o[0]` for node state, `o[1]` for
+  edge state. Edge state persists across blocks, so Orb has a genuine edge
+  residual stream that the e3nn models lack.
+- The two per-edge attention gates (`_send_attn`, `_receive_attn`, shape
+  `[E, 1]`) are the closest thing to attention patterns in the zoo, and a
+  natural hook point for "which neighbours does atom i listen to".
+- Heads read `node_features` directly: energy via a pooled MLP plus a
+  linear per-element reference, forces as a direct per-node MLP (direct
+  model), confidence as 50 bins per atom.
+- Loader returns `(model, atoms_adapter)`; the calculator needs both.
+  orb-models sets `torch.set_default_dtype(float32)` globally at import.
+- Environment: orb-models pins `dm-tree==0.1.8`, which has no cp313 wheel
+  and fails to build from source; run the introspection script with
+  `uv run --python 3.12` (the script's `requires-python` cap alone did not
+  stop uv from picking 3.13 on 2026-09-10).
+
+## UMA (`fairchem-core`), `uma-s-1p1`, Cu 2x2x2 (32 atoms, 2496 edges)
+
+Loaded with an HF token via `pretrained_mlip.get_predict_unit`. Wrapping:
+`MLIPPredictUnit.model` is an `AveragedModel` around a `HydraModel` with
+`backbone: eSCNMDMoeBackbone` and `output_heads: ModuleDict`. Backbone
+config for uma-s-1p1: `num_layers=4, lmax=2, mmax=2, sphere_channels=128,
+num_experts=32`.
+
+```
+backbone.sphere_embedding            Embedding          out (32,128)     initial scalar node embedding
+backbone.charge/spin/dataset_embedding, mix_csd         out (1,128)      per-system conditioning (csd)
+backbone.composition_embedding       Embedding          out (32,128)
+backbone.routing_mlp                 Sequential         out (1,32)       MoLE expert mixing coefficients, per system
+backbone.distance_expansion          GaussianSmearing   out (2496,64)
+backbone.envelope                    PolynomialEnvelope out (2496,)
+backbone.source/target_embedding     Embedding          out (2496,128)
+backbone.edge_degree_embedding       EdgeDegreeEmbedding out (32,9,128)  first node state in grid layout
+backbone.blocks.{0..3}               eSCNMD_Block       out (32,9,128)   <- node_feats after block i
+  .norm_1 .edge_wise (Edgewise) .norm_2 .atom_wise (SpectralAtomwise | GridAtomwise)
+backbone.norm                        EquivariantRMSNormArraySphericalHarmonicsV2  out (32,9,128)
+backbone                             -> {"node_embedding": (32,9,128), "batch": (32,)}
+output_heads.energyandforcehead.head MLP_EFS_Head       -> {"energy": (1,), "embeddings": (32,9,128), "forces": (32,3), "stress": (1,9)}
+output_heads.energyandforcehead      DatasetSpecificSingleHeadWrapper -> {"{task}_energy": ..., "{task}_embeddings": ..., ...} for every task
+```
+
+Facts that matter:
+
+- **Default inference settings silence hooks.** `get_predict_unit(name)`
+  uses the `"default"` preset: `compile=True, merge_mole=True`. With that,
+  only the pre-compile embeddings and the head fired; nothing inside the
+  backbone did. Passing
+  `InferenceSettings(compile=False, activation_checkpointing=False, merge_mole=False)`
+  (or the `"traineval"` preset) makes every block boundary hookable. Energy
+  agreed to 1e-6 eV between the two modes. The UMA adapter must construct
+  the predict unit itself with hooks-safe settings.
+- Node features `[N, (lmax+1)^2, C]` = `[32, 9, 128]`; the middle axis holds
+  SH coefficients, degree-l block at `l^2:(l+1)^2`. `blocks.i` output is
+  pre-`balance_channels`; `balance_channels` runs in the backbone forward
+  loop after each block and is not a module, so the post-balance state is
+  only observable as the next block's input (`register_forward_pre_hook`).
+- The final `norm` output equals the backbone's `node_embedding`; the head
+  also returns it as `embeddings`. This is the tensor Chorna et al. and
+  Platonic-representation work used.
+- MoLE: `routing_mlp` produces one 32-vector of expert coefficients per
+  system from the charge/spin/dataset embedding; effective linear weights in
+  every block are a system-dependent mixture. Activation hooks are
+  unaffected. Weight-level analysis must fix a system (or a task) first;
+  `merge_mole=True` materialises the merged weights but breaks hooks.
+- The head wrapper emits per-task outputs (`omat_energy`, `oc20_energy`,
+  `omol_energy`, ...) for all tasks in one forward; the calculator selects
+  by `task_name`. Useful for cross-task comparison from a single cache.
+- Custom Wigner-D rotation (`Jd.pt`, `WignerDataModule`) and optional
+  Triton kernels; no e3nn. Weights gated on HuggingFace (`facebook/UMA`).
+
+## Cross-cutting conclusions
+
+1. Three feature layouts: e3nn flat (MACE, SevenNet, NequIP), SO(3) grid
+   (UMA/eSEN/eqV2), plain (Orb, MatterSim). One `blocks()` view over all.
+2. Three output conventions from hooked modules: plain tensor (MACE
+   products), tuple (MACE interactions, Orb encoder), mutated dict
+   (SevenNet, fairchem). `HookPoint.select` plus an optional write-back
+   handles all three.
+3. Never hook below block boundaries: TorchScript, cuEquivariance,
+   OpenEquivariance, and Triton all rewrite block internals.
+4. Every family has its own graph construction and dtype expectations;
+   `from_atoms` must delegate to the family's own code.

matterlens-0.0.1/docs/case-studies.md

@@ -0,0 +1,99 @@
+# Case studies: how matterlens's abstractions map onto prior work
+
+Written 2026-09-10 as pitch context. Each section takes a published or in-house
+result, describes what the authors actually built to get at model internals,
+and shows which matterlens abstraction replaces each piece. The science in these
+papers (anchor selection, optimal-transport metrics, distillation to algebra,
+transfer experiments) is outside matterlens's scope; what matterlens removes is the
+bespoke extraction layer under it and the silent inconsistencies that layer
+introduces.
+
+## Platonic representation of foundation MLIPs (Li and Walsh, NMI 2026)
+
+Code: https://github.com/WMD-group/PlatonicRep, file `mlip_embedding_extractor.py`
+(1166 lines). What they built, per family, to obtain one vector per atom:
+
+| Family | Mechanism | Glue | Which tensor | Invariance |
+|---|---|---|---|---|
+| MACE | built-in `get_descriptors` | ~35 lines | all layers concatenated (256-d) by default; notebook uses layer 1 only (128-d) | selected by MACE's default `invariants_only=True` |
+| Orb v3 | reach into private `.model`, call `node_features` | ~57 lines | after all 5 blocks, before heads | not equivariant, nothing to select |
+| SevenNet | forward hook on the module before the readout, plus manual iteration over `_modules` bypassing `forward` to skip the force pass | ~95 lines | last gate output | inherited: SevenNet's last layer is built with lmax 0 |
+| NequIP | rebuild the `SequentialGraphNetwork` with inserted `SaveForOutput` modules | ~120 lines | node features key | not stated |
+
+Other properties: one forward pass per structure, no batching; three separate
+conda environments; no pinned versions; the Orb import path used no longer
+exists in current orb-models, so the extractor is already broken by drift.
+
+What changes with matterlens:
+
+- The four mechanisms collapse to one call per family:
+  `cache = load(family, name).run_with_cache(hm.from_atoms(atoms), names=["head.node_feats"])`
+  then `cache.scalars("head.node_feats")`. Roughly 15 lines total.
+- "Which tensor" becomes an explicit, comparable choice. The paper compares
+  MACE's all-layers concatenation against Orb's last layer; both happen to be
+  256 wide, which hides that they are different objects. With canonical names
+  the comparison is `head.node_feats` on every model, or `blocks.i.node_feats`
+  for a layer-resolved version they could not easily run.
+- "Scalars only" becomes a declared operation (`.scalars()`) rather than an
+  accident of each architecture's defaults.
+- The adapter test suite catches the two failure modes this code is exposed
+  to: import drift (fidelity test fails to load) and a hook that silently
+  never fires (UMA's compile default did exactly this on 2026-09-10).
+- The activation store emits the Zenodo-style artifact with per-atom
+  metadata for free.
+- New questions become cheap: at which block does cross-model alignment
+  emerge; do the degree-1 (vector) channels of MACE and SevenNet align, not
+  just the scalars; does patching an anchor direction from one model into
+  another change its energy.
+
+## Comparing latent features of uMLIPs (Chorna, Ceriotti et al., 2026)
+
+Paper: https://arxiv.org/abs/2512.05717. Materials Cloud record has per-model
+extraction scripts in separate directories with separate conda environments.
+UMA's mixture-of-experts routing was fixed to the OMat expert by hand;
+single-atom structures were excluded as incompatible with UMA. The MACE
+feature width is reported as 144, which is not any single hook point width in
+MACE-MP-0b3 (128 channels); which tensor it is cannot be recovered from the
+paper.
+
+With matterlens: the UMA adapter owns the task selection and the compile/MoLE
+handling once; `head.node_feats` is a named, documented tensor with a
+recorded irreps string, so "144-d MACE features" would instead read as a
+hook name and a layout that anyone can reproduce.
+
+## Skala exchange-correlation functional (Viswanath and Blaiszik, in-house draft)
+
+Draft: `docs/draft.tex`. This is the in-house existence proof for the whole
+program, and its methods section is the matterlens workflow executed by hand on a
+model that is not an MLIP:
+
+| Step in the draft | MatterLens abstraction |
+|---|---|
+| Load the traced graph and re-derive the model as equations, verifying manual re-execution to machine precision | the adapter fidelity test (wrapped output equals native output) and the introspection scripts that dump the module tree |
+| Reduce the nonlocal branch to per-atom moments organised as 16 radial shells times spherical harmonics up to degree 3 | a `HookPoint` on the moment tensor with an irreps/layout declaration; `cache.blocks()` to view it by degree |
+| Truncate the moments to their l=0 part, keeping everything else, and observe the sign flip on multiply-bonded N/O systems | `hooks({"nonlocal.moments": keep_degree(0)})`, an intervention returning a masked tensor |
+| The same surgery shell by shell to find bond and lone-pair detectors | the same hook with a mask over the radial-shell axis |
+| Sigma-bond count linearly decodable from internal moments at leave-one-molecule-out R^2 of 0.90 | `probe.linear` with held-out and random-init controls |
+| Distil the mechanism to algebra and transplant it onto r2SCAN | outside matterlens; this is the science |
+
+Two observations for the pitch. First, the draft itself notes the nonlocal
+branch rewrites exactly as the Atomic Cluster Expansion, which is the same
+mathematical object as a MACE interaction block. The block-boundary plus
+irreps abstraction therefore covers "grid-based functionals" in the program
+plan without a new design, which is why the plan's ambition to span model
+families is realistic. Second, Skala adds a fourth activation scope beyond
+node, edge, and graph: per grid point (the pointwise meta-GGA branch). That
+is a v1 extension of the cache's scope field, not a redesign.
+
+## The pitch in one paragraph
+
+Every group that has looked inside an MLIP so far wrote its own extraction
+layer: hundreds of lines per model family, private attributes, hand-selected
+tensors whose identity is not recoverable from the paper, and environments
+that break within months. None of that code is reused by the next group.
+MatterLens replaces it with named, tested hook points and an activation cache
+that knows what its tensors mean, so that probing, ablation, patching, and
+SAE training run identically across MACE, UMA, SevenNet, and Orb. The Skala
+result shows the workflow produces real physics when done by hand once;
+matterlens is what makes it repeatable for the d-band, dipole, and
+superconductivity targets that follow.

matterlens-0.0.1/docs/prior-art.md

@@ -0,0 +1,115 @@
+# Prior art (surveyed 2026-09-10)
+
+Condensed from two web surveys run on 2026-09-10. Claims marked "uncertain"
+were not verified beyond abstracts.
+
+## The gap is real
+
+No hooking or lens library exists for MLIPs or geometric GNNs. Searches for
+"TransformerLens for GNNs" and "MatterLens" return nothing. GNN
+"explainability" tooling exists but is the attribution/XAI tradition, not
+mechanistic interpretability:
+
+- PyG `torch_geometric.explain` (GNNExplainer, PGExplainer, CaptumExplainer):
+  node/edge masks. https://pytorch-geometric.readthedocs.io/en/latest/tutorial/explain.html
+- GraphXAI (Zitnik lab): benchmark datasets and metrics, Sci. Data 2023. https://github.com/mims-harvard/GraphXAI
+
+No SAE-on-MLIP paper found for MACE, UMA, Equiformer, SevenNet, Orb,
+MatterSim, or NequIP. No d-band or oxidation-state probe on MLIP latents
+found. Creed et al., "Six Open Questions in MLIP Foundation Models"
+(2026-06) names interpretability as an open pillar. https://arxiv.org/abs/2606.07327
+
+## Closest existing results on MLIP internals
+
+- **Platonic representation of foundation MLIPs**, Li and Walsh, Nat. Mach.
+  Intell. 2026-05: MACE-MP-0 S/M/L, MACE-omat, SevenNet-omat, Orb-v3;
+  embeddings from the invariant component before energy readout; element
+  embeddings cluster by periodic group; equivariant models give compact
+  clusters. Code: https://github.com/WMD-group/PlatonicRep .
+  https://www.nature.com/articles/s42256-026-01235-7
+- **Comparing latent features of uMLIPs**, Chorna, Ceriotti, Chong et al.,
+  Adv. Intell. Syst. 2026: MACE-MP-0b3, PET-MAD, DPA-3.1, UMA-S-1p1; feature
+  reconstruction error shows models encode chemistry very differently
+  (cross-model GFRE ~0.66); fine-tuning preserves pre-training features.
+  https://arxiv.org/abs/2512.05717
+- **Goodfire and Radical AI on MatterGen** (2026-04): linear probes on
+  activations predict band gap; probe-guided accept/reject in the denoising
+  loop raises the hit rate. Probes, not SAEs; no code or preprint found.
+  https://www.goodfire.com/research/self-correcting-search
+- **Sparsity-promoting fine-tuning of MACE**, Cho et al., ICLR 2026: 0.5 to
+  3 percent of weights suffice; sparsity patterns show enhanced d-orbital
+  contributions for transition metals. https://arxiv.org/abs/2606.18691
+- **PROBE**, Mehdi, Cho, Isayev, 2026-05: attention classifier on frozen
+  per-atom embeddings gives reliability scores and per-atom importance.
+  https://arxiv.org/abs/2605.00640
+- **Equivariant SAEs**, Erdogan and Lucic, 2025 (rev. 2026-08): SAEs on
+  symmetric data are unidentifiable; equivariant SAEs probe better despite
+  worse reconstruction. Whether any target model is an MLIP is uncertain.
+  https://arxiv.org/abs/2511.09432 , code https://github.com/ege-erdogan/equivariant-sae
+- **GNN-LRP on force fields**: JCTC 2025 (SchNet/PaiNN)
+  https://pubs.acs.org/doi/10.1021/acs.jctc.4c01424 ; Nat. Commun. 2025-11
+  https://www.nature.com/articles/s41467-025-65863-0
+- Vault science briefing on what to look for (dipoles, Born charges, LES):
+  `rootstock-scratch/notes/Mechanistic Interpretability of MLIPs.md`.
+
+## Lens libraries for LLMs: what to copy and what to avoid
+
+**TransformerLens** https://github.com/TransformerLensOrg/TransformerLens
+
+- `HookPoint` identity modules at every interesting tensor; names are module
+  paths (`blocks.0.attn.hook_z`); `run_with_cache` returns an
+  `ActivationCache` with residual-stream helpers assuming `[batch, pos, d_model]`.
+- 1.x reimplemented every architecture and converted HF weights into it.
+  The v3 migration guide calls this "a latent source of bugs"; v3's
+  `TransformerBridge` keeps the native HF module and maps it onto canonical
+  component names with legacy aliases. `HookedTransformer` is deprecated.
+  https://transformerlensorg.github.io/TransformerLens/content/migrating_to_v3.html
+- Weight processing (`fold_ln`, `center_writing_weights`) are
+  computation-preserving rewrites specific to LN + residual streams; no
+  obvious equivariant analogue.
+
+**nnsight / NDIF** https://github.com/ndif-team/nnsight
+
+- Wraps any `nn.Module`; `.output`/`.input` on any submodule inside
+  `with model.trace(...)`; `.source` AST-instruments a module's forward to
+  expose intra-module tensors; `remote=True` runs on NDIF. Inherits HF
+  naming fragmentation; **nnterp** (https://arxiv.org/abs/2511.14465) adds a
+  renaming layer plus init-time validation tests. NDIF hosts LLMs, diffusion,
+  VLMs; nothing scientific found.
+
+**pyvene** https://github.com/stanfordnlp/pyvene
+
+- Per-model-type tables mapping canonical component names to
+  `(module path template, hook type, optional post-split)`. The cheapest
+  pattern for canonical names over heterogeneous modules; adopted here.
+
+**baukit** https://github.com/davidbau/baukit : `Trace`/`TraceDict` by dotted
+module name, retain/edit output. Minimal; the mechanism under everything.
+
+**torchlens** https://github.com/johnmarktaylor91/torchlens : logs every
+tensor op in a forward pass, validated, tested on GNNs. Useful for
+discovering where a needed tensor lives when it is not at a module boundary.
+
+**torch.fx**: fails on input-dependent control flow and un-wrapped
+functions; MLIPs hit both (neighbour lists, scatter). Not used.
+
+## SAE tooling interface
+
+- SAELens needs a hook name, `d_in`, and an activation source; inference works
+  on any torch model. https://github.com/jbloomAus/SAELens
+- `dictionary_learning` uses an `ActivationBuffer` over nnsight submodules;
+  many trainer variants. https://github.com/saprmarks/dictionary_learning
+- EleutherAI `sparsify` takes plain module names with globs and trains
+  many SAEs at once, no activation store. https://github.com/EleutherAI/sparsify
+- SAEDashboard hard-requires `HookedTransformer`. https://github.com/jbloomAus/SAEDashboard
+
+Implication: the activation store only has to yield `[n_units, d_in]` rows
+with metadata; "token" becomes "atom in structure".
+
+## Hosting and community
+
+- HuggingFace hosts MACE (`mace-foundations`), UMA (`facebook/UMA`, gated),
+  OMAT24 models, OrbMol; the Hugging Science community (2026) curates
+  science models. No HF-authored MLIP interp work found.
+- NeurIPS 2026 "Interpretability for Discovery" workshop (Dec, Atlanta).
+  https://interpretability4discovery.github.io/

matterlens-0.0.1/docs/releasing.md

@@ -0,0 +1,24 @@
+# Releasing
+
+Publishing is automated by `.github/workflows/publish.yml` using uv and PyPI
+trusted publishing, so no API token is stored anywhere.
+
+One-time setup (2026-09-10, not yet done):
+
+1. On pypi.org, log in as the account that will own the project, go to
+   "Your projects" -> "Publishing" -> "Add a new pending publisher", and enter:
+   PyPI project name `matterlens`, owner `Garden-AI`, repository `matterlens`,
+   workflow name `publish.yml`, environment name `pypi`.
+2. In the GitHub repo, Settings -> Environments -> create `pypi`. Optionally
+   restrict it to tag deployments and require a reviewer.
+
+Every release afterwards:
+
+```bash
+uv version --bump patch   # or set the version in pyproject.toml and __init__.py
+git commit -am "Release vX.Y.Z" && git tag vX.Y.Z && git push --tags
+```
+
+The workflow builds the sdist and wheel, smoke-imports the wheel in a clean
+environment, and publishes. Keep `__version__` in `src/matterlens/__init__.py`
+in step with `pyproject.toml` until a dynamic-version setup is adopted.