Deployment of the MACE Model Without Using ASE

[Hsuchein]

First of all, I would like to express my gratitude for your efforts in developing the excellent MACE architecture.

Recently, I have been considering whether it is possible to package MACE into a .pt file using Torch's JIT technology and then deploy it using C++. This could potentially improve the inference speed of MACE and seamlessly integrate it with my C++ code, which is exactly what I am aiming for.

However, after going through the tutorial, it seems that MACE can only function as a calculator that is invoked by the atoms object in ASE for computation. My project does not require ASE, and the presence of ASE creates issues when packaging the model.

Ideally, I would like to achieve something like outputs = model.forward({species, coordinates}), where MACE acts as a standalone model. I would only need to input the atomic types (species) and atomic coordinates (coordinates) to compute the energy and potential.

Below is the code I have tried, but it resulted in an error. Could you please advise if MACE can be used independently of ASE, saved as a .pt model, and deployed with the Libtorch library to handle inputs and output energy and force?

Thank you for your attention to this matter.

``
`import torch
from mace.calculators import mace_off
from ase import Atoms

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

species = ['C', 'C']
coordinates = [
[0.0, 0.0, 0.0],
[1.0, 0.0, 0.0]
]

class Wrapper(torch.nn.Module):
def init(self):
super().init()
self.calc = mace_off(model="medium", device='cuda')

def forward(self, species, coordinates):
    
    atoms = Atoms(symbols=species, positions=coordinates)
    atoms.set_calculator(self.calc)
    energy = atoms.get_potential_energy()
    forces = atoms.get_forces()
    
    return energy, forces

wrapped_model = Wrapper()

energy, forces = wrapped_model(species, coordinates)

print(f"Energy: {energy}")
print(f"Forces: {forces}")

traced_model = torch.jit.script(wrapped_model)

traced_model.save("/home/wayne/data/mace/mace_medium.pt")`

[CheukHinHoJerry]

ASE calculator is not a must as long as you provide a dictionary that is correct for calling model.forward. Take a look at
https://github.com/ACEsuit/mace/blob/main/mace/data/atomic_data.py and how it is used in testing scripts etc.

If you want to use it in C++, maybe you can look at how it is used in lammps.

[bernstei]

Isn't the neighbor list calculated in python, which would therefore make it impossible to run in pure C++ without some significant supporting code?

[CheukHinHoJerry]

Yes. I would imagine you will have to write an interface for it though not something easy. (please correct me if I am wrong.)

[ilyes319]

Hey @CheukHinHoJerry,

Yes it is totally possible to run MACE in C++ using torchscript. This exactly how we made the lammps interface, see here: https://github.com/ACEsuit/lammps/blob/mace/src/ML-MACE/pair_mace.cpp#L308.
As @CheukHinHoJerry you will just need to provide the right input dictionary. To construct the NL you should be able to use the C++ matscipy NN.

[Hsuchein]

I guess you're right. I find a code in eval_configs.py, which use output = model(batch.to_dict(), compute_stress=args.compute_stress) to obtain the force and energy. Is that means what I should do is to provide the model with right dictionary, does it make sense?

Furthermore, I am not quite sure about nblist you guys mentioned, should I provide the model with nblist?

[CheukHinHoJerry]

This is the neighborlist that mace uses which is implemented in C++

https://github.com/libAtoms/matscipy

The mace model needs three things form matscipy neighborlist: edge_index, shifts and unit_shifts

https://github.com/ACEsuit/mace/blob/main/mace/data/neighborhood.py.

Calling this equivalently in C++, and then follow how mace serialize the input should work the same as what python gives. Although I would imagine this requires some work.

[wcwitt]

As @CheukHinHoJerry and @ilyes319 have said, it is possible to do what you want with a medium amount of effort, particularly if you are an experienced C++ programmer. The exact procedure required will depend on the details of your code (e.g., whether you need a neighbor list, like that available from matscipy). A few things to consider:

• We do exactly this when calling MACE from LAMMPS. The link @ilyes319 provided is a good place to start https://github.com/ACEsuit/lammps/blob/mace/src/ML-MACE/pair_mace.cpp#L308 because it shows how the torch dictionary inputs are created in C++. However, note that it uses a custom LAMMPS wrapper of the normal MACE model, which slightly extends the functionality. You likely only need the normal MACE object.
• When JIT-compiling the MACE model, you will need the e3nn JIT extension, as we do here: https://github.com/ACEsuit/mace/blob/main/mace/cli/create_lammps_model.py.
• Your C++ program will need to link against libtorch, which can be challenging the first time. For this, these docs are a good starting point: https://pytorch.org/tutorials/advanced/cpp_export.html.

[Hsuchein]

Thank you so much for your supportive help!

[CheukHinHoJerry]

Just to add - I recently found https://github.com/Luthaf/vesin to be quite easy to use in C++ that they provide a single-file and you can include and link to your own package. Matscipy binds to python interface a lot so would be difficult to call in C++. I don’t have much experience with that yet but my feeling is that using vesin would be easier.

[Hsuchein]

Hi guys, I found another question.

It seems that I can obtain the force by setting compute_force=True

out = self.model( data, training=False, compute_force=False, compute_virials=False, compute_stress=False, compute_displacement=False, )

Therefore, what is the difference betweent compute_force=True and using torch.autograd ?

[ilyes319]

it is using autograd internally.