[DRAFT] AddingMagneticMACE #1244
Conversation
|
Hey @CheukHinHoJerry thank you very much for that. Yes the model should be in the extensions.py, the new blocks themselves can be in blocks.py. For the random rotations, could please create a new file in data folder called augmentations.py where you put it? thank you |
|
No problem - I will make those changes and let you know. Thank you for the quick comment. |
|
@ilyes319 I turned on precommit checks but there are too many lines in |
MagneticMACE MagneticMACE
|
@CheukHinHoJerry Yep you can disable that for that file! |
# Conflicts: # mace/calculators/mace.py # mace/data/atomic_data.py # mace/modules/__init__.py # mace/modules/extensions.py # mace/tools/train.py
BREAKING: rename the magnetic calculator result key from mace_magmom to MACE_magmoms. Old .model files are not expected to reload through the new MagneticMACECalculator surface. Wire compute_edge_forces and compute_atomic_stresses through MagneticScaleShiftMACE.forward, mirroring the ScaleShiftMACE output contract (returns edge_forces, atomic_virials, atomic_stresses). Remove the stale MagneticSCFMACE data["mace_magmom"] dead write — nothing downstream consumes that key. Add tests/test_magmace.py smoke coverage for the MagneticSCFMACE forward path. test_run_eval_magnetic_mace is currently @pytest.mark.skip pending an eval_configs.py update to plumb magmom into the data dict post-merge; tracked as follow-up. test_run_train_magnetic_mace is left in the file but not exercised by this commit (training is heavy and out of scope for prep smoke). Pre-commit-config bump --max-module-lines=1500 -> 2000 already covered by the merge commit. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
…CE keys Restore magmom support in mace/cli/eval_configs.py so magnetic MACE models can be evaluated end-to-end: - Add --magmom_key (default REF_magmom) and --return_magforces. - Thread KeySpecification with the magmom_key through config_from_atoms so AtomicData.from_config populates data["magmom"]. - Plumb compute_magforces through get_model_output and write MACE_magforces to atoms.arrays when --return_magforces is set. Canonicalize the magmom keys across the package so input and output do not share a buffer: - INPUT reference data: REF_magmom (alongside REF_energy/REF_forces) - OUTPUT model prediction: MACE_magmoms (alongside MACE_energy/MACE_forces) MagneticMACECalculator default magmom_key changes MACE_magmoms -> REF_magmom. The test_run_train_magnetic_mace fixture no longer needs to copy REF_magmom into MACE_magmoms before calling get_potential_energy(). Unskip test_run_eval_magnetic_mace and add --return_magforces=True to its Namespace. Both test_run_eval_magnetic_mace and test_run_magnetic_scf now pass in clean-pr-test. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
…hift In _build_model the MagneticScaleShiftMACE branch hardcoded atomic_inter_shift=[0.0]*len(heads), so --mean was silently ignored for magnetic models. Replace with _determine_atomic_inter_shift(args.mean, heads), matching the regular MACE/ScaleShiftMACE branches. The helper already returns [0.0]*len(heads) when args.mean is None, so the previous behaviour is preserved for callers that do not pass --mean. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
…ce_model Add MagneticScaleShiftMACE to the whitelist in mace/tools/scripts_utils.extract_config_mace_model. Round-trip the magnetic hyperparameters needed to rebuild an identical instance: - m_max (per-element saturation, registered buffer) - max_m_ell (max ell of the magmom spherical-harmonics basis) - num_mag_radial_basis (Chebyshev count in mag_radial_embedding) - use_magmom_one_body - num_mag_radial_basis_one_body (only when use_magmom_one_body=True, read from onebody_magmombasis_coeffs.shape[1]) Adds tests/test_magmace.py::test_extract_config_magnetic_round_trip covering the keys above. Required upstream of any foundation-FT work that wants to reconstruct a magnetic foundation from its config. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Add multihead_tools.inherit_magnetic_hyperparameters_from_foundation: copies m_max, max_m_ell, num_mag_radial_basis, and num_mag_radial_basis_one_body from the foundation checkpoint onto args, so users do not have to pass them by hand when fine-tuning a saved magnetic foundation. Ported from 482c250. run_train calls it right after args.r_max = model_foundation.r_max.item(). No-op for non-magnetic foundations (extract_config_mace_model does not emit m_max etc. for them). tests/test_magmace.py covers the helper directly with a synthetic foundation_config fixture. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
# Conflicts: # mace/data/__init__.py # mace/modules/__init__.py
Previously --m_max required a positional list of N floats matching the
order of atomic_numbers — fragile across datasets with different element
coverage. Mirror the --E0s pattern by also accepting a dict literal mapping
atomic number to m_max value:
--m_max '{26: 1.8, 28: 1.2}'
Only listed elements are required; unspecified ones default to 1.0.
scripts_utils.resolve_m_max handles the conversion to a per-element list
ordered by atomic_numbers. It accepts:
* None -> None (caller decides default)
* list[float] -> fast path, validated length
* single dict-literal string token -> ordered by atomic_numbers
* single float-literal string token -> broadcast across all elements
* legacy list of float-literal strings -> parsed and length-validated
argparse now uses type=str, nargs="+" so both the dict form and the
legacy space-separated float form keep working. The resolver is called
in model_script_utils._build_model's MagneticScaleShiftMACE branch right
before the constructor call, using model_config["atomic_numbers"] as the
ordering.
tests/test_magmace.py adds 7 focused unit tests covering all input
shapes and the error cases (wrong length, unknown element).
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
The previous resolve_m_max raised ValueError on any dict key that wasn't
in atomic_numbers. That breaks the natural use case of passing a generic
over-spec dict (e.g. all Z=1..94) on a dataset whose z_table only covers
a subset — which is what we want for FT-ready foundation training.
Change: silently ignore extra dict keys (log INFO listing them), only
use the dict entries that match the current z_table; missing entries
fall back to `default`.
Also normalize atomic_numbers to python int before set-comparison so
np.int64 z_tables (the common case) work the same as python int.
Tests:
* test_resolve_m_max_unknown_element_raises -> replaced with
test_resolve_m_max_extra_dict_keys_ignored (asserts the new behavior)
* new test_resolve_m_max_numpy_atomic_numbers covers np.int64 zs
Reported from a real MATPES_V2_BADER training run where the dict covered
Z=1..94 but the train data lacked Po (84), At (85), Rn (86), Fr (87), Ra (88).
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Upstream main added `use_edge_irreps_first` to ScaleShiftMACE/MACE (commit 6dd9134) and the foundation_filtered _build_model branches pass it through to whatever __init__ they call. MagneticMACE.__init__ on this branch didn't accept it, so MATPES_V2_BADER training failed at model construction with: TypeError: MagneticMACE.__init__() got an unexpected keyword argument 'use_edge_irreps_first' Mirror upstream commit a13910d ("add unused use_edge_irreps_first arg to models that don't inherit from MACE baseclass"): add the kwarg with a False default and the pylint:unused-argument silence. The magnetic stack doesn't have the edge-irreps-first behavior so this is intentionally a silent ignore. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
…on gate The post-merge run_train.py asserts model.__class__.__name__ is in a whitelist before calling run_e3nn_to_cueq / run_e3nn_to_oeq. Add MagneticScaleShiftMACE to both whitelists. The e3nn->cueq conversion is a generic walk-and-replace over tensor products and works on the magnetic class as well — the source branch (magnetic-fix-nonSOC) has been training MAGNETIC models with --enable_cueq=True for months; it just had the assertion commented out there. Whitelisting is cleaner than commenting out. Same change for the oeq path for consistency. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
The radial magmom transform `1 - 2*(|m|/m_max)**2` assumed `|m|/m_max ∈ [0, 1]` so the result stays in `[-1, 1]` (where the Chebyshev basis is numerically stable). It was unclamped, so any sample with |bader_magmom| > m_max for its element pushed the transform below -1, where the Chebyshev recurrence is unstable and produced single-step gradient blowups (loss → 2.85e+176 at step 348 of a MATPES bader run on 2026-05-24, then NaN forever after). Clamping the ratio to [0, 1] is the targeted fix; the MATPES bader training started 2026-05-24 10:00 has been running cleanly since. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Reconcile magnetic-pr with upstream develop. Conflicts resolved additively: - mace/cli/run_train.py: keep both foundation-model gates (magnetic hyperparameter inheritance and finetune_dipoles_polarizabilities check); include both MagneticScaleShiftMACE and AtomicDielectricMACE in the CUEQ conversion allow-list. - mace/tools/model_script_utils.py: import both load_foundations_mdp and resolve_m_max. - mace/tools/scripts_utils.py: extract_config_mace_model accepts both MagneticScaleShiftMACE and AtomicDielectricMACE; the MagneticScaleShiftMACE-specific config keys are preserved alongside the new atomic_energies/scale_shift/MLP_irreps handling. Also add a pylint import-error disable on the guarded torch_geometric import in mace/data/augmentation.py so the lint step passes cleanly.
The magforces branches in take_step() and evaluate() indexed output_args
with a hard ['magforces'] key, which raised KeyError for any caller that
didn't populate the key (e.g. tests/unit/test_lora.py::
test_lora_evaluate_preserves_frozen_state, which passes only forces/
virials/stress). Use .get('magforces', False) so magforces stays opt-in
and legacy callers keep working.
The magnetic-pr magforces path required args.return_magforces on the Namespace, but callers that construct partial Namespaces (e.g. the LES extension tests: test_run_eval_with_bec, test_run_eval_no_bec) hit AttributeError. Normalize once at the top of run() so magforces stays opt-in and legacy callers keep working.
…cMACECalculator MagneticMACECalculator reads model.magmom_mace.<attr> in six places, but that attribute only exists when the loaded model is a MagneticSCFMACE wrapper. Training via mace_run saves a raw MagneticScaleShiftMACE, so loading a trained checkpoint blew up with: AttributeError: 'MagneticScaleShiftMACE' object has no attribute 'magmom_mace' Replace the six accesses with getattr(model, 'magmom_mace', model).<attr> so raw and SCF-wrapped models both work. Rejected the alternative of wrapping raw models in MagneticSCFMACE(use_scf=False) at load time — that would push every model into the equilibrated_magmom branch of calculate() and break committee inference (assert len(self.models) == 1).
|
I can't reproduce the pylint error locally... any suggested fix? |
aacostadiaz
left a comment
There was a problem hiding this comment.
Thanks for this, really nice to see magnetic MACE taking shape!!
Don't worry about the red lint check, that one isn't on you. It's a pre-existing import cycle in mace.data that only surfaced because pylint isn't pinned and your new file shifted the module ordering. I'll fix it separately, so just ignore it here.
One heads-up: if you take the --max-module-lines suggestion below, add the per-file # pylint: disable=too-many-lines to extensions.py at the same time, otherwise lint just goes red on C0302 instead.
|
Thanks for the quick reply and confirming the pylint issue! I will take a look at those comment and fix it soon. |
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
- utils.py: add `or compute_edge_forces` to two magforces branches of get_outputs at lines 294 & 315 (was only for the hessian path). Keeps the autograd graph alive when magforces + edge_forces are requested together. - extensions.py: swap `if self.use_magmom_one_body:` for `if hasattr(self, "one_body_cheb_basis_with_const"):` at the four forward- body sites (two per magnetic model class). Keeps the __init__ construction guard as-is. Makes TorchScript compile / model reload robust when a checkpoint saved with use_magmom_one_body=False is loaded. - augmentation.py: create_random_rotation_loader now distinguishes DistributedSampler (pass through) from ordinary samplers (fall back to shuffle=True). Fixes the single-process path where data_aug_magmom=True previously reused a RandomSampler and the loader stopped re-shuffling between epochs. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
| import numpy as np # noqa: F401 (already imported at top of file) | ||
|
|
There was a problem hiding this comment.
| import numpy as np # noqa: F401 (already imported at top of file) |
| if model.radial_embedding.bessel_fn.__class__.__name__ == "BesselBasis": | ||
| model.radial_embedding.bessel_fn.bessel_weights = torch.nn.Parameter( | ||
| model_foundations.radial_embedding.bessel_fn.bessel_weights.clone() | ||
| ) |
There was a problem hiding this comment.
does a similar thing need to be done for GaussianBasis and gaussian_weights?
| if model.interactions[i].__class__.__name__ in [ | ||
| "MagneticRealAgnosticSpinOrbitCoupledDensityInteractionBlock", | ||
| ]: | ||
| # import pdb; pdb.set_trace() |
There was a problem hiding this comment.
did you mean to keep the commented code here?
| model.products[i].conv_tp.weight = torch.nn.Parameter( | ||
| model_foundations.products[i].conv_tp.weight.clone() | ||
| ) | ||
| # model.products[i].conv_tp_weights.weight = torch.nn.Parameter( |
There was a problem hiding this comment.
should this also be retained?
There was a problem hiding this comment.
While the LES and Polar require separate requirements.txt files, you could probably but the magnetic extras in the setup.cfg files given these dependencies exist on PyPI unlike the LES and Polar extras which point towards github repositories e.g.
In setup.cfg:
[options.extras_require]
...
magnetic =
sphericart-torch==1.0.9
torch-geometric
* build: drop Python 3.9, make 3.12 the CI default Python 3.9 reached EOL in October 2025. The concrete cost in this repo was not the EOL date: torch stopped publishing cp39 wheels after 2.8.0, so the `unit (py3.9)` job silently resolved to torch <=2.8 while every other job ran 2.9+. That leg was validating a different stack than the one we ship. - python_requires >=3.10; drop the 3.9 classifier and the README floor. - unit / workflows / workflows-full matrices: 3.10-3.13. - PR and release frontiers move to 3.10 + 3.13. - setup-mace default, the lint job and the release build move to 3.12. The per-job pins that spread across 3.10/3.11/3.12 are left alone on purpose: collapsing them onto 3.12 would shrink the set of versions CI exercises outside the matrices. - pylint: pin py-version to the floor (3.10) so lint tracks what we support rather than whichever interpreter the job runs on. Verified message-set identical between py-version 3.10 and 3.12. - mypy: python_version 3.8 -> 3.10, which was stale under the old floor too. - ruff: py38 -> py39, deliberately still below the floor. FA102 is the only enabled rule and it degrades to a no-op at py310; py39 keeps `X | Y` flagged because TorchScript cannot parse it. * refactor: drop the python-3.9 annotation workaround in polar_density_cube f358b99 added `from __future__ import annotations` to polar_density_cube because its PEP-604/PEP-585 annotations crashed the entry point at import time on python 3.9. With the floor at 3.10 that workaround is dead: `str | Path`, `tuple[...]` and `list[Path]` all evaluate natively. Verified by importing all 12 console-script modules on 3.10 and 3.12. Left alone on purpose: - mace/cli/plot_train.py keeps its future import. It is NOT a 3.9 workaround: matplotlib/pandas are import-guarded and `pd` is None when they are missing, so `data: pd.DataFrame` would be evaluated at def time and raise AttributeError — precisely the clean-install crash the guard exists to prevent. - tests/helpers.py keeps its `list | None` future import; tests are excluded from lint/format anyway. Removing the import made ruff's FA102 demand it back, because ruff.toml was pinned to py39 globally. That was too broad: the constraint is torch, not python. TorchScript rejects `X | Y` on torch < 2.2 (and the future import does not rescue it), so the guard only belongs on code carrying @compile_mode("script"). Global target-version now states the real floor (py310) and per-file-target-version keeps py39 on the scripted surface. Verified with a positive control (planting `int | None` in modules/blocks.py still trips FA102) and a negative one (the same line in mace/cli does not). * build: pin pylint The pre-commit pylint hook uses `language: system`, so its `rev: pylint-2.5.2` is decorative and CI lints with whatever pylint the `dev` extra happens to install. A new pylint release can therefore turn any PR red with no code change — pylint 4 is what surfaced the long-standing `mace.data` import cycle that currently blocks #1244. This belongs in this PR rather than a standalone one: pylint 4.0.6 requires Python >= 3.10.0, so a single-line pin is only satisfiable once the floor moves to 3.10. On the old 3.9 floor it would need a python_version marker, and the `unit (py3.9)` / `workflows (py3.9)` jobs — which install `extras: dev` — would fail dependency resolution outright. Verified `.[dev]` resolves to pylint==4.0.6 on 3.10, 3.11, 3.12 and 3.13.
|
@CheukHinHoJerry thanks for addressing the comments. Could you merge the latest changes from |
- Factor radial-weight transfer into `_copy_radial_weights` helper and call it from all three loaders (default / magnetic / mdp). The helper uses `.data.copy_(...)` in-place so a target registered as `register_buffer` (BesselBasis/GaussianBasis default `trainable=False`) is no longer silently promoted to a Parameter. Adds GaussianBasis parallel to BesselBasis handling. - Drop dead `.reshape/.flatten/.clone/*num_species_foundations` comment blocks around magmom_skip_tp, skip_tp, and conv_tp_weights inside `load_foundations_elements_magnetic`. - test_magmace.py test_resolve_m_max_numpy_atomic_numbers: drop redundant `import numpy as np` (already imported at module top). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Address @AntObi's review comment on PR ACEsuit#1244: sphericart-torch and torch-geometric are on PyPI so magnetic can live in `[options.extras_require]` rather than a separate requirements file (LES/Polar still need requirements.txt because they point at github URLs). - setup.cfg: new `magnetic` extra listing sphericart-torch==1.0.9 and torch-geometric. - ci-extensions.yaml: magnetic job now uses `extras: dev, magnetic` instead of `pip-packages: -r requirements/magnetic.txt`; drop the now-stale `'requirements/magnetic.txt'` path-filter entry (setup.cfg is already in the shared filter, so extras edits still trigger CI). - requirements/magnetic.txt: deleted. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…re bug fix Cleanup pass driven by 3 iterations of codex-style review of the PR (P1/P2 findings only, style-only nits skipped). - utils.py: assert magmoms -> ValueError (assert stripped under python -O); fix "Magnetic momenet must be inputed" -> "moment must be provided". - extensions.py: rename magmom_lenghts -> magmom_lengths (5 sites); drop misleading pretraining-hook comment and stale "specicies dependent transform" comment (species-dep scaling is implemented in forward). - calculators/mace.py: 3x print -> logging.info to match sibling MACECalculator; delete commented dev leftovers; use DefaultKeys.MAGMOM.value for magmom_key default. - eval_configs.py: use DefaultKeys.MAGMOM.value instead of hard-coded "REF_magmom" (single source of truth). - default_keys.py: add MAGMOM/MAGFORCES entries so keydict() picks them up alongside energy/forces/etc. - arg_parser.py: fix "agumentation"/"manetic" typos in --data_aug_magmom help; switch --magmom_key / --magforces_key to DefaultKeys. - augmentation.py: sample quaternion in fp32 when magmom.dtype is bf16/fp16 then cast R back, so mixed-precision runs get proper SO(3) sampling instead of ~256-discrete-angle bf16 rotations; drop the misleading "hemisphere rotation" docstring (implementation is uniform SO(3)). - train.py: collapse take_step.closure duplicated model() branch into a single call with conditional compute_magforces kwarg; apply the same pattern to take_step_lbfgs.closure -- the LBFGS closure previously hardcoded the no-magforces branch, so magforces were silently dropped during LBFGS steps even when --loss requested them. - tests/extensions/magnetic/test_magmace.py: drop 3 debug print(...) in test bodies. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
|
Thanks for the extra round of review. There are still some improvements to be done on how pre-train models are loaded but those can come later by a follow up PR I think. |
aacostadiaz
left a comment
There was a problem hiding this comment.
Thanks for addressing the comments. Two small things, then this looks good to merge from my side.
… magmom/magforces - scale MagneticMACECalculator.node_energy by energy_units_to_eV - default magmom/magforces to zeros in atomic_data (match forces/charges) - gate MagFs metric with filter_nonzero_weight so unlabeled configs (magforces_weight=0) do not contaminate rel_mae/rel_rmse Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This PR implements
MagneticMACEas abstract class with other spin-informed variant of magnetic MACE. Major changes areThis PR is still WIP because I am not sure how one would like to place the code and I am current at the stage of adding some tests , happy to make changes.
Questions in my mind for @ilyes319 atm:
extension.pytoo? Or should they be placed where they are supposed to be outsideextension.py?