Skip to content

Optimisation with an ASE calculator

chmpy.opt relaxes molecules and crystals with a trust-region optimiser, and can compute elastic constants and force constants on top. It does not need ASE, but most machine-learned potentials ship as ASE calculators. You can use any of them by wrapping it once:

from chmpy.calc import Calculator

# any ase.calculators.calculator.Calculator works here
from mace.calculators import mace_mp
calc = Calculator.from_ase(mace_mp(model="medium", default_dtype="float64"))

To try the examples below without installing a model, use ASE's built-in effective-medium potential: from ase.calculators.emt import EMT and calc = Calculator.from_ase(EMT()). EMT covers a handful of metals (Al, Cu, Ag, Au, Ni, Pd, Pt) and gives only rough results for H, C, N and O.

For PET-MAD and the other UPET models there is also a native route that skips ASE entirely. It takes the same dtype and device options:

from chmpy.calc.adapters.metatomic import upet

calc = upet("pet-mad-s", dtype="float64")   # or Calculator.from_ase(UPETCalculator(...))

Use float64 if you plan to compute elastic constants or phonons: both take differences of forces, and float32 noise is amplified by them.

The same models also run on pet-kokkos, a C++ implementation that needs no torch and is several times faster on small cells. pip install 'chmpy[pet]' gives the CPU build; for an NVIDIA GPU install the CUDA one from its own index instead (pip install pet-kokkos --extra-index-url https://peterspackman.github.io/pet-kokkos/cu12):

from chmpy.calc.adapters.pet_kokkos import PetKokkosCalculator

calc = PetKokkosCalculator("pet-mad-s")      # downloaded on first use
calc = PetKokkosCalculator("pet-mad-s", gemm="tf32", attention="tf32")  # faster on GPUs

pet-kokkos computes in fp32 (with fp64 sums), so the same caveat applies to elastic constants and phonons. The tf32 modes are good for relaxations and MD, but their energies scatter by about 1e-5 relative, which the optimiser is told about and works around.

The ASE wrapper keeps one ase.Atoms alive and writes each new geometry into it, rather than copying a structure on every call. Results are cached, so asking for the same geometry twice costs nothing.

Molecules

from chmpy import Molecule
from chmpy.opt import relax

water = Molecule.load("water.xyz")
result = relax(water, calc, fmax=0.01)
result
# <Relaxation converged in 6 steps (7 evaluations), E=1.879225 eV, fmax=0.009192>

relaxed = result.structure   # a Molecule, like the input
relaxed.save("water_relaxed.xyz")

fmax is the largest force in eV/Å. relax takes the same arguments whatever kind of structure you pass. progress=True prints one line per step (see Progress), and fixed= takes a boolean mask of atoms to hold still.

Some models need a charge or a spin as input. Pass them with info, and they are copied into atoms.info on every call:

result = relax(water, calc, info={"charge": 0, "spin": 1})

Crystals

By default a Crystal is relaxed inside its space group. The degrees of freedom are the asymmetric unit plus the cell strains the point group allows, so the symmetry holds exactly at every step rather than being restored afterwards. This needs about 1/|G| as many variables as a P1 relaxation, and noisy forces from an MLIP cannot push the structure into a lower symmetry.

import numpy as np
from chmpy import Crystal
from chmpy.core import Element
from chmpy.crystal import AsymmetricUnit, SpaceGroup, UnitCell

copper = Crystal(
    UnitCell.from_lengths_and_angles((3.7, 3.7, 3.7), np.radians((90, 90, 90))),
    SpaceGroup(225),
    AsymmetricUnit([Element["Cu"]], np.array([[0.0, 0.0, 0.0]])),
)
result = relax(copper, calc, fmax=1e-4, smax=1e-3)
result
# <Relaxation converged in 3 steps (4 evaluations), E=-0.028146 eV, fmax=9.562e-15, smax=0.0006596>
result.structure
# <Crystal Cu Fm-3m>
result.structure.unit_cell.lengths
# [3.5898, 3.5898, 3.5898]

smax is the largest stress component in GPa. Forces and stress are checked separately, each in its own units.

Structures loaded from files work the same way:

crystal = Crystal.load("acetic_acid.cif")

relax(crystal, calc)                     # atoms and cell, in Pna2_1
relax(crystal, calc, cell=False)         # atoms only, at fixed cell
relax(crystal, calc, pressure=1.0)       # minimise the enthalpy at 1 GPa
relax(crystal, calc, symmetry=False)     # P1: every atom, all six strains
relax(crystal, calc, stages="two-stage") # fixed cell first, then everything

result = relax(crystal, calc)
result.structure.save("acetic_acid_relaxed.cif")  # still in Pna2_1

If a structure has no freedoms at all, relax raises rather than spending a calculation to find that out. Examples are fcc copper at fixed cell, or any crystal whose atoms all sit on fixed special positions.

Starting from an ase.Atoms

An Atoms object carries no space group, so it is relaxed in P1. Convert it to a Crystal first if you want symmetry; alternatively, wrap it as a System:

from chmpy.calc import System

system = System.from_ase(atoms)
result = relax(system, calc)        # atoms, plus the cell if periodic
relaxed_atoms = result.structure.to_ase()

Elastic constants and phonons

Both start from a relaxed structure and use its symmetry to reduce the number of calculations:

from chmpy.opt import elastic_tensor
from chmpy.vib import force_constants

relaxed = relax(copper, calc, fmax=1e-4, smax=1e-3).structure

elastic = elastic_tensor(relaxed, calc)
elastic
# <ElasticResult 3 independent constants, 2 strains, 5 evaluations, ...>
elastic.c_voigt[:, 0]
# [172.6, 115.5, 115.5, 0, 0, 0]   GPa
elastic.tensor.is_stable()
# True

constants = force_constants(relaxed, calc, cutoff=6.0)
constants.frequencies([0.5, 0.5, 0.5], units="thz")

For a cubic crystal, elastic_tensor applies two strains rather than six. The responses to the others are rotations of these two. At each strained geometry, the atoms relax within the subgroup that the strain leaves intact. force_constants displaces one atom per symmetry orbit, and only along the directions its site symmetry does not already determine. Pass symmetry=False to either function to do every calculation explicitly, which is a useful cross-check.

For a noisy model, elastic_tensor(..., strain="auto") keeps increasing the strain until the tensor's own noise estimate is acceptable. The result reports noise_fraction, and it logs a warning if the tensor should not be trusted.

Mode symmetry

gamma_modes gives the irreducible representations of a crystal's modes at Gamma, and which are infrared or Raman active, from its symmetry alone. label_gamma_modes labels the modes computed from a set of force constants. For molecular crystals, rigid_body_fractions gives how much of each mode is whole-molecule translation or libration:

from chmpy.vib import gamma_modes, label_gamma_modes, rigid_body_fractions

print(gamma_modes(naphthalene))
# optical: 27Ag + 27Bg + 26Au + 25Bu
# infrared: 26Au + 25Bu
# raman: 27Ag + 27Bg
# silent: none

constants = force_constants(naphthalene, calc)
modes = label_gamma_modes(naphthalene, constants)
fractions = rigid_body_fractions(naphthalene, modes)
for mode, (translation, libration) in zip(modes[:4], fractions[:4]):
    print(mode, f"{translation:.2f} {libration:.2f}")
#       0.00  Au + 2Bu             1.00 0.00
#      34.01  Bg         Raman     0.00 1.00
#      42.89  Au         IR        0.99 0.00
#      57.18  Ag         Raman     0.00 1.00

The acoustic modes are the uniform translations, projected out before the dynamical matrix is diagonalised. A soft or imaginary optical mode is never taken for an acoustic one, and keeps its label: in cubic BaTiO3 the unstable mode is the polar T1u. lattice_modes gives the lattice modes expected by the correlation method, here 3Ag + 3Bg librations and 2Au + Bu translations. In a centred cell, modes that belong to the zone boundary of the primitive cell are folded onto Gamma, and they are left unlabelled.

Vibrational free energies

thermal_properties sums the harmonic free energy over a q-mesh, using a set of force constants measured in a supercell:

from chmpy.vib import thermal_properties

constants = force_constants(relaxed, calc)
thermal_properties(constants, [100.0, 298.15], mesh=(8, 8, 8)).free_energy
# eV per unit cell

For molecular crystals, the supercell needed to converge this can run to thousands of atoms. The Einstein-Debye model of Nickerson and Johnson (PCCP 2026, DOI 10.1039/D6CP00089D) needs only the unit cell. It treats the optical modes at Gamma as Einstein oscillators, and the acoustic modes as a Debye solid whose cutoff comes from the sound velocities of the elastic tensor:

from chmpy.vib import einstein_debye

result = einstein_debye(relaxed, calc, [100.0, 298.15])
result.per_molecule("free_energy")
# kJ/mol per molecule, for ranking polymorphs
result.debye.temperature, result.n_imaginary

If you already have the elastic tensor, pass it with elastic= to save computing it again. For paracetamol (HXACAN01) with PET-MOLS-S, the result is within 0.05 kJ/mol per molecule of a 2x2x3 supercell on an 8x8x8 mesh, at about a tenth of the cost.

Lattice energies

result = calc.lattice_energy(Crystal.load("acetic_acid.cif"))

This relaxes the crystal and each symmetry-unique molecule, then takes the difference per molecule. See chmpy.opt.lattice_energy.

Progress

relax, elastic_tensor, force_constants and lattice_energy all take a progress argument. It is silent by default; progress=True prints:

calc.lattice_energy(crystal, progress=True)
# [1/2] relaxing the crystal (<Crystal C8H9NO2 P2_1/n>, Z=4)
#   relax <SymmetryAdapted P2_1/n: 60 atomic + 4 cell dof for 80 atoms>
#        1 +    E=   -485.71...  fmax=  0.29... smax=  0.41...  delta=...
#        ...
# [1/2] crystal: <Relaxation converged in 41 steps ...>
# [2/2] relaxing molecule 0 (C8H9NO2)
#   ...

Pass a callable to get a chmpy.opt.progress.Progress for every event instead. Each event has task, stage and message fields, an index out of a total when the count is known, and a depth with the parents it is nested inside. Events from inside a relaxation also carry the optimiser's step. Use these to drive a progress bar without parsing text:

from tqdm import tqdm

bar = None

def show(event):
    global bar
    if event.depth == 0 and event.total and not event.done:
        if bar is None:
            bar = tqdm(total=event.total)
        bar.set_description(event.stage)
        bar.update(1)

elastic_tensor(relaxed, calc, progress=show)

Going the other way

A chmpy calculator can be handed to code that expects an ASE one:

from ase.optimize import BFGS

atoms = relaxed.to_ase_atoms()
atoms.calc = calc.as_ase()
BFGS(atoms).run(fmax=0.01)

Complete examples

Two command-line scripts in examples/ put this together with MACE-Polar:

  • examples/lattice_energy.py structure.cif [--json] [-d out/] relaxes a molecular crystal and its molecules, and reports the lattice energy with its interaction and strain parts.
  • examples/elastic_tensor.py structure.cif [--strain auto] [--json] [-d out/] relaxes the crystal, then computes the relaxed-atom elastic tensor, its averaged moduli, and the Young's modulus, linear compressibility and sound velocities along each cell axis.