Examples¶

from __future__ import annotations
from chgnet.model.dynamics import CHGNetCalculator
import numpy as np
from ase import Atoms
from pymatgen.io.ase import AseAtomsAdaptor
from ase.calculators.calculator import Calculator, all_changes, all_properties


class AMSCHGNetCalculator(CHGNetCalculator):
    """
    Override the CHGNetCalculator calculate() method.

    Modifications made:

    * add a lattice for non-periodic systems with a gap of at least 10 angstrom

    * only calculate the stress tensor if requested

    The original code can be found at https://github.com/CederGroupHub/chgnet
    """

    def calculate(
        self,
        atoms: Atoms | None = None,
        properties: list | None = None,
        system_changes: list | None = None,
    ) -> None:
        """Calculate various properties of the atoms using CHGNet.

        Args:
            atoms (Atoms | None): The atoms object to calculate properties for.
            properties (list | None): The properties to calculate.
                Default is all properties.
            system_changes (list | None): The changes made to the system.
                Default is all changes.
        """

        if sum(atoms.get_pbc()) == 0:
            delta = np.max(atoms.get_positions()) - np.min(atoms.get_positions())
            atoms.set_cell([delta + 10, delta + 10, delta + 10])

        properties = properties or all_properties
        system_changes = system_changes or all_changes
        Calculator.calculate(
            self,
            atoms=atoms,
            properties=properties,
            system_changes=system_changes,
        )

        # Run CHGNet
        structure = AseAtomsAdaptor.get_structure(atoms)
        graph = self.model.graph_converter(structure)
        task = "e"
        if "forces" in properties:
            task += "f"
        if "stress" in properties:
            task += "s"
        if "magmoms" in properties:
            task += "m"
        model_prediction = self.model.predict_graph(graph.to(self.device), task=task)

        # Convert Result
        factor = 1 if not self.model.is_intensive else structure.composition.num_atoms

        self.results = dict()
        self.results["energy"] = model_prediction["e"] * factor
        self.results["free_energy"] = model_prediction["e"] * factor
        if "f" in model_prediction:
            self.results["forces"] = model_prediction["f"]
        if "s" in model_prediction:
            self.results["stress"] = model_prediction["s"] * self.stress_weight
        if "m" in model_prediction:
            self.results["magmoms"] = model_prediction["m"]


def get_calculator(use_device: str = "cpu", **kwargs):
    return AMSCHGNetCalculator(use_device=use_device, **kwargs)

import numpy
import ase
import gpaw


class ASEGPAWCalculator(ase.calculators.calculator.Calculator):
    def __init__(
        self,
        pbc=[True, True, True],
        cancel_total_force=False,
        charge=0,
        name="atoms",
        **gpaw_kwargs,
    ):
        self.name = name
        self.counter = 1
        self.pbc = pbc
        self.cancel = cancel_total_force
        self._kwargs = gpaw_kwargs
        ase.calculators.calculator.Calculator.__init__(self)
        self._setup_gpaw(charge)

    def calculate(self, atoms=None, properties=None, system_changes=None):
        atoms.center()
        atoms.set_pbc(self.pbc)
        super().calculate(atoms, properties, system_changes)

        self.gpaw.calculate(atoms, properties, system_changes)
        self.results = self.gpaw.results
        # remove total force on the molecule
        if self.cancel:
            molecule_force = self.results["forces"].sum(axis=0)
            per_atom_force = molecule_force / self.results["forces"].shape[0]
            self.results["forces"] -= per_atom_force

    def _setup_gpaw(self, charge):
        self.charge = charge
        txt = self.name
        if self.counter > 1:
            txt = txt + f"_{self.counter}"
        txt = txt + ".txt"
        self.gpaw = gpaw.GPAW(txt=txt, **self._kwargs)
        self.counter += 1

    @property
    def implemented_properties(self):
        return self.gpaw.implemented_properties


# AMS looks for "get_calculator"
get_calculator = ASEGPAWCalculator

import os
from ase.calculators.calculator import Calculator, all_changes
import subprocess
import pickle

""" 
Example for using xtb (GFN2-xTB) from a different python environment together with AMS 

IMPORTANT: Set the xtb_python_environment variable to the conda python executable where xtb is installed!

Ideally xtb-python would be installed into the AMS python environment, then one
could just access the XTB calculator directly without having to go the
roundabout way via files. But it's easier to install xtb into a separate conda environment.

To install xtb and xtb-python into a conda environment, do:
conda config --add channels conda-forge
conda install xtb
conda install xtb-python
conda install ase

"""

xtb_python_environment = "/home/hellstrom/anaconda3/bin/python3"


class MyCalculator(Calculator):
    """this class is used inside the AMS python environment"""

    implemented_properties = ["energy", "forces", "charges"]

    def __init__(self, python: str, file: str):
        Calculator.__init__(self)
        self.python = python
        self.file = file

    def calculate(self, atoms=None, properties=["energy"], system_changes=all_changes):
        with open("atoms.pkl", "wb") as f:
            pickle.dump(atoms, f)
        subprocess.run([self.python, self.file])
        with open("results.pkl", "rb") as f:
            self.results = pickle.load(f)
        if os.path.exists("atoms.pkl"):
            os.remove("atoms.pkl")
        if os.path.exists("results.pkl"):
            os.remove("results.pkl")


def get_calculator(**kwargs):
    """this function is imported in the amspython environment"""
    return MyCalculator(python=os.path.abspath(xtb_python_environment), file=os.path.abspath(__file__))


def run_xtb_and_pickle():
    """this is run in the python environment where xTB is installed"""
    from xtb.ase.calculator import XTB

    with open("atoms.pkl", "rb") as f:
        atoms = pickle.load(f)
    atoms.calc = XTB(method="GFN2-xTB")
    atoms.get_forces()
    with open("results.pkl", "wb") as f:
        pickle.dump(atoms.calc.results, f)


if __name__ == "__main__":
    """this is run in the python environment where xTB is installed"""
    run_xtb_and_pickle()

#!/usr/bin/env amspython
import scm.plams as plams
from os.path import abspath

"""
Example showing how to run the external xtb program using GFN2-xTB.

NOTE 1: To instead run GFN1-xTB, you can use the DFTB engine inside AMS.

NOTE 2: You must modify the path below to give the correct path to the xtb_calculator.py file

NOTE 3: To run this script, use the command 
$AMSBIN/amspython run_xtb.py
"""


def main():
    plams.init()
    molecule = plams.from_smiles("O")
    s = plams.Settings()
    s.input.ams.Task = "GeometryOptimization"
    s.input.ASE.File = abspath("/home/hellstrom/xtb_calculator.py")
    s.runscript.nproc = 1  # make sure to run AMS in serial!
    job = plams.AMSJob(settings=s, molecule=molecule, name="opt_GFN2-xTB")
    job.run()

    optimized_molecule = job.results.get_main_molecule()

    plams.log(f"Energy: {job.results.get_energy(unit='eV'):.3f} eV")

    # access atomic charges
    # the charges are stored in "Other%charges" on ase.rkf
    plams.log("Atomic charges:")
    try:
        charges = job.results.get_charges()
        for at, charge in zip(optimized_molecule, charges):
            plams.log(f"{at.symbol} {charge:.3f}")
    except KeyError:
        plams.log("Couldn't get charges")


if __name__ == "__main__":
    main()