Python interface

Jupyter Notebook tutorials illustrating the use of currently available calculations can be found in the tutorials documentation directory. This currently includes examples for:

• Single Point

• Geometry Optimization

• Molecular Dynamics

• Equation of State

• Phonons

• Nudged Elastic Band

These make use of WEAS Widget for visualisation, which can be installed using janus-core’s visualise extra. For example:

pip install janus-core[mace,visualise]

or

pip install janus-core[all]

Calculation outputs

By default, calculations performed will modify the underlying ase.Atoms object to store information in the Atoms.info and Atoms.arrays dictionaries about the MLIP used.

Additional dictionary keys include arch, corresponding to the MLIP architecture used, and model, corresponding to the model path, name or label.

Results from the MLIP calculator, which are typically stored in Atoms.calc.results, will also, by default, be copied to these dictionaries, prefixed by the MLIP arch.

For example:

from janus_core.calculations.single_point import SinglePoint

single_point = SinglePoint(
    struct="tests/data/NaCl.cif",
    arch="mace_mp",
    model="tests/models/mace_mp_small.model",
)

single_point.run()
print(single_point.struct.info)

will return

{
    'spacegroup': Spacegroup(1, setting=1),
    'unit_cell': 'conventional',
    'occupancy': {'0': {'Na': 1.0}, '1': {'Cl': 1.0}, '2': {'Na': 1.0}, '3': {'Cl': 1.0}, '4': {'Na': 1.0}, '5': {'Cl': 1.0}, '6': {'Na': 1.0}, '7': {'Cl': 1.0}},
    'model': 'tests/models/mace_mp_small.model',
    'arch': 'mace_mp',
    'mace_mp_energy': -27.035127799332745,
    'mace_mp_stress': array([-4.78327600e-03, -4.78327600e-03, -4.78327600e-03,  1.08000967e-19, -2.74004242e-19, -2.04504710e-19]),
    'system_name': 'NaCl',
}

Note

If running calculations with multiple MLIPs, arch and mlip_model will be overwritten with the most recent MLIP information. Results labelled by the architecture (e.g. mace_mp_energy) will be saved between MLIPs, unless the same arch is chosen, in which case these values will also be overwritten.

D3 Dispersion

A PyTorch implementation of DFTD2 and DFTD3, using the TorchDFTD3Calculator, can be used to add dispersion corrections to MLIP predictions.

The required Python pacakge is included with mace_mp, but can also be installed as its own extra:

pip install janus-core[d3]

Once installed, dispersion can be added through calc_kwargs through the dispersion keyword, with dispersion_kwargs used to pass any further keywords to the TorchDFTD3Calculator:

from ase import units

from janus_core.calculations.single_point import SinglePoint

single_point = SinglePoint(
    struct="tests/data/NaCl.cif",
    arch="mace_mp",
    model="tests/models/mace_mp_small.model",
    calc_kwargs={"dispersion": True, "dispersion_kwargs": {"cutoff":  95.0 * units.Bohr}}
)

Note

In most cases, defaults for dispersion_kwargs are those set within TorchDFTD3Calculator, but in the case of mace_mp, we mirror the corresponding defaults from the mace.calculators.mace_mp function.

The TorchDFTD3Calculator can also be added to any existing calculator if required:

Note

Keyword arguments for TorchDFTD3Calculator should be passed directly here, as shown with cutoff. This will not have access to mace_mp default values, so will always use defaults from TorchDFTD3Calculator.

from ase import units

from janus_core.helpers.mlip_calculators import add_dispersion, choose_calculator

mace_calc = choose_calculator("mace_mp")
calc = add_dispersion(mace_calc, device="cpu", cutoff=95 * units.Bohr)

Additional Calculators

Although janus-core only directly supports the MLIP calculators listed in Getting started, any valid ASE calculator can be attached to a structure, including calculators for currently unsupported MLIPs.

This structure can then be passed to janus-core calculations, which can be run as usual.

For example, performing geometry optimisation using the (ASE built-in) Lennard Jones potential calculator:

from janus_core.calculations.geom_opt import GeomOpt
from ase.calculators.lj import LennardJones
from ase.io import read

struct = read("tests/data/NaCl-deformed.cif")
struct.calc = LennardJones()

geom_opt = GeomOpt(
    struct=struct,
    fmax=0.001,
)
geom_opt.run()