TorchXPotential

class TorchXPotential(dtype=None, device=None, file=None, mapping=None, enforceLTX=None)

Constructor of the potential.

Parameters:

  • dtype (str) – The dtype of the involved torch tensors [TorchXPotential.float32, TorchXPotential.float64]. Default is TorchXPotential.float32.

  • device (str) – The device to evaluate the torchscript model. E.g.: cpu, cuda, cuda:0, cuda:1. Default is to choose available cuda devices.

  • file (str) – The name of the file that contains the torchscript model parameters.

  • mapping (dict with (ParticleIdentifier : Number)) – A mapping from ParticleIdentifiers to TorchXModel Particle Numbers

  • enforceLTX (bool) – Enforce to use also in sequential case the LTX version if available LTX versions support local stress calculation. Note that, in parallel MPI decomposition mode LTX (if available) is used be default anyway.

classmethod getAllParameterNames()

Return the names of all used parameters as a list.

getAllParameters()

Return all parameters of this potential and their current values as a <parameterName / parameterValue> dictionary.

static getDefaults()

Get the default parameters of this potential and return them in form of a dictionary of <parameter name, default value> key-value pairs.

getParameter(parameterName)

Get the current value of the parameter parameterName.

setParameter(parameterName, value)

Set the parameter parameterName to the given value.

Parameters:

  • parameterName (str) – The name of the parameter that will be modified.

  • value – The new value that will be assigned to the parameter parameterName.

Usage Examples

Set up a TorchXPotential and add it to the TremoloXPotentialSet. In this example the dtype is switched from float32 to float64.

potential_set = TremoloXPotential_set(name='TorchX-example')

# Add particle types for the needed elements, in this case Si and H.
potential_set.addParticleType(
    ParticleType(symbol='H')
)
potential_set.addParticleType(
    ParticleType(symbol='Si')
)

# Add the TorchXPotential with the torch-script file trained via QuantumATK.
_potential = TorchXPotential(
    file='MACE_training_example.qatkpt',
    dtype='float64',
)
potential_set.addPotential(_potential)
calculator = TremoloXCalculator(parameters=potential_set)

Set up a pre-trained MACE model and switch the dtype to float64. Enforce LTX-mode (even on single GPU) to enable calculation of local stress.

# Set up a pre-trained MACE potential with float64 precision.
# Enable LTX-mode to calculate local stress.
potential_set = TorchX_MACE_MP_0_L0_2023(dtype='float64', enforceLTX=True)
calculator = TremoloXCalculator(parameters=potential_set)

bulk_configuration.setCalculator(calculator)

# Calculate and print the local stress with the MACE ML-FF.
local_stress = LocalStress(bulk_configuration)
nlprint(local_stress)

Notes

The TorchXPotential class can be used to include torch-based machine-learned FF models, such as MACE [1]. The primary use case is to run simulations with user trained ML-FF models, after training the model using the MachineLearnedForceFieldTrainer. The .qatkpt file which results from the training workflow can be loaded in the TorchXPotential, either in script or using the MachineLearnedForceField block in the Workflow Builder. The path to the torch-script should be passed via the file argument, when constructing the TorchXPotential object.

For external user-trained models, currently only models based on the MACE architecture are supported. A given model can be converted using the function convertMACEModelToQATKFormat.

The dtype keyword can be used to switch the floating point precision that is used for calculating the energy, forces, and stress. The default float32 value is normally sufficient for MD simulations, as it provides a better performance. However, for calculations that require a higher accuracy, e.g. DynamicalMatrix, OptimizeNudgedElasticBand, etc. it is recommended to use float64 instead.

The enforceLTX keyword can be used to manually switch between standard serial mode (enforceLTX=False) and multi-GPU mode (enforceLTX=True), which enables MPI communication of node features between message-passing layers of the ML-FF (currently supported only for MACE and MatterSim models). By default, the modes are selected automatically, depending on whether the calculation is run with GPU and multiple MPI processes or not. One can use this flag to enforce the multi-GPU mode even when running serial by setting enforceLTX=True. This enables calculation of LocalStress, which is otherwise not supported for TorchXPotential.

The device keyword can be used to manually disable or enable GPU acceleration. However, the recommended approach is to leave this parameter at the default value None, and enable the GPU acceleration as described in the technical notes.

The mapping keyword can be used to specify a dict between atomic symbols and atomic numbers used to encode the elements in most ML-FF backends. The mapping is not needed for models trained with QuantumATK, or converted using convertMACEModelToQATKFormat.