Training model with Energy, Force and Virial

[dominicvarghese]

Hi everyone,

I have a DFT dataset with energy, force and virials (not stress) of 1705 configurations. Is it possible to train a model using the virials as well in addition to energy and force?

I didnt find any tag similar to compute_stress tag for the virials. Is there an option to set compute_virials=True?

To train the model, I was using the following tags:

  --compute_forces=True \
  --compute_stress=True \
  --energy_weight=1.0 \
  --forces_weight=100.0 \
  --virials_weight=1.0 \
  --loss="weighted" \
  --energy_key="energy" \
  --forces_key="force" \
  --virials_key="virial" \
  --error_table=PerAtomRMSEstressvirials 

But in the run, it was not showing the RMSE virial suggesting this method was not using the virial data?

2026-05-27 15:52:12.875 INFO: Epoch 39: head: Default, loss=0.00091506, RMSE_E_per_atom=    0.99 meV, RMSE_F=   18.52 meV / A, RMSE_stress=None meV / A^3
2026-05-27 15:52:12.936 INFO: Training complete
2026-05-27 15:52:12.936 INFO: 
2026-05-27 15:52:12.936 INFO: ===========RESULTS===========
2026-05-27 15:52:13.004 INFO: Loading checkpoint: ./checkpoints/STO_finetuned_MATPES-run_run-56_epoch-39.pt
/jet/home/dvarghes/mace/mace/tools/checkpoint.py:187: UserWarning: Environment variable TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD detected, since the`weights_only` argument was not explicitly passed to `torch.load`, forcing weights_only=False.
  torch.load(f=checkpoint_info.path, map_location=device),
2026-05-27 15:52:13.602 INFO: Loaded Stage one model from epoch 39 for evaluation
2026-05-27 15:52:13.603 INFO: Saving model to checkpoints/STO_finetuned_MATPES-run_run-56.model
2026-05-27 15:52:16.388 INFO: Compiling model, saving metadata to STO_finetuned_MATPES-run_compiled.model
2026-05-27 15:52:20.289 INFO: Computing metrics for training, validation, and test sets
2026-05-27 15:52:20.290 INFO: Skipping evaluation for heads: ['pt_head']
2026-05-27 15:52:20.290 INFO: Evaluating train_Default ...
2026-05-27 15:53:07.030 INFO: Skipping evaluation of train_pt_head (in skip_heads list)
2026-05-27 15:53:07.030 INFO: Evaluating valid_Default ...
2026-05-27 15:53:09.545 INFO: Skipping evaluation of valid_pt_head (in skip_heads list)
2026-05-27 15:53:09.831 INFO: Error-table on TRAIN and VALID:
+---------------+---------------------+------------------+-------------------+---------------------------------------+
|  config_type  | RMSE E / meV / atom | RMSE F / meV / A | relative F RMSE % | RMSE Stress (Virials) / meV / A (A^3) |
+---------------+---------------------+------------------+-------------------+---------------------------------------+
| train_Default |            1.0      |         18.7     |          2.08     |                  None                 |
| valid_Default |            1.0      |         18.5     |          2.09     |                  None                 |
+---------------+---------------------+------------------+-------------------+---------------------------------------+
2026-05-27 15:54:05.803 INFO: Done

However, in the initial part of reading the data, we had:

Total Training set [energy: 1705, stress: 0, virials: 1705, dipole components: 0, head: 1705, elec_temp: 0, total_charge: 0, polarizability: 0, total_spin: 0, forces: 1705, charges: 0]

Does this mean the model is using the virial data? If not, how to use the virial data for training the model?

Thanks
Dominic

[ilyes319]

Hey @dominicvarghese, to train on virials you need to use the virials loss:

  atoms.info["REF_virials"] = virials

  Then train with:

  mace_run_train \
    --train_file dataset_with_virials.xyz \
    --loss virials \
    --virials_key REF_virials
    --error_table PerAtomMAEstressvirials

We are in the process of deprecating the virials, and convert everything to stresses. The easiest for you is to convert from virials to stresses.

  import numpy as np
  from ase.io import read, write

  atoms_list = read("dataset_with_virials.xyz", index=":")

  for atoms in atoms_list:
      volume = atoms.get_volume()
      virials = np.asarray(atoms.info["REF_virials"])
      atoms.info["REF_stress"] = -virials / volume

and then

 mace_run_train \
    --train_file dataset_with_stresses.xyz \
    --loss stress \
    --stress_key REF_stress

[dominicvarghese]

I used the above script to convert the virials to stress and was training a model from scratch.

2026-06-03 18:31:01.960 INFO: Epoch 397: head: Default, loss=0.00057025, RMSE_E_per_atom=    0.64 meV, RMSE_F=   13.58 meV / A, RMSE_stress=    2.52 meV / A^3
2026-06-03 18:33:33.062 INFO: Epoch 398: head: Default, loss=0.00057358, RMSE_E_per_atom=    0.58 meV, RMSE_F=   13.68 meV / A, RMSE_stress=    2.51 meV / A^3
2026-06-03 18:36:03.328 INFO: Epoch 399: head: Default, loss=0.00058015, RMSE_E_per_atom=    0.62 meV, RMSE_F=   13.88 meV / A, RMSE_stress=    2.50 meV / A^3
2026-06-03 18:36:03.328 INFO: Training complete
2026-06-03 18:36:03.328 INFO: 
2026-06-03 18:36:03.329 INFO: ===========RESULTS===========
2026-06-03 18:36:03.522 INFO: Loading checkpoint: ./checkpoints/STO_MACELES-run_run-56_epoch-397.pt
2026-06-03 18:36:04.338 INFO: Loaded Stage one model from epoch 397 for evaluation
2026-06-03 18:36:04.338 INFO: Saving model to checkpoints/STO_MACELES-run_run-56.model
2026-06-03 18:36:06.679 INFO: Compiling model, saving metadata to STO_MACELES-run_compiled.model
2026-06-03 18:36:11.297 INFO: Computing metrics for training, validation, and test sets
2026-06-03 18:36:11.298 INFO: Skipping evaluation for heads: ['pt_head']
2026-06-03 18:36:11.298 INFO: Evaluating train_Default ...
2026-06-03 18:37:04.697 INFO: Evaluating valid_Default ...
2026-06-03 18:37:07.873 INFO: Error-table on TRAIN and VALID:
+---------------+---------------------+------------------+-------------------+---------------------------------------+
|  config_type  | RMSE E / meV / atom | RMSE F / meV / A | relative F RMSE % | RMSE Stress (Virials) / meV / A (A^3) |
+---------------+---------------------+------------------+-------------------+---------------------------------------+
| train_Default |            0.7      |         14.5     |          1.64     |                     2.6               |
| valid_Default |            0.6      |         13.6     |          1.58     |                     2.5               |
+---------------+---------------------+------------------+-------------------+---------------------------------------+

The errors seems to be small after training the model for 400 epochs.

But, when I was checking the Parity plots, there seems to be mismatch in the stress errors.

Will the stress get improved, when trained for longer epochs?

  --compute_forces=True \
  --compute_stress=True \
  --energy_weight=10.0 \
  --forces_weight=100.0 \
  --stress_weight=10.0 \
  --loss="universal" \
  --energy_key="REF_energy" \
  --forces_key="REF_forces" \
  --stress_key="REF_stress" \
  --lr=0.005 \
  --valid_fraction=0.05 \
  --batch_size=10 \
  --max_num_epochs=400 

[gabor1]

There are some configs (or stress components) that are really good, and some have a constant shift. Can you identify what distinguishes these groups? This is not a training issue but some sort of data issue. Do you think you might have the kinetic contribution in there as well, rather than just the derivative of the energy with respect to the deformation tensor?

[dominicvarghese]

I was using the dataset from the paper - https://journals.aps.org/prb/pdf/10.1103/PhysRevB.108.115435

Here, the authors had STO datset in the cubic and monolayer structures.

First, ab initio MD (AIMD) simulations [41] with a time step of 2 fs and 500 steps were performed for supercells of cubic SrTiO3 and freestanding thin films with thicknesses of 5.97, 10.02, and 13.94 Å at 300, 500, and 700 K. Only the Γ point was used for the 𝑘-point integration in AIMD simulations. The initial training set was selected to include configurations every 10 steps. Then, classical MD simulations using the lammps package [42] with the pretrained potential and a total of 1 ns for the active learning process were conducted for all systems and temperatures. Finally, 1705 configurations were selected for accurate DFT calculations with a 3×3×3 ( 3×3×1) 𝑘-point mesh for the bulk (thin film) SrTiO3 to obtain the energies, atomic forces, and virials. In addition, 189 configurations including 50 configurations of the 2.2-nm-thick thin film not in the training set were selected from SrTiO3 MD trajectories to form the cross-validation set.

I combined their hidden set of 189 configurations with their initial training set - making the total configurations 1894.

Maybe I should remove that extra 189 from my training set as from the stress parity plot, we see a small chunk being good fits and a large part having the shift?

[gabor1]

This is useful background info, but please do what I asked. Identify which configurations or stress components are in the group on your plot that look really accurate and such are the ones that are shifted by a constant. There are at least two distinct groups of points that are shifted