Issue creating LAMMPS model from EnergyDipoleMACE model

[lemessuriernb]

I would like to test my inital MACE model, but I keep getting an error when using the create_lammps_model.py to convert an EnergyDipoleMACE model. Is it possible to use this type of MACE model with LAMMPS?

Error:
Traceback (most recent call last):
File "/home/lemessuriernb/miniconda3/envs/mace/bin/mace_create_lammps_model", line 8, in
sys.exit(main())
^^^^^^
File "/home/lemessuriernb/miniconda3/envs/mace/lib/python3.12/site-packages/mace/cli/create_lammps_model.py", line 16, in main
lammps_model_compiled = jit.compile(lammps_model)
^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/lemessuriernb/miniconda3/envs/mace/lib/python3.12/site-packages/e3nn/util/jit.py", line 99, in compile
compile(
File "/home/lemessuriernb/miniconda3/envs/mace/lib/python3.12/site-packages/e3nn/util/jit.py", line 111, in compile
mod = torch.jit.script(mod, **script_options)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/lemessuriernb/miniconda3/envs/mace/lib/python3.12/site-packages/torch/jit/_script.py", line 1432, in script
return _script_impl(
^^^^^^^^^^^^^
File "/home/lemessuriernb/miniconda3/envs/mace/lib/python3.12/site-packages/torch/jit/_script.py", line 1146, in _script_impl
return torch.jit._recursive.create_script_module(
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/lemessuriernb/miniconda3/envs/mace/lib/python3.12/site-packages/torch/jit/_recursive.py", line 559, in create_script_module
return create_script_module_impl(nn_module, concrete_type, stubs_fn)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/lemessuriernb/miniconda3/envs/mace/lib/python3.12/site-packages/torch/jit/_recursive.py", line 636, in create_script_module_impl
create_methods_and_properties_from_stubs(
File "/home/lemessuriernb/miniconda3/envs/mace/lib/python3.12/site-packages/torch/jit/_recursive.py", line 468, in create_methods_and_properties_from_stubs
concrete_type._create_methods_and_properties(
File "/home/lemessuriernb/miniconda3/envs/mace/lib/python3.12/site-packages/torch/jit/_recursive.py", line 1004, in try_compile_fn
return torch.jit.script(fn, _rcb=rcb)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/lemessuriernb/miniconda3/envs/mace/lib/python3.12/site-packages/torch/jit/_script.py", line 1432, in script
return _script_impl(
^^^^^^^^^^^^^
File "/home/lemessuriernb/miniconda3/envs/mace/lib/python3.12/site-packages/torch/jit/_script.py", line 1204, in _script_impl
fn = torch._C._jit_script_compile(
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError:
python value of type 'float' cannot be used as a value. Perhaps it is a closed over global variable? If so, please consider passing it in as an argument or use a local varible instead.:
File "/home/lemessuriernb/miniconda3/envs/mace/lib/python3.12/site-packages/mace/modules/utils.py", line 411
num_graphs: int,
) -> torch.Tensor:
mu = positions * charges.unsqueeze(-1) / (1e-11 / c / e) # [N_atoms,3]
~ <--- HERE
return scatter_sum(
src=mu, index=batch.unsqueeze(-1), dim=0, dim_size=num_graphs
'compute_fixed_charge_dipole' is being compiled since it was called from 'EnergyDipolesMACE.forward'
File "/home/lemessuriernb/miniconda3/envs/mace/lib/python3.12/site-packages/mace/modules/models.py", line 1035
dim_size=num_graphs,
) # [n_graphs,3]
baseline = compute_fixed_charge_dipole(
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
charges=data["charges"],
~~~~~~~~~~~~~~~~~~~~~~~~
positions=data["positions"],
~~~~~~~~~~~~~~~~~~~~~~~~~~~~
batch=data["batch"],
~~~~~~~~~~~~~~~~~~~~
num_graphs=num_graphs,
~~~~~~~~~~~~~~~~~~~~~ <--- HERE
) # [n_graphs,3]
total_dipole = total_dipole + baseline

[ilyes319]

Hey,

It is not currently possible to use the EnergyDipole model in lammps. Why ASE is not sufficient for you?

[lemessuriernb]

Hi,
I see. Is this a feature that may be added in the future?
I will see if ASE has the functionality I need, but my initial interest in MACE was because it interfaced with LAMMPS.
Which MACE models can be used with LAMMPS?
Is EnergyDipole only usable with ASE? Or will it work with OpenMM at all?
Thank you.

[ilyes319]

The MACE and ScaleShiftMACE models are the only one with a LAMMPS interface. If you have scripts that run in LAMMPS I recommend you train a normal MACE force field, do your trajectory and then evaluate the DipoleMACE on the trajectory to get all the dipoles. It will be identical.
It is not in our priorities at the moment.