Spin/Charge System Conditioning by JonathanSchmidt1 · Pull Request #1080 · metatensor/metatrain

JonathanSchmidt1 · 2026-03-22T12:57:13Z

System conditioning for PET (charge & spin)

Adds per-system charge and spin multiplicity conditioning to the PET architecture,
allowing a single model to be trained and evaluated across multiple charge and spin
states. The feature is activated through architecture.model.system_conditioning: true.
System Conditioning is separated into its own SystemConditioningEmbedding module (pet/modules/conditioning.py). The resulting embedding is added to
PET's node features via a zero-initialised gated projection, so the model starts as
the unconditioned baseline and learns to use charge/spin information only as needed.

Charge and spin are supplied as mtt::charge (integer, elementary charges) and
mtt::spin (integer, spin multiplicity 2S+1) in the extra_data section of the
dataset config or via atoms.info in ASE (requires merging of a PR into metatrain).

Changes

New files

src/metatrain/pet/modules/conditioning.py — SystemConditioningEmbedding module
and get_system_conditioning_transform (re-exported from utils/system_data.py)
src/metatrain/utils/system_data.py — generic get_system_data_transform callable
for attaching per-system scalar TensorMaps to System objects in a CollateFn
src/metatrain/pet/tests/test_conditioning.py — test suite for the feature

PET model (pet/model.py)

Adds system_conditioning hyper; if enabled, builds SystemConditioningEmbedding
and injects the embedding into node features during both initial featurisation and
residual updates
Declares mtt::charge / mtt::spin in requested_inputs() so the exported model
communicates its requirements to downstream tools (ASE calculator, eval pipeline)

Training (pet/trainer.py)

Reads model.system_conditioning.required_data_keys and registers
get_system_conditioning_transform as a CollateFn callable so charge/spin are
attached to System objects during training

Checkpoint upgrade (pet/checkpoints.py)

v11 → v12 upgrade detects the presence of system_conditioning.* weights in the
state dict to auto-enable the hyper for checkpoints trained with conditioning
(avoids silent neutral-singlet fallback when loading old muon-branch checkpoints)

Eval (cli/eval.py + utils/system_data.py)

mtt eval now reads extra_data from the dataset config and routes any keys
present in the model's requested_inputs() through get_system_data_transform,
so charge/spin reach the model during evaluation
The transform raises early if a TensorMap is per-atom rather than per-system,
preventing silent index errors on mixed datasets

Hypers (pet/documentation.py, share/base_hypers.py)

New system_conditioning block in ModelHypers:
system_conditioning: bool, max_charge: int = 10, max_spin: int = 10

Training config example

architecture:
  model:
    system_conditioning: true
    max_charge: 5   # embeds charges in [-5, +5]
    max_spin: 5     # embeds multiplicities in [1, 5]

training_set:
  - path: dataset.mtt
    extra_data:
      mtt::charge:
        field: charge
      mtt::spin:
        field: spin

Contributor (creator of pull-request) checklist

Tests updated (for new features and bugfixes)?
Documentation updated (for new features)?
Issue referenced (for PRs that solve an issue)?

Maintainer/Reviewer checklist

CHANGELOG updated with public API or any other important changes?
GPU tests passed (maintainer comment: "cscs-ci run")?

📚 Documentation preview 📚: https://metatrain--1080.org.readthedocs.build/en/1080/

Adds an `extra_data_options` parameter to `MemmapDataset` that loads per-system scalar arrays from `.bin` files alongside the training targets. Each key (e.g. `mtt::charge`) maps to a `TensorMap` in the sample namedtuple and is forwarded to the `extra` argument of `CollateFn` callables. Also adds `get_extra_data_info()` to expose `TargetInfo` metadata for the extra_data keys, mirroring `get_target_info()`. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

- Drop stale system.add_data() calls for charge/spin from MemmapDataset.__getitem__ (data now flows through extra_data_options + get_system_conditioning_transform) - Remove charge/spin fields from SystemsHypers in base_hypers.py (config now lives under extra_data: {mtt::charge: {key: ...}}) Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

…rade Old muon2 checkpoints already contain system_conditioning.* weights. Setting system_conditioning=False was dropping them silently. Now the upgrade checks for the presence of those weights and enables the hyper automatically, so converted checkpoints use the embedding they were trained with. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

or per-atom extra data is never passed to the system transform - system_data: raise if a TensorMap is per-atom (samples != ["system"]) rather than silently misindexing into the systems list - model: validate charge/spin are integer-valued before .long() conversion; log.debug when a system falls back to default 0/1 - conditioning: document zero-init gate design intent

…add test that confirms eval is working with spin/charge

…nstead of memmap

JonathanSchmidt1 · 2026-03-22T13:19:51Z

What's the policy now on how to update the classifier/llpr checkpoints?

JonathanSchmidt1 · 2026-03-23T15:15:08Z

I updated the checkpoints so the only test that is still failing will require the metatomic part to be merged (we could also switch that to metatomic but I I think it's good to have it here so we will notice directly if we mess up the interaction with metatomic)

pfebrer

The failing test, isn't it a test that should go to metatomic?

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

sofiia-chorna

very interesting 😊 thanks a lot for implementation!

i left some comments. would be nice to fix the use of "charge"/"spin" and "mtt::charge"/"mtt::spin" (those metatomic PRs should be merged first I suppose? metatensor/metatomic#183 metatensor/metatomic#189) to make sure the tests verify not the default values but actual inputs

sofiia-chorna · 2026-04-08T06:28:08Z

+
+    :param checkpoint: The checkpoint to update.
+    """
+    import logging


would be nice to have all imports at the top

sofiia-chorna · 2026-04-08T06:31:12Z



+@with_config(ConfigDict(extra="forbid", strict=True))
+class SystemDataKeyHypers(TypedDict):


seems to be not used anywhere

sofiia-chorna · 2026-04-08T06:52:47Z

+        the range ``[1, max_spin]``.
+    """
+
+    required_data_keys: List[str] = ["mtt::charge", "mtt::spin"]


why mtt:: here? in all other places you use simply "charge" and "spin". that's why CI job failes i guess

sofiia-chorna · 2026-04-08T07:33:06Z

+            objects describing each per-system scalar array.
+        """
+        extra_data_info_dict: Dict[str, TargetInfo] = {}
+        if not self.extra_data_config:


it seems we don't initialize self.extra_data_config anywhere in the DiskDataset

sofiia-chorna · 2026-04-08T07:43:45Z

+        checkpoint["model_data"]["model_hypers"]["max_charge"] = 10
+    if "max_spin" not in checkpoint["model_data"]["model_hypers"]:
+        checkpoint["model_data"]["model_hypers"]["max_spin"] = 10
+    # Rename edge_linear -> edge_embedder (muon2 branch used edge_linear)


is muon2 branch merged? or why we need to rename it?

sofiia-chorna · 2026-04-08T08:01:17Z

+                    self.system_conditioning(
+                        inputs["charge"],
+                        inputs["spin"],
+                        inputs["system_indices"],
+                    )


we currently call self.system_conditioning(...) with the same inputs at every gnn layer => output is the same => we can move it outside of the loop over gnn layers and then simply do:

output_node_embeddings = output_node_embeddings + cond_embedding

same for _residual_featurization_impl

sofiia-chorna · 2026-04-08T08:14:58Z

+        # the `extra` argument of CollateFn callables
+        extra_data_dict = {}
+        for key, arr in self.extra_data_arrays.items():
+            is_per_atom = arr.shape[0] == self.na[-1]


looks fragile, it is possible to have a case when the total number of atoms equals the number of systems... we should check ["per_atom"] from self.extra_data_config here, seems to be available

sofiia-chorna · 2026-04-08T08:22:44Z

+            # Extract per-system charge and spin for conditioning
+            if self.system_conditioning is not None:


i would move extraction outside to a separate function to remain the forward function readable!

sofiia-chorna · 2026-04-08T08:35:12Z

+    def validate(self, charge: torch.Tensor, spin: torch.Tensor) -> None:
+        """Check that charge and spin values are within the supported range.
+
+        Call this outside of ``torch.compile`` regions to get descriptive errors.


validate is called inside scripted forward, have you tried to call the exported (mtt export) model? i am just wondering to make sure the validate call survives export

sofiia-chorna · 2026-04-08T09:39:04Z

        collate_fn_train = CollateFn(
            target_keys=list(train_targets.keys()),
            callables=[
                rotational_augmenter.apply_random_augmentations,
                get_system_with_neighbor_lists_transform(requested_neighbor_lists),
+                *conditioning_callables,
                get_remove_additive_transform(additive_models, train_targets),
                get_remove_scale_transform(scaler),
            ],
            batch_atom_bounds=self.hypers["batch_atom_bounds"],
        )
        collate_fn_val = CollateFn(
            target_keys=list(train_targets.keys()),
            callables=[  # no augmentation for validation
                get_system_with_neighbor_lists_transform(requested_neighbor_lists),
+                *conditioning_callables,
                get_remove_additive_transform(additive_models, train_targets),
                get_remove_scale_transform(scaler),
            ],
            batch_atom_bounds=self.hypers["batch_atom_bounds"],
        )


sorry it gives me urge for a small refactor 😁

Suggested change

base_callables = [

get_system_with_neighbor_lists_transform(requested_neighbor_lists),

*conditioning_callables,

get_remove_additive_transform(additive_models, train_targets),

get_remove_scale_transform(scaler),

]

collate_fn_train = CollateFn(

target_keys=list(train_targets.keys()),

callables=[rotational_augmenter.apply_random_augmentations, *base_callables],

batch_atom_bounds=self.hypers["batch_atom_bounds"],

)

collate_fn_val = CollateFn(

target_keys=list(train_targets.keys()),

callables=base_callables,

batch_atom_bounds=self.hypers["batch_atom_bounds"],

)

…oning-rebased # Conflicts: # pyproject.toml # src/metatrain/pet/trainer.py # tests/utils/data/test_dataset.py

…ntity Metatomic's standard per-system input name is `spin_multiplicity` (`spin` is only the short ASE info key on the calculator side). To make exported PET models pluggable into MetatomicCalculator without an extra prefix, rename throughout: `required_data_keys`, `system.get_data` reads, the hyperparameter `max_spin` → `max_spin_multiplicity`, the internal embedding attribute `spin_embedding` → `spin_multiplicity_embedding`, validate/forward parameter names, the test fixtures, and the MemmapDataset extra_data example. Also restore the rename of `required_data_keys` from `mtt::charge`/`mtt::spin` (which had been accidentally reverted) to the unprefixed standard names. Bump the model checkpoint version to 13 and add `model_update_v12_v13` that renames the `max_spin` hyperparameter and the `system_conditioning.spin_embedding.*` state-dict keys so existing v12 checkpoints continue to load. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

JonathanSchmidt1 · 2026-04-28T12:01:14Z

Let's take a look at this only once charge/spin are fully merged into metatomic

- checkpoints.py: hoist `import logging` to module top. - share/base_hypers.py: drop unused `SystemDataKeyHypers` TypedDict. - utils/data/dataset.py: * initialize `self.extra_data_config = {}` in `DiskDataset.__init__` so `get_extra_data_info()` no longer AttributeErrors when called on a disk dataset that did not set it externally. * inside `MemmapDataset.__getitem__`, replace the fragile `arr.shape[0] == self.na[-1]` per-atom heuristic with the explicit `self.extra_data_config[key]["per_atom"]` flag, which is unambiguous when n_atoms == n_systems. - pet/model.py: hoist `system_conditioning(...)` out of both featurization GNN loops. Inputs (charge, spin_multiplicity, system_indices) are loop-invariant, so compute the embedding once and add it inside the loop. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Round 2 of metatensor#1080 review (sofiia-chorna): - pet/model.py: extract per-system charge/spin_multiplicity reading from `forward()` into a module-level helper `_extract_charge_spin_multiplicity` so the core forward stays readable. - pet/trainer.py: factor out `base_callables` shared between the train and validation `CollateFn`s. Train just prepends `rotational_augmenter` to the base list. - pet/tests/test_conditioning.py: add `test_export_with_conditioning_preserves_validate` to regression-test that the in-forward `validate(...)` call survives TorchScript compilation in `model.export()`. Drives the saved model end-to-end via `MetatomicCalculator` and checks both the valid path and the out-of-range error path. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

JonathanSchmidt1 and others added 26 commits March 19, 2026 13:36

add get_extra_data_info and tests for MemmapDataset extra_data

0d81173

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

respect per atom/subtargets

31eca6a

addition tests, mtt::charge name and heterogeneous atom counts

a592bf6

add charge and spin to dataset

8ebc205

add spin/charge hyperparameters to systemhypers

a0c5b86

small fixes

1562557

add new checkpoint

26a7144

new checkpoint file

8e6bdd3

revert toml change

402167e

change hyperparameters for system embedding

4dae87b

remove duplicate extra_data_arrays registration

9a23028

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

add requested inputs

fa0f1e2

reverse change in naming introduced in muon branch

a7144ba

add extra data to eval, separate from conditioning

7586e0b

add extra data to eval, separate from conditioning

47293b8

fix spin/charge extraction

9c9f4f4

fix absolute to relative indices

8388d00

fix torch script string issues

d6a9a56

warning on missing conditioning incompatible with torch script

4c3fe36

requested inputs not working in eval, revert back to all extra data, …

2b84edd

…add test that confirms eval is working with spin/charge

add check for existing system data, necessary when reading from xyz i…

34e3c3b

…nstead of memmap

lint

354ab91

JonathanSchmidt1 requested a review from frostedoyster March 22, 2026 12:57

JonathanSchmidt1 requested a review from abmazitov as a code owner March 22, 2026 12:57

missing return type in system-conditioning-transform

9e8fd4a

update llpr/classification checkpoints

a30431a

JonathanSchmidt1 requested a review from SanggyuChong as a code owner March 23, 2026 15:04

JonathanSchmidt1 added 2 commits March 23, 2026 15:28

correct function parameters for linter

82c67c8

correct function parameters for linter

d184944

pfebrer reviewed Mar 24, 2026

View reviewed changes

Comment thread src/metatrain/cli/eval.py Outdated

JonathanSchmidt1 and others added 6 commits March 26, 2026 10:15

replace extra_data_keys with model.requested_inputs() in eval

278c862

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

remove system-conditioning-transform

be39760

remove conditioning keys

bd57efd

remove test that depends on metatomic calculator, will go into metatomic

346d634

mtt::spin/charge to spin/charge

b067e5f

add standard output charge/spin

b9e3ed0

sofiia-chorna reviewed Apr 8, 2026

View reviewed changes

JonathanSchmidt1 and others added 2 commits April 28, 2026 11:37

Merge remote-tracking branch 'upstream/main' into only-system-conditi…

e40d4af

…oning-rebased # Conflicts: # pyproject.toml # src/metatrain/pet/trainer.py # tests/utils/data/test_dataset.py

JonathanSchmidt1 and others added 2 commits April 28, 2026 12:08



		@with_config(ConfigDict(extra="forbid", strict=True))
		class SystemDataKeyHypers(TypedDict):

		# Extract per-system charge and spin for conditioning
		if self.system_conditioning is not None:

Conversation

JonathanSchmidt1 commented Mar 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

System conditioning for PET (charge & spin)

Changes

Training config example

Contributor (creator of pull-request) checklist

Maintainer/Reviewer checklist

Uh oh!

JonathanSchmidt1 commented Mar 22, 2026

Uh oh!

JonathanSchmidt1 commented Mar 23, 2026

Uh oh!

pfebrer left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

sofiia-chorna left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

JonathanSchmidt1 commented Apr 28, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

JonathanSchmidt1 commented Mar 22, 2026 •

edited

Loading

sofiia-chorna left a comment •

edited

Loading