convertMACEModelToQATKFormat

Included in QATK.MLFF

convertMACEModelToQATKFormat(model_path, conversion_format='e3nn', model_head_number=-1)

Convert a open-source trained MACE model to a QATK compatible format.

Parameters:

  • model_path (str | Path) – The full path to a MACE model file.

  • conversion_format (str) – The format to convert the model to. Possible values are MLParameterOptions.MODEL_CONVERSION options (MLParameterOptions.MODEL_CONVERSION.E3NN, MLParameterOptions.MODEL_CONVERSION.CUE_FP32, or MLParameterOptions.MODEL_CONVERSION.CUE_FP64) or strings. For cuequivariance formats, the datatype is integrated in the saved model and cannot be set at initialization of a TorchXPotential. For the E3NN format, the datatype can be chosen when initializing the TorchXPotential.
    Default: MLParameterOptions.MODEL_CONVERSION.E3NN

  • model_head_number (int) – The head number of the model to use for the QATK compatible model.
    Default: The last head of the input model.

Usage Examples

MACE models set up via MACEFittingParameters and trained via MachineLearnedForceFieldTrainer are automatically converted into QATK-compatible formats. For MACE models acquired from elsewhere, it is possible to convert them into a TremoloX-native format by using this method. The conversion is required in order to use MACE models with QuantumATK’s native force field framework.

  1. Converting to standard e3nn format

Assume that we have a MACE model file named mace_model.model with only a single model head in our current working directory. Using the following script, the MACE model can be converted into a QATK-compatible format with a .qatkpt extension.

# Convert the MACE model to standard e3nn QATK format
convertMACEModelToQATKFormat("mace_model.model")

After the conversion has saved the model in a QuantumATK compatible format, the MACE model can be loaded as a force field potential and used to set up a TremoloXCalculator that can be used as usual in QuantumATK.

# Load the interatomic potential into a TremoloXCalculator that can be
# attached to and used with QATK configurations...
model_path = "mace_model.qatkpt"

potentialSet = TremoloXPotentialSet(name="my_mace_potential")
potential = TorchXPotential(
    dtype="float64",
    file=model_path,
)

for symbol in potential.get_symbols(model_path):
    potentialSet.addParticleType(ParticleType(symbol))
potentialSet.addPotential(potential)

calculator = TremoloXCalculator(parameters=potentialSet)

  1. Converting to CuEquivariance-accelerated formats

MACE models can also be converted to CuEquivariance-accelerated formats, which provide optimized GPU performance. These formats are available in both floating-point 32 and floating-point 64 precision. Note that the precision is fixed at conversion time and cannot be changed afterward.

# Alternatively, convert to CuEquivariance-accelerated formats for optimized GPU performance
# Convert to fp32 CuEquivariance format
convertMACEModelToQATKFormat(
    "mace_model.model", conversion_format=MLParameterOptions.MODEL_CONVERSION.CUE_FP32
)
# Set up calculator similarly using 'mace_model_cue_fp32.qatkpt' file
# ...

# Convert to fp64 CuEquivariance format
convertMACEModelToQATKFormat(
    "mace_model.model", conversion_format=MLParameterOptions.MODEL_CONVERSION.CUE_FP64
)
# Set up calculator similarly using 'mace_model_cue_fp64.qatkpt' file

The example scripts are available for download: conversion_example.py

Notes

  • The conversion utility is available for .model files containing one or more model heads generated with the MACE framework [1][2].

  • Three conversion formats are available via the conversion_format parameter (MLParameterOptions.MODEL_CONVERSION constants or string equivalents - MLParameterOptions):

    • MLParameterOptions.MODEL_CONVERSION.E3NN (or 'e3nn', default): Converts to the standard e3nn format with .qatkpt extension. This format is flexible and can be used with different floating-point precisions in a TorchXPotential when a TremoloXCalculator is set up.

    • MLParameterOptions.MODEL_CONVERSION.CUE_FP32 (or 'cue_fp32'): Converts to CuEquivariance-accelerated format with floating-point 32 precision. The output file will have a _cue_fp32.qatkpt extension. This format provides optimized GPU performance but the precision is fixed and cannot be changed after conversion.

    • MLParameterOptions.MODEL_CONVERSION.CUE_FP64 (or 'cue_fp64'): Converts to CuEquivariance-accelerated format with floating-point 64 precision. The output file will have a _cue_fp64.qatkpt extension. This format provides optimized GPU performance with higher precision, but lower performance than the fp32 format. The precision is fixed and cannot be changed after conversion.

  • The output model file will have the same base name as the input file path but with the extension replaced according to the chosen conversion format.

  • It is only possible to convert one head from a .model file into a .qatkpt file at a time. If several heads from a .model file are required, the conversion needs to be repeated with different model_head_number values.

  • CuEquivariance-accelerated formats (cue_fp32 and cue_fp64) are optimized for GPU inference on Linux systems and provide significant performance improvements compared to the standard e3nn format. However, once converted to a specific precision, the model cannot be used with a different precision without reconverting from the original .model file.