Print out long-range and short-range part during prediction with MACELES

[FadillaRizka]

Dear developers,

I am interested to use MACE with long-range corrections (MACE-LES).
I would like to request a guidance on how to print out the short-range and long-range part of the total energy, as well as the atomic charge during the predictions.

As reference, here in the input file I use to train the MACE-LES potential and the script to do total energy and force predictions.

TRAINING_INPUT_FILE

model: "MACELES" 
num_interactions: 2 
num_channels: 128
max_L: 1
correlation: 3 
r_max: 6.0
max_ell: 3 
name: "mace01" 
model_dir: "MACE_models" 
log_dir: "MACE_models" 
checkpoints_dir: "MACE_models" 
results_dir: "MACE_models" 
train_file: "../database/h2o-all.xyz" 
valid_fraction: 0.20 
test_file: "../database/h2o-all.xyz" 
E0s: '{1:-1.122582090, 8:-2.39899122}'
energy_key: "energy" 
forces_key: "forces" 
device: cuda 
batch_size: 5 
max_num_epochs: 6000 
distributed: True
enable_cueq: True
swa: True
seed: 123
restart_latest: True

PREDICTION_SCRIPT

from ase.io import read, write
import numpy as np
from ase import Atoms
from ase.io.trajectory import Trajectory
from ase.calculators.singlepoint import SinglePointCalculator
from mace.calculators import MACECalculator

calc = MACECalculator(model_path='../MACE_models/mace01.model', device='cuda')

traj = read('dft.xyz', index=':')
for atoms in traj:
    atoms.calc = calc
    atoms.get_potential_energy()
    atoms.write('mace.xyz', append=True)

Thank you for your help.

[ilyes319]

Hey, that is currently not possible from the ASE calculator atm. We will add that in a future update. If you need it now, you can try to hack it by directly evaluating the model and looking at the "les_energy"

mace/mace/modules/extensions.py

Line 269 in 42cd495

"les_energy": les_energy,