Add uma model - #22
Open
declanwhchan wants to merge 38 commits into
Open
Conversation
… in pyproject.toml and README.md
…LI together; add new arguments for parsing for UMA; add new cli pytests for MACE/UMA
…a into optional dependencies in pyproject.toml; update cli in README.md for uma support
…emove imports in __init__.py to avoid mace/uma conflicts
…tent argument parsing; separate imports for mace/uma to avoid conflicts
… in relaxation.py; combined MACE/UMA logging; fix "require" to "required" typo in make_attack.py for MACE/UMA --model CLI
…y to customize MACE vs UMA attacks
…termine model; change file imports for attacks based on MACE or UMA folder
…MACE or UMA virtual environment
…oth MACE and UMA attacks together in pytests
…storage of fmax, task, charge, spin using .json instead of atoms.info
… path requirements and set default values for charge and spin
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
…or mace-mh; added mace-mh support
… MACE-MH and UMA modelsp
…calculator.py to separate MACE and UMA
There was a problem hiding this comment.
Pull request overview
This PR extends mlff_attack to support Meta’s UMA models alongside MACE/MACE-MH, unifying the relaxation and attack workflows under shared CLI entry points and a shared “ASE-compatible calculator” attack implementation.
Changes:
- Unified single-structure relaxation CLI into
calc-single(replacingmace-calc-single/uma-calc-single) and expandedmake-attackto accept UMA task/charge/spin and MACE-MH head selection. - Refactored calculator setup into
src/mlff_attack/calculator.pyand updatedsetup_calculator()to route to MACE vs UMA setup. - Replaced MACE-specific FGSM/PGD implementations with unified ASE-calculator-based
FGSM_ASE/PGD_ASE, plus extensive test and documentation updates.
Reviewed changes
Copilot reviewed 33 out of 36 changed files in this pull request and generated 9 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/test_visualization.py | Removes unused header imports while keeping visualization tests intact. |
| tests/test_uma_calc_single.py | Adds UMA-focused unit tests guarded by optional dependency skipping. |
| tests/test_pgd.py | Updates PGD tests to PGD_ASE and adds UMA + MACE-MH coverage via optional deps. |
| tests/test_mace_calc_single.py | Updates MACE calc tests for new calculator routing and adds MACE-MH head test. |
| tests/test_fgsm.py | Updates FGSM tests to FGSM_ASE and adds UMA + MACE-MH coverage via optional deps. |
| tests/test_examples.py | Updates example tests to FGSM_ASE and guards with optional MACE dependency. |
| tests/test_cli.py | Updates CLI tests for calc_single.py + new UMA/head argument validation. |
| tests/test_attacks.py | Updates attack tests to unified attack classes and adds module-level optional-dep skip logic. |
| src/mlff_attack/visualization.py | Improves logging levels and minor formatting fixes. |
| src/mlff_attack/relaxation.py | Routes calculator setup through new calculator helpers and adds new CLI-facing args. |
| src/mlff_attack/metrics.py | Promotes an error log to logger.error. |
| src/mlff_attack/grad_based/pgd.py | Replaces PGD_MACE with unified PGD_ASE supporting MACE autograd + ASE-force fallback. |
| src/mlff_attack/grad_based/fgsm.py | Replaces FGSM_MACE with unified FGSM_ASE supporting MACE autograd + ASE-force fallback. |
| src/mlff_attack/examples/example_fgsm_attack.py | Updates examples to new FGSM_ASE class. |
| src/mlff_attack/cli/visualize_traj.py | Updates CLI description and minor formatting. |
| src/mlff_attack/cli/uma_calc_single.py | Removes UMA-only calc CLI (superseded by calc_single.py). |
| src/mlff_attack/cli/make_attack.py | Adds UMA + MACE-MH arguments and backend inference/validation. |
| src/mlff_attack/cli/mace_calc_single.py | Removes MACE-only calc CLI (superseded by calc_single.py). |
| src/mlff_attack/cli/calc_single.py | Adds unified single-structure relaxation CLI supporting MACE/MACE-MH/UMA. |
| src/mlff_attack/cli/init.py | Updates CLI exports to reference calc_single. |
| src/mlff_attack/calculator.py | New helper module encapsulating MACE and UMA calculator setup. |
| src/mlff_attack/attacks.py | Updates make_attack() to use unified attack classes and new calculator routing. |
| src/mlff_attack/init.py | Updates package exports to include new calculator module. |
| README.md | Documents dual-backend support, uv-based install flow, and the MACE/UMA e3nn conflict. |
| pyproject.toml | Makes MACE/UMA optional extras, adds uv conflicts, and reorganizes notebook/dev deps. |
| notebooks/_1_general_workflow.ipynb | Updates notebook references from FGSM_MACE to FGSM_ASE. |
| grad_based/README_FGSM.md | Updates docs to the new FGSM_ASE class name and examples. |
| docs/source/mlff_attack.relaxation.rst | Updates description to include UMA. |
| docs/source/mlff_attack.cli.rst | Updates CLI docs to reference calc_single. |
| docs/source/installation.rst | Adds detailed install guidance (uv + two-venv strategy for MACE vs UMA). |
| docs/source/index.rst | Updates overview and installation quick-start for MACE/UMA split environments. |
| docs/source/examples.rst | Updates examples to FGSM_ASE and calc-single. |
| docs/source/conf.py | Switches Sphinx theme to Furo and updates theme options. |
| .gitignore | Broadens venv ignores and adds model/weight extensions. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+3
to
+7
| mace = pytest.importorskip( | ||
| "mace", | ||
| reason="test_attacks.py requires mace-torch dependencies (switch to virtual environment that supports MACE)", | ||
| ) | ||
|
|
Comment on lines
+111
to
+115
| atoms = setup_calculator( | ||
| atoms, | ||
| model_path, | ||
| device, | ||
| calculator=calculator_kind, |
Comment on lines
+92
to
+96
| if calculator not in {None, "mace", "uma"}: | ||
| logger.error("[ERROR] Invalid calculator '%s'. Use 'mace' or 'uma'.", calculator) | ||
| atoms = None | ||
|
|
||
| elif calculator == "uma": |
Comment on lines
+103
to
+113
| model_name = Path(args.model).name.lower() | ||
| is_mace_mh = model_name.startswith("mace-mh") | ||
|
|
||
| if model_name.startswith("uma"): | ||
| calculator = "uma" | ||
| elif model_name.startswith("mace"): | ||
| calculator = "mace" | ||
| else: | ||
| raise SystemExit( | ||
| "--model must start with 'uma' for UMA or 'mace' for MACE" | ||
| ) |
Comment on lines
+154
to
+163
| model_name = Path(args.model).name.lower() | ||
| is_mace_mh = model_name.startswith("mace-mh") | ||
| if model_name.startswith("uma"): | ||
| calculator = "uma" | ||
| elif model_name.startswith("mace"): | ||
| calculator = "mace" | ||
| else: | ||
| raise SystemExit( | ||
| "--model basename must start with 'uma' for UMA or 'mace' for MACE" | ||
| ) |
Comment on lines
259
to
261
| except (ValueError, NotImplementedError, RuntimeError) as exc: | ||
| logger.info("[ERROR] Failed to generate attack: %s", exc) | ||
| logger.error("[error] Failed to generate attack. Run calc-single first and use the same model to generate attack: %s", exc) | ||
| return 1 |
Comment on lines
+23
to
+27
| except ImportError: | ||
| logger.info( | ||
| "[ERROR] MACE requires mace-torch. Install it with: pip install -e \".[mace]\"" | ||
| ) | ||
| return None |
Comment on lines
+37
to
+42
| logger.info( | ||
| "[ERROR] Invalid MACE-MH head '%s'. Choose one of: %s", | ||
| mace_head, | ||
| ", ".join(model_path.heads), | ||
| ) | ||
| return None |
Comment on lines
+96
to
+103
| valid_uma_tasks = {"oc20", "oc22", "oc25", "omat", "omol", "odac", "omc"} | ||
| if uma_task not in valid_uma_tasks: | ||
| logger.info( | ||
| "[ERROR] Invalid UMA uma_task '%s'. Choose one of: %s", | ||
| uma_task, | ||
| ", ".join(sorted(valid_uma_tasks)), | ||
| ) | ||
| return None |
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
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.
Added UMA support to work with MACE for relaxation and attacks:
Added MACE-MH support: