Issue when running Lammps with Kokkos GPU and Allegro potential for calculating Elastic Constants

[ArcXeon]

I have trained an allegro potential model for a triclinic structure.

After training I am trying to calculate its Elastic Constants using Lammps. I am using the Elastic Constants at 0K script provided by Lammps Examples.

I am running Lammps using GPU Kokkos. But for every run, I am getting different values for Elastic Constants and the values for Elastic Tensor calculated is machanically unstable.

What is the solution to this?

[cw-tan]

Is there an analogous computation you can do in ASE (https://nequip.readthedocs.io/en/latest/integrations/ase.html) to verify whether the problem is the potential or whatever you're running?

[cw-tan]

Can you write the analogous elastic constants computation script to use ASE? The point is to isolate if it's a model-specific problem, or a problem with how you're running things.

[ArcXeon]

import numpy as np
from ase.io import read as ase_read
from ase.data import atomic_numbers
from ase.optimize import GoodOldQuasiNewton, FIRE
from ase import units
import cuequivariance_torch
from nequip.ase import NequIPCalculator
import torch

# -------------------------
# Inputs
# -------------------------
compiled_model = "deployed_allegro_ase.nequip.pth"
lammps_data = "relaxed.data"
lammps_units = "metal"
lammps_atom_style = "atomic"

mapping_symbols = "Al S P Li".split()
Z_of_type = {i + 1: atomic_numbers[sym] for i, sym in enumerate(mapping_symbols)}

# Strain sweep parameters
delta = 5e-3                 # engineering strain magnitude (0.5%)
relax_fmax = 1e-3            # eV/Å; 0 K internal relaxation threshold
relax_steps = 1000            # max steps per internal relaxation
optimizer_cls = GoodOldQuasiNewton  

# ---------------
# Read structure 
# ---------------
atoms0 = ase_read(
    lammps_data,
    format="lammps-data",
    units=lammps_units,
    atom_style=lammps_atom_style,
    Z_of_type=Z_of_type,
)
atoms0.set_pbc(True)
atoms0.wrap(eps=1e-12)

# -------------------------
# Initialize calculator (species from structure)
# -------------------------
calc = NequIPCalculator.from_compiled_model(
    compile_path=compiled_model,
    chemical_symbols=["Li", "Al", "P", "S"],
    device="cuda" if torch.cuda.is_available() else "cpu",
)
atoms0.calc = calc

# -------------------------
# Helpers
# -------------------------
voigt_labels = ["xx", "yy", "zz", "yz", "xz", "xy"] 

def build_strain_tensor(voigt_index, gamma):
    """Return small symmetric strain tensor for a given Voigt mode.
    For shear modes, gamma is the engineering shear strain (E_xy = gamma/2)."""
    eps = np.zeros((3, 3), dtype=float)
    if voigt_index == 0:   # xx
        eps[0, 0] = gamma
    elif voigt_index == 1: # yy
        eps[1, 1] = gamma
    elif voigt_index == 2: # zz
        eps[2, 2] = gamma
    elif voigt_index == 3: # yz
        eps[1, 2] = eps[2, 1] = 0.5 * gamma
    elif voigt_index == 4: # xz
        eps[0, 2] = eps[2, 0] = 0.5 * gamma
    elif voigt_index == 5: # xy
        eps[0, 1] = eps[1, 0] = 0.5 * gamma
    return eps

def apply_strain_and_relax(cell0, eps, atoms_template):
    """Apply strain to cell, scale atoms affinely, relax internal positions, return stress (Voigt)."""
    # Right-multiply to deform lattice vectors; valid for small strains
    F = np.eye(3) + eps
    new_cell = cell0 @ F
    atoms = atoms_template.copy()
    atoms.set_cell(new_cell, scale_atoms=True)
    atoms.calc = atoms_template.calc

    # 0 K internal relaxation at fixed cell
    with optimizer_cls(atoms, logfile="/dev/null") as opt:
        opt.run(fmax=relax_fmax, steps=relax_steps)

    # Stress in Voigt order: [sxx, syy, szz, syz, sxz, sxy]
    sig_voigt = atoms.get_stress(voigt=True)  # eV/Å^3
    return np.array(sig_voigt, dtype=float)

# -------------------------
# Finite-difference stiffness matrix
# -------------------------
C = np.zeros((6, 6), dtype=float)  # eV/Å^3

cell0 = atoms0.cell.array.copy()

for i in range(6):
    # +delta
    eps_p = build_strain_tensor(i, +delta)
    sig_p = apply_strain_and_relax(cell0, eps_p, atoms0)

    # -delta
    eps_m = build_strain_tensor(i, -delta)
    sig_m = apply_strain_and_relax(cell0, eps_m, atoms0)

    # Central difference: C[:, i] = d sigma / d (engineering strain component i)
    C[:, i] = (sig_p - sig_m) / (2.0 * delta)

# Enforce minor symmetry numerically
C = 0.5 * (C + C.T)

# Convert to GPa
C_GPa = C / units.GPa

# -------------------------
# Report
# -------------------------
np.set_printoptions(precision=3, suppress=True)
print("Voigt order [xx, yy, zz, yz, xz, xy]")
print("Elastic stiffness tensor C (GPa):")
print(C_GPa)


C11, C22, C33 = C_GPa[0,0], C_GPa[1,1], C_GPa[2,2]
C12, C13, C23 = C_GPa[0,1], C_GPa[0,2], C_GPa[1,2]
C44, C55, C66 = C_GPa[3,3], C_GPa[4,4], C_GPa[5,5]
print(f"C11={C11:.2f}, C22={C22:.2f}, C33={C33:.2f}, C12={C12:.2f}, C13={C13:.2f}, C23={C23:.2f}, "
    f"C44={C44:.2f}, C55={C55:.2f}, C66={C66:.2f}")

@cw-tan So I used this script to calculate the Elastic Constants using the same potential model, and it is giving me a mechanically stable elastic tensor and the output is same for every time I run this code.

[cw-tan]

Then maybe it's about how the LAMMPS script is being run? Maybe more details about that? Ideally the result you get from ASE and LAMMPS match

[ArcXeon]

Could it be because I am using float32 model? Does Lammps work properly only with float64 models?

[cw-tan]

Unlikely. All inputs the model takes are in float64 and our LAMMPS integrations are configured to pass data in float64 to the model. float32 models will cast down to float32 internally after the initial embeddings. So it's unlikely to be some mismatch between the model precision and LAMMPS. If you think it's some numerical reason though, then I don't know. All I can comment on is that we've typically found float32 models to be precise enough for various workflows including geometry relaxations and phonon calculations.