MACEModels.jl

MACE interface with support for ensemble of models and batch size selection for potentially faster inference.

DOFs currently hardcoded at 3.

Fields

• model_paths::Vector{String}: Model paths to load. If multiple paths are given, an ensemble of models is used and mean energies/forces used to propagate dynamics.
• models::Vector: Loaded models.
• device::Vector{String}: Device to use for inference. If a single device is given, all models will be loaded on that device. If multiple devices are given, each model will be loaded on the corresponding device. Default is "cpu".
• default_dtype::Type{T}: Default data type for inference. Default is Float32.
• batch_size::Union{Int,Nothing}: Batch size for inference. Default is 1.
• cutoff_radius::T: Cutoff radius for the models.
• last_eval_cache::MACEPredictionCache: Cache for the last evaluation.
• z_table: Table of atomic numbers.
• atoms::Atoms: Atoms object for the system.
• cell::PeriodicCell: Periodic cell for the system.
• ndofs::Int: Number of degrees of freedom in the system.
• mobile_atoms::Vector{Int}: Indices of mobile atoms in the system.

MACEModel(
atoms::Atoms,
cell::AbstractCell,
model_paths::Vector{String};
device::Union{String,Vector{String}}="cpu",
default_dtype::Type=Float32,
batch_size::Int=1,
mobile_atoms::Vector{Int}=collect(1:length(atoms)),

) modelpaths::Vector{String}, device::Union{String, Vector{String}}="cpu", defaultdtype::Type=Float64, batch_size::Int=1, atoms, cell, )

Interface to MACE machine learning interatomic potentials with support for ensemble of models and batch size selection for potentially faster inference.

Arguments

• atoms: Atoms object for the system. predict! can be called with a Vector of different Atoms objects to evaluate different structures.
• cell: Cell object for the system. predict! can be called with a Vector of different Cell objects to evaluate different structures.
• model_paths::Vector{String}: Model paths to load. If multiple paths are given, an ensemble of models is used and mean energies/forces used to propagate dynamics.
• device::Union{String, Vector{String}}: Device to use for inference. If a single device is given, all models will be loaded on that device. If multiple devices are given, each model will be loaded on the corresponding device. Default is "cpu".
• default_dtype::Type: Default data type for PyTorch (usually Float32)
• batch_size::Int: Batch size for inference. Can be set to nothing to always adapt batch size to the number of structures provided. Ring-polymer methods benefit from this.
• mobile_atoms::Vector{Int}: Indices of mobile atoms in the system. Default is all atoms.

Cache which stores the results of the last evaluation of the MACE model. This is used to speed up concurrent energy and force evaluations, as well as to give access to ensemble predictions.

Use e.g. get_energy_mean(MACEModel.last_eval_cache) to access results of the last prediction made by the model.

get_energy_ensemble(mace_cache::MACEPredictionCache)

Returns the potential energy evaluated by each model in the ensemble in units of Hartree.

get_energy_mean(mace_cache::MACEPredictionCache)

Returns the mean potential energy of the structures stored in the evaluation cache.

Energy is returned in units of Hartree.

get_energy_variance(mace_cache::MACEPredictionCache)

Returns the variance of the potential energy in Hartree².

get_forces_ensemble(mace_cache::MACEPredictionCache)

Returns the forces evaluated by each model in the ensemble in units of Hartree/Bohr.

get_forces_mean(mace_cache::MACEPredictionCache)

Returns the mean forces of the structures stored in the evaluation cache. Forces are returned in units of Hartree/Bohr.

get_forces_variance(mace_cache::MACEPredictionCache)

Returns the force variance of the structures stored in the evaluation cache. Forces are returned in units of Hartree²/Bohr².

mace_configuration_from_nqcd_configuration(
atoms::Atoms,
cell::Union{InfiniteCell, PeriodicCell},
R::AbstractMatrix,

)

Converter into a single mace.data.utils.Configuration to make use of MACE's data loading.

predict!(
mace_interface::MACEModel,
atoms::Union{Vector{<:Atoms}, Atoms},
R::Vector{<:AbstractMatrix},
cell::Union{Vector{<:AbstractCell}, AbstractCell},
)

Evaluate the MACE model on a set of structures if they haven't already been evaluated. The results are stored in mace_interface.last_eval_cache.

Use methodswith(MACEPredictionCache) to see what data is provided from the results

Arguments

mace_interface: MACEModel to evaluate with.

atoms: Atoms object for the structures passed for evaluation. Can be either a vector of different Atoms objects or a single Atoms object.

R: Vector of structures to evaluate.

cell: Cell object for the structures passed for evaluation. Can be either a vector of different Cell objects or a single Cell object.

predict(
mace_interface::MACEModel,
atoms::Union{Vector{<:Atoms}, Atoms},
R::Vector{<:AbstractMatrix},
cell::Union{Vector{<:AbstractCell}, AbstractCell},
)

Evaluate the MACE model on a set of structures if they haven't already been evaluated. The results are returned as a MACEPredictionCache.

Use methodswith(MACEPredictionCache) to see what data can be provided from the results.

Arguments

mace_interface: MACEModel to evaluate with.

atoms: Atoms object for the structures passed for evaluation. Can be either a vector of different Atoms objects or a single Atoms object.

R: Vector of structures to evaluate.

cell: Cell object for the structures passed for evaluation. Can be either a vector of different Cell objects or a single Cell object.

NQCModels.derivative(model::MACEModel, atoms::Atoms, R::AbstractArray{T,3}, cell::Union{InfiniteCell, PeriodicCell})

This variant of NQCModels.derivative can make use of batch evaluation to speed up inference for multiple structures.

NQCModels.derivative(model::MACEModel, atoms::Atoms, R::Vector{<:AbstractMatrix}, cell::Union{InfiniteCell, PeriodicCell})

This variant of NQCModels.derivative can make use of batch evaluation to speed up inference for multiple structures.

NQCModels.potential(model::MACEModel, atoms::Atoms, R::Vector{<:AbstractMatrix}, cell::AbstractCell)

This variant of NQCModels.potential can make use of batch evaluation to speed up inference for multiple structures.

Ensemble algorithm which enforces the split between dynamics and inference workers provided.

Even when specifying this EnsembleAlgorithm, you still need to manually dispatch

Containers for energy and forces, to be used in the RemoteModel.

MultiProcessConfig(runners, evaluators, model_load_function::Function, positions_prototype; model_listener::Function=batch_evaluation_loop)

Configuration for distributed execution dynamics propagation and model evaluation.

Arguments

• runners::Vector{Int}: Process IDs of the NQCDynamics workers.
• evaluators::Vector{Int}: Process IDs of the model evaluators.
• model_load_function::Function: Function ()->x<:Model that can load the model to be evaluated on any process. (i.e. not specific to a process)
• model_listener::Function: Function that listens for structures on the input channels and outputs the results on the output channels. Default is batch_evaluation_loop. (model, input_channels, output_channels)->()
• positions_prototype: Sample atomic positions to determine system size.

RemoteModel(config::MultiProcessConfig, structure_prototype)

Creates a RemoteModel that can be used to evaluate structures on a remote process.

batch_evaluation_loop(model<:Model, input_channels::Vector{RemoteChannel}, output_channels::Vector{RemoteChannel}; user_model_function::Function=mace_batch_predict, max_delay = 1000)

This doesn't know how to use your model most efficiently, it only handles I/O between processes!

To use this communication wrapper, you need to define a function that evaluates your model and returns a tuple of energies and forces, as shown in the example:

function user_model_function(model<:Model, structures::AbstractVector{AbstractMatrix{Number}})
    MACEModels.predict!(model, model.atoms, structures, model.cell)
    energies = MACEModels.get_energy_mean(model.last_eval_cache) # This is a Vector{Number} type
    forces = MACEModels.get_forces_mean(model.last_eval_cache) # This is a Vector{Matrix{Number}} type
    return (energies, forces)
end

It polls the input channels for new structures, and when it finds them, it evaluates them using the user-defined model function. The results are then put back into the correct output channels.

Arguments

• model::Model: The model that will be used for evaluation.
• input_channels::Vector{RemoteChannel}: The input channels that will be used to receive structures.
• output_channels::Vector{RemoteChannel}: The output channels that will be used to send the results.
• user_model_function::Function: The user-defined function that will be used to evaluate the model.