SchnetPackCalculator

Calculator subclass to accomodate machine learning models trained using the SchnetPack package.

class mlcalcdriver.calculators.schnetpack.SchnetPackCalculator(model_dir, available_properties=None, device='cpu', units={'dipole_moment': 'Debye', 'energy': 'eV', 'positions': 'angstrom'}, md=False, dropout=False)[source]

Bases: mlcalcdriver.calculators.calculator.Calculator

Calculator based on a SchnetPack model

Parameters

  • model_dir (str) – Path to the stored model on which the calculator will be based. If $MODELDIR is defined, the path can be relative to it. If not, the path must be absolute or relative to the working directory.

  • available_properties (str or list of str) – Properties that the model can predict. If None, they automatically determined from the model. Default is None.

  • device (str) – Can be either “cpu” to use cpu or “cuda” to use “gpu”

  • units (dict) – Dictionnary containing the units in which the calculator makes predictions. Default is mlcalcdriver.globals.eVA for a SchnetPackCalculator.

  • md (bool) – Whether the calculator is used with ASE to do molecular dynamics. Default is False and should be changed through the AseSpkCalculator object.

  • dropout (bool) – Whether the calculator should use the dropout layers to estimate a confidence interval for the prediction. Default is False. No effect if the model hasn’t been trained with dropout layers.

run(property, posinp=None, batch_size=128)[source]

Central method to use when making a calculation with the calculator.

Parameters

  • property (str) – Property to be predicted by the calculator

  • posinp (Posinp) – Atomic configuration to pass to the model

  • batch_size (int) – Batch sizes. Default is 128.

Returns

predictions – Corresponding prediction by the model.

Return type

numpy.ndarray