janus-core

Tools for machine learnt interatomic potentials

Contents

• Getting started
• Features
• Python interface
• Command line interface
• Docker/Podman images
• Development
• License
• Funding

Getting started

Dependencies

All required and optional dependencies can be found in pyproject.toml.

Installation

The latest stable release of janus-core, including its dependencies, can be installed from PyPI by running:

python3 -m pip install janus-core

To get all the latest changes, janus-core can also be installed from GitHub:

python3 -m pip install git+https://github.com/stfc/janus-core.git

By default, no machine learnt interatomic potentials (MLIPs) will be installed with janus-core. These can be installed separately, or as extras.

For example, to install MACE, CHGNet, and SevenNet, run:

python3 -m pip install janus-core[mace,chgnet,sevennet]

Warning

We are unable to support for automatic installation of all combinations of MLIPs, or MLIPs on all platforms. Please refer to the installation documentation for more details.

To install all MLIPs currently compatible with MACE, run:

python3 -m pip install janus-core[all]

Individual extras are listed in Getting Started, as well as in pyproject.toml under [project.optional-dependencies].

Further help

Please see Getting Started, as well as guides for janus-core's Python and command line interfaces, for additional information, or open an issue if something doesn't seem right.

Features

Unless stated otherwise, MLIP calculators and calculations rely heavily on ASE.

Current and planned features include:

• Support for multiple MLIPs
  • MACE
  • M3GNet
  • CHGNet
  • ALIGNN
  • SevenNet
  • NequIP
  • DPA3
  • Orb
  • MatterSim
  • GRACE
  • EquiformerV2
  • eSEN
  • UMA
  • PET-MAD
• Single point calculations
• Geometry optimisation
• Molecular Dynamics
  • NVE
  • NVT (Langevin(Eijnden/Ciccotti flavour) and Nosé-Hoover (Melchionna flavour))
  • NPT (Nosé-Hoover (Melchiona flavour))
• Nudged Elastic Band
• Phonons
  • Phonopy
• Equation of State
• Training ML potentials
  • MACE
• Fine-tuning MLIPs
  • MACE
• MLIP descriptors
  • MACE
• Data preprocessing
  • MACE
• Rare events simulations
  • PLUMED
• Elasticity

Python interface

Calculations can also be run through the Python interface. For example, running:

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",
)

results = single_point.run()
print(results)

will read the NaCl structure file and attach the MACE-MP (medium) calculator, before calculating and printing the energy, forces, and stress.

Tutorials

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

• Single Point

• Geometry Optimization

• Molecular Dynamics

• Equation of State

• Phonons

• Nudged Elastic Band

• Elasticity

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.

This is also the case the calculations performed using the CLI, with the same information written to extxyz output files.

Tip

For complete provenance tracking, calculations and training can be run using the aiida-mlip AiiDA plugin.

Command line interface

All supported MLIP calculations are accessible through subcommands of the janus command line tool, which is installed with the package:

janus singlepoint
janus geomopt
janus md
janus phonons
janus eos
janus neb
janus train
janus descriptors
janus preprocess
janus elasticity

For example, a single point calcuation (using the MACE-MP "small" force-field) can be performed by running:

janus singlepoint --struct tests/data/NaCl.cif --arch mace_mp --model small

A description of each subcommand, as well as valid options, can be listed using the --help option. For example,

janus singlepoint --help

prints the following:

Usage: janus singlepoint [OPTIONS]

Perform single point calculations and save to file.


╭─ Options ───────────────────────────────────────────────────────────────────────────╮
│ --config        TEXT  Path to configuration file.                                   │
│ --help                Show this message and exit.                                   │
╰─────────────────────────────────────────────────────────────────────────────────────╯
╭─ MLIP calculator ───────────────────────────────────────────────────────────────────╮
│ *  --arch               [mace|mace_mp|mace_off|m3gne  MLIP architecture to use for  │
│                         t|chgnet|alignn|sevennet|neq  calculations.                 │
│                         uip|dpa3|orb|mattersim|grace  [required]                    │
│                         |esen|equiformer|pet_mad|uma                                │
│                         |mace_omol]                                                 │
│    --device             [cpu|cuda|mps|xpu]            Device to run calculations    │
│                                                       on.                           │
│                                                       [default: cpu]                │
│    --model              TEXT                          MLIP model name, or path to   │
│                                                       model.                        │
│    --calc-kwargs        DICT                          Keyword arguments to pass to  │
│                                                       selected calculator. Must be  │
│                                                       passed as a dictionary        │
│                                                       wrapped in quotes, e.g.       │
│                                                       "{'key': value}".             │
╰─────────────────────────────────────────────────────────────────────────────────────╯
╭─ Calculation ───────────────────────────────────────────────────────────────────────╮
│ *  --struct            PATH                          Path of structure to simulate. │
│                                                      [required]                     │
│    --properties        [energy|stress|forces|hessia  Properties to calculate. If    │
│                        n]                            not specified, 'energy',       │
│                                                      'forces' and 'stress' will be  │
│                                                      returned.                      │
│    --out               PATH                          Path to save structure with    │
│                                                      calculated results. Default is │
│                                                      inferred from `file_prefix`.   │
╰─────────────────────────────────────────────────────────────────────────────────────╯
╭─ Structure I/O ─────────────────────────────────────────────────────────────────────╮
│ --file-prefix         PATH  Prefix for output files, including directories. Default │
│                             directory is ./janus_results, and default filename      │
│                             prefix is inferred from the input stucture filename.    │
│ --read-kwargs         DICT  Keyword arguments to pass to ase.io.read. Must be       │
│                             passed as a dictionary wrapped in quotes, e.g. "{'key': │
│                             value}". By default, read_kwargs['index'] = ':', so all │
│                             structures are read.                                    │
│ --write-kwargs        DICT  Keyword arguments to pass to ase.io.write when saving   │
│                             any structures. Must be passed as a dictionary wrapped  │
│                             in quotes, e.g. "{'key': value}".                       │
╰─────────────────────────────────────────────────────────────────────────────────────╯
╭─ Logging/summary ───────────────────────────────────────────────────────────────────╮
│ --log                                  PATH  Path to save logs to. Default is       │
│                                              inferred from `file_prefix`            │
│ --tracker         --no-tracker               Whether to save carbon emissions of    │
│                                              calculation                            │
│                                              [default: tracker]                     │
│ --summary                              PATH  Path to save summary of inputs,        │
│                                              start/end time, and carbon emissions.  │
│                                              Default is inferred from               │
│                                              `file_prefix`.                         │
│ --progress-bar    --no-progress-bar          Whether to show progress bar.          │
│                                              [default: progress-bar]                │
╰─────────────────────────────────────────────────────────────────────────────────────╯

Please see the user guide for examples of each subcommand.

Tutorials

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

• Single Point

• Geometry Optimization

• Molecular Dynamics

• Phonons

• Nudged Elastic Band

• Elasticity

Using configuration files

Default values for all command line options may be specifed through a Yaml 1.1 formatted configuration file by adding the --config option. If an option is present in both the command line and configuration file, the command line value takes precedence.

For example, with the following configuration file and command:

struct: "NaCl.cif"
properties:
  - "energy"
out: "NaCl-results.extxyz"
arch: mace_mp
model-path: medium
calc-kwargs:
  dispersion: True

janus singlepoint --arch mace_mp --struct KCl.cif --out KCl-results.cif --config config.yml

This will run a singlepoint energy calculation on KCl.cif using the MACE-MP "medium" force-field, saving the results to KCl-results.cif.

Note

properties must be passed as a Yaml list, as above, not as a string.

Minimal and full example configuration files for all calculations can be found here.

Docker/Podman images

You can use janus_core in a JupyterHub or marimo environment using docker or podman. We provide regularly updated docker/podman images, which can be dowloaded by running:

docker pull ghcr.io/stfc/janus-core/jupyter:amd64-latest

docker pull ghcr.io/stfc/janus-core/marimo:amd64-latest

or using podman

podman pull ghcr.io/stfc/janus-core/jupyter-amd64:latest

podman pull ghcr.io/stfc/janus-core/marimo-amd64:latest

for amd64 architecture, if you require arm64 replace amd64 with arm64 above, and next instructions.

To start, for marimo run:

podman run --rm --security-opt seccomp=unconfined -p 8842:8842 ghcr.io/stfc/janus-core/marimo:amd64-latest

or for JupyterHub, run:

podman run --rm --security-opt seccomp=unconfined -p 8888:8888 ghcr.io/stfc/janus-core/jupyter:amd64-latest

For more details on how to share your filesystem and so on you can refer to this documentation: https://summer.ccp5.ac.uk/introduction.html#run-locally.

Development

We recommend installing uv for dependency management when developing for janus-core:

1. Install uv
2. Install janus-core with dependencies in a virtual environment:

git clone https://github.com/stfc/janus-core
cd janus-core
uv sync --extra all # Create a virtual environment and install dependencies
source .venv/bin/activate
pre-commit install  # Install pre-commit hooks
pytest -v  # Discover and run all tests

License

BSD 3-Clause License

Funding

Contributors to this project were funded by