compose.mace_jax

Loads a MACE model from PyTorch via MACE-JAX.

Foundation Models

load_foundational_model(family='mp', version='medium-0b3')

Loads a foundational MACE model from PyTorch.

Parameters:

• family (str) – Model family to load, e.g., “mp” or “off”.

• version (str) – Model version to load, e.g., “medium-0b3”.

Returns:

A tuple of the loaded PyTorch model and its configuration dictionary.

load_torch_model(model_file)

Loads a foundational MACE model from PyTorch.

Parameters:

model_file (str) – Filename / path of PyTorch model to load.

Returns:

A tuple of the loaded PyTorch model and its configuration dictionary.

mace_jax_neighborlist_from_torch(config, torch_model, displacement, max_edge_multiplier=1.25, per_particle=False, scale_pos=0.1, scale_pot=96.485, species_mapping=<chemtrain.compose.mace_jax.SpeciesMapping object>, equivariance_config=None, use_custom_batch_fn=False, comm=None, head=None)

Convert a PyTorch MACE model for neighbor-list property prediction.

Parameters:

• config – Configuration dictionary for the MACE model.

• torch_model – The PyTorch MACE model to convert.

• displacement – jax-md displacement function.

• max_edge_multiplier – Multiplier to limit the maximum number of edges per particle.

• per_particle – Return per-particle energies instead of total energy.

• scale_pos – Scaling factor for positions, i.e., to convert units.

• scale_pot – Scaling factor for potentials, i.e., to convert units.

• species_mapping – Mapping for species to model-compatible indices.

• equivariance_config – Backend-neutral equivariance configuration.

• use_custom_batch_fn – Use chemtrain’s custom batching transform even when the selected backend does not require it.

• comm – Optional deployment communication interface. It must provide a gather method used between message-passing blocks.

• head – Head name or index for a multi-head model. The first head is used when no head is given.

Returns:

Converted variables and a JIT-compiled apply function.

Species Mappings

class SpeciesMapping

Identity mapping for species.

class AtomicNumberMapping(max_number=100)

Maps atomic numbers to MACE-JAX species.

Build a Model

mace_jax_neighborlist(*, displacement, r_cutoff, n_species=100, positions_test=None, neighbor_test=None, max_edge_multiplier=1.25, edges_per_particle=None, avg_num_neighbors=None, mode='energy', per_particle=False, equivariance_config=None, use_custom_batch_fn=False, mace_config=None)

Initialize a MACE-JAX model for neighbor-list property prediction.

The initialization follows the chemutils implementation and converts the resulting model to the MACE-JAX representation.

Parameters:

• displacement – jax-md displacement function.

• r_cutoff – Radial cutoff distance for the model and neighbor list.

• n_species – Number of different atom species the network is supposed to process.

• positions_test – Sample positions to estimate max_edges / max_angles. Needs to be provided to enable capping.

• neighbor_test – Sample neighborlist to estimate max_edges / max_angles. Needs to be provided to enable capping.

• max_edge_multiplier – Multiplier for initial estimate of maximum edges.

• avg_num_neighbors – Average number of neighbors per particle. Guessed if positions_test and neighbor_test are provided.

• mode – Prediction mode of the model.

• per_particle – Return per-particle energies instead of total energy.

• equivariance_config – Backend-neutral equivariance configuration.

• use_custom_batch_fn – Use chemtrain’s custom batching transform even when the selected backend does not require it.

• mace_config – Kwargs to change the default structure of MACE. For definition of the kwargs, see MACE.

Returns:

Initialized variables, a JIT-compiled apply function, and the model configuration.