Problems with MACE and UMA examples

[cwalker7]

Hello - I ran into some issues when running the examples - could you share which versions of MACE and Fairchem they were tested on?

I am using the following:

ase                       3.27.0
torch                     2.9.1+rocm6.4 
mace-torch                0.3.15
fairchem-core             2.13.0

UMA Fairchem example

The functions init_edge_rot_mat and rotation_to_wigner have been renamed to eulers_to_wigner and init_edge_rot_euler_angles at some point.

File /autofs/nccs-svm1_home1/cuw/software/DistMLIP/DistMLIP/DistMLIP/implementations/uma/escn_md.py:23
     21 from DistMLIP.implementations.uma.compute import generate_graph
     22 from fairchem.core.models.base import HeadInterface
---> 23 from fairchem.core.models.uma.common.rotation import (
     24     init_edge_rot_mat,
     25     rotation_to_wigner,
     26 )
     27 from fairchem.core.models.uma.common.rotation_cuda_graph import RotMatWignerCudaGraph
     28 from fairchem.core.models.uma.common.so3 import CoefficientMapping, SO3_Grid

ImportError: cannot import name 'init_edge_rot_mat' from 'fairchem.core.models.uma.common.rotation' (/lustre/orion/mat226/scratch/cuw/lammps_nnp/lib/python3.12/site-packages/fairchem/core/models/uma/common/rotation.py)

MACE example

For MACE, I attempted to run the example on 1 node with 8 AMD GPUs, and am getting the following error:

cuequivariance or cuequivariance_torch is not available. Cuequivariance acceleration will be disabled.
There are 54000 atoms.
Using Materials Project MACE for MACECalculator with /ccs/home/cuw/.cache/mace/20231210mace128L0_energy_epoch249model
Using float32 for MACECalculator, which is faster but less accurate. Recommended for MD. Use float64 for geometry optimization.
Traceback (most recent call last):
  File "/autofs/nccs-svm1_home1/cuw/software/DistMLIP/DistMLIP/examples/test_mace.py", line 33, in <module>
    dist_calc.calculate(atoms)
  File "/lustre/orion/mat226/scratch/cuw/lammps_mace/lib/python3.12/site-packages/DistMLIP/implementations/mace/mace.py", line 85, in calculate
    ret_tensors = self._create_result_tensors(
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
TypeError: MACECalculator._create_result_tensors() missing 1 required positional argument: 'out'

Thanks! This will be very useful for my work and much simpler than LAMMPS to set up.

[AegisIK]

The versions we used are associated with specific GitHub commits and can be found in https://github.com/AegisIK/DistMLIP/blob/main/pyproject.toml

For installation, I would recommend following the readme here: https://github.com/AegisIK/DistMLIP/blob/main/README.md.

If there are some additional features in the later versions of MACE and Fairchem that people would like to see added, we can certainly consider adding them. Let me know if you have any additional questions!

[cwalker7]

Ok, thank you. I got the MACE example to run with that specific commit of mace-torch, though I get OOM errors for the 54k atom system, but that's a separate issue. The UMA is not a priority for me.

I would like to use the new MACE-POLAR model to simulate aqueous ions, which needs to import mace.modules.extensions from a recent version. I can use a single GPU implementation for now, but any support for this would be appreciated!