Adding a MACE-MDP foundation model calculator and supporting cueq=True for AtomicDielectricMACE#1439
Merged
Merged
Conversation
Mace mu alpha
Cueq support for AtomicDielectricMACE
transpose_mul_ir converts CuEq linear output ir_mul→mul_ir before the block, so hardcode layout="mul_ir" instead of get_layout(). Fixes cueq=True/False inconsistency for models with non-scalar MLP_irreps. Update cueq matrix test to use MLP_irreps with 1o+2e components and add rotation equivariance checks for dipole and polarizability.
…ut blocks GatedEquivariantBlock natively supports both "mul_ir" and "ir_mul" layouts, so the transpose sandwich (ir_mul→mul_ir before gate, mul_ir→ir_mul after) is unnecessary. Use get_layout(cueq_config) directly instead.
Adding the final version of MACE-MDP as a foundation model calculator and AtomicDielectricMACE with cueq=True working
New --finetune_dipoles_polarizabilities flag enables fine-tuning a pretrained MACE-MDP checkpoint on dipoles and polarizabilities: - Validates that --model=AtomicDielectricMACE and --foundation_model are set; forces --loss=dipole_polar - Loads the MDP checkpoint via the existing foundation_model path and copies all weights (incl. readouts) with load_foundations() - After configure_model(), freezes all parameters except model.readouts - Reuses existing DipolePolarLoss, data pipeline, and training loop without modification Also adds tests/test_mdp_finetune.py verifying that non-readout params are frozen and readout params update after a short fine-tuning run.
…order irreps - Guard --E0s=foundation/estimated in MDP fine-tuning path (fall back to average) - Replace load_foundations() call with new load_foundations_mdp() that handles higher-order irreps (0e+1o+2e) in skip_tp and transfers all angular momentum channels in symmetric_contractions for all product layers - Guard scale_shift access in load_foundations_elements against MDP models
ilyes319
added a commit
that referenced
this pull request
May 3, 2026
Adding Documentation for MACE-MDP for PR #1439
Contributor
|
Hey @Nilsgoe, the deletion of import torch._dynamo as dynamo in mace/calculators/mace_torchsim.py might unsafe — that import has the side effect of initializing Dynamo before torch.compile runs. Could you revert that one-line removal here? Everything else in the PR looks good. |
Contributor
Author
|
Hey @ilyes319, I added the |
Introducing first ways for finetuning MACE-MDP, still some features missing and more testing needed
Contributor
Author
|
With the newest changes also fine-tuning of MACE-MDP would be possible, it is still naive only ... |
Collaborator
|
LGTM! Thanks for the PR |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This PR introduces [MACE-MDP](https://github.com/Nilsgoe/MACE-MDP) as a possible calculator for directly computing dipole moments and polarizabilities in organic systems (molecules and condensed phases).
Additionally, it supports
cueq=Truethanks to the newGatedEquivariantBlockadded in the development branch; to ensure correctness, I have included an additional test.