Custom Relaxers

ASSYST’s relaxation step is designed to be extensible. If your preferred energy/force engine does not expose an ASE-compatible ase.calculators.calculator.Calculator, you can still plug it into the workflow by subclassing assyst.relaxations.Relax and overriding the relax() method.

When to subclass Relax

The built-in relax() implementation drives minimization through ASE’s LBFGS optimizer and therefore requires an ASE calculator to be attached to the Atoms object. A custom subclass is the right tool when:

• the external code has its own minimizer (e.g. a force-field engine with native geometry optimisation), or

• the energy/force interface is not easily wrapped as an ASE calculator.

The contract your override must satisfy:

1. Accept a single Atoms object (the structure to relax).

2. Call assyst.utils.update_uuid() on the returned structure so that provenance tracking continues to work.

3. Return a new Atoms object with a SinglePointCalculator carrying the final energy, forces, and stress.

Toy example

The snippet below shows a minimal custom relaxer that delegates geometry optimisation to a hypothetical external library myengine.

from dataclasses import dataclass

import numpy as np
from ase import Atoms
from ase.calculators.singlepoint import SinglePointCalculator

from assyst.relaxations import Relax
from assyst.utils import update_uuid


# ---------------------------------------------------------------------------
# Toy stand-in for a non-ASE energy/force engine
# ---------------------------------------------------------------------------

class MyEngineError(RuntimeError):
    pass


def myengine_run_relaxation(positions, cell, numbers, max_steps, ftol):
    """Pretend external relaxation routine.

    In practice this would call out to a C extension, a subprocess,
    or a REST API.  Here it simply returns the input unchanged with
    made-up energetics so the example runs without any real dependency.

    Returns
    -------
    relaxed_positions : np.ndarray, shape (N, 3)
    final_energy      : float   (eV)
    final_forces      : np.ndarray, shape (N, 3)  (eV/Å)
    final_stress      : np.ndarray, shape (6,)    (eV/ų, Voigt order)
    """
    # --- replace this block with actual engine calls ---
    n_atoms = len(positions)
    relaxed_positions = positions.copy()
    final_energy = -float(n_atoms)          # 1 eV/atom binding
    final_forces = np.zeros((n_atoms, 3))   # converged → forces ≈ 0
    final_stress = np.zeros(6)
    # ---------------------------------------------------
    return relaxed_positions, final_energy, final_forces, final_stress


# ---------------------------------------------------------------------------
# Custom Relax subclass
# ---------------------------------------------------------------------------

@dataclass(frozen=True, eq=True)
class MyEngineRelax(Relax):
    """Relax structures using ``myengine``'s native geometry optimiser.

    Inherits ``max_steps`` and ``force_tolerance`` from
    :class:`~assyst.relaxations.Relax`.
    """

    def relax(self, structure: Atoms) -> Atoms:
        # 1. Run the external relaxation
        relaxed_pos, energy, forces, stress = myengine_run_relaxation(
            positions=structure.get_positions(),
            cell=structure.get_cell(),
            numbers=structure.get_atomic_numbers(),
            max_steps=self.max_steps,
            ftol=self.force_tolerance,
        )

        # 2. Build the output Atoms object
        relaxed = structure.copy()
        relaxed.set_positions(relaxed_pos)

        # 3. Attach a SinglePointCalculator with the final energetics
        relaxed.calc = SinglePointCalculator(
            relaxed,
            energy=energy,
            forces=forces,
            stress=stress,
        )

        # 4. Update provenance (UUID / lineage) — do not skip this step
        update_uuid(relaxed)

        return relaxed

Using the custom relaxer in the workflow

Once defined, MyEngineRelax is a drop-in replacement anywhere Relax is accepted:

from assyst.relaxations import relax as assyst_relax

settings = MyEngineRelax(max_steps=200, force_tolerance=5e-4)

relaxed_structures = list(
    assyst_relax(
        structures=my_structures,
        settings=settings,
        calculator=None,   # calculator is unused by MyEngineRelax.relax
    )
)

Note

The top-level assyst.relaxations.relax() function attaches an ASE calculator to each structure before calling settings.relax. If your custom relax method does not need an ASE calculator you can pass calculator=None and iterate over the structures directly, bypassing the helper function altogether:

relaxer = MyEngineRelax()
relaxed_structures = [relaxer.relax(s) for s in my_structures]

No metadata is lost by this approach, provided your relax implementation follows the contract above (in particular, the update_uuid() call).