from __future__ import annotations
from pathlib import Path
from typing import Iterable, Literal, Sequence
from scm.pisa.block import DriverBlock,EngineBlock,FixedBlock,FreeBlock,InputBlock,VerbatimBlock
from scm.pisa.key import BoolKey,FloatKey,FloatListKey,IntKey,IntListKey,MultipleChoiceKey,PathStringKey,StringKey,BoolType
[docs]class MLPotential(EngineBlock):
r"""
:ivar Backend: The machine learning potential backend.
:vartype Backend: Literal["FAIRChem", "M3GNet", "MACE", "NequIP", "TorchANI"]
:ivar Device: Device on which to run the calculation (e.g. cpu, cuda:0).
If empty, the device can be controlled using environment variables for TensorFlow or PyTorch.
:vartype Device: Literal["", "cpu", "cuda:0", "cuda:1", "mps"]
:ivar MLDistanceUnit: Unit of distances expected by the ML backend (not the ASE calculator). The ASE calculator may require this information.
:vartype MLDistanceUnit: Literal["Auto", "angstrom", "bohr"]
:ivar MLEnergyUnit: Unit of energy output by the ML backend (not the unit output by the ASE calculator). The ASE calculator may require this information.
:vartype MLEnergyUnit: Literal["Auto", "Hartree", "eV", "kcal/mol", "kJ/mol"]
:ivar Model: Select a pre-parameterized or custom model.
AIMNet2-(wB97MD3/B973c): best for fast calculations of small, drug-like molecules; limited to aperiodic systems of 14 elements (H, B, C, N, O, F, Si, P, S, Cl, As, Se, Br, I).
ANI-(1x/1ccx/2x): best for very fast calculations of organic molecules; limited to elements H, C, N, O (ANI-1x/1ccx), F, S, Cl (ANI-2x).
eSEN-S-Con-OMol: best for highly accurate calculations of diverse organic and bio-relevant molecules; not intended for calculations on periodic inorganic bulk materials.
M3GNet-UP-2022: best for fast calculations of inorganic crystalline materials; not designed for accurately modeling small organic molecules or biomolecules.
MACE-MP-0-(Small/Medium/Large): best for accurate periodic calculations of inorganic materials; size trades speed/accuracy; not designed for accurately modeling small organic molecules or biomolecules. MACE-MPA-0 has improved accuracy vs MP-0.
UMA-S-1.1 variants: best for high accuracy calculations on a broad range of systems; choose from OC20 (adsorption and surface chemistry), ODAC (adsorption in porous frameworks), OMat (inorganic materials), OMC (organic molecular crystals), OMol (molecules, biomolecules, metal complexes, electrolytes); can be computationally expensive compared to other, more targeted models.
Set Custom to choose a backend and provide your own parameters.
:vartype Model: Literal["Custom", "AIMNet2-B973c", "AIMNet2-wB97MD3", "ANI-1ccx", "ANI-1x", "ANI-2x", "eSEN-S-Con-OMol", "M3GNet-UP-2022", "MACE-MP-0-Large", "MACE-MP-0-Medium", "MACE-MP-0-Small", "MACE-MPA-0", "UMA-S-1.1-OC20", "UMA-S-1.1-ODAC", "UMA-S-1.1-OMat", "UMA-S-1.1-OMC", "UMA-S-1.1-OMol"]
:ivar NumThreads: Number of threads.
If not empty, OMP_NUM_THREADS will be set to this number; for PyTorch-engines, torch.set_num_threads() will be called.
:vartype NumThreads: str | StringKey
:ivar ParameterDir: Path to a set of parameters for the backend, if it expects to read from a directory.
:vartype ParameterDir: str | Path | StringKey
:ivar ParameterFile: Path to a set of parameters for the backend, if it expects to read from a file.
:vartype ParameterFile: str | Path | StringKey
:ivar UnpairedElectrons: The number of unpaired electrons in the system for a spin unrestricted calculation. The spin multiplicity is taken as this value plus one.
:vartype UnpairedElectrons: int | IntKey
:ivar Unrestricted: Enables spin unrestricted calculations, passing spin information to the machine learning model. Only applicable to 'UMA-S-1.1-OMol' and custom FAIRChem models.
:vartype Unrestricted: BoolType | BoolKey
:ivar FAIRChem: Options for the FAIRChem machine learning potential backend.
:vartype FAIRChem: MLPotential._FAIRChem
:ivar MACE: Options for the MACE machine learning potential backend.
:vartype MACE: MLPotential._MACE
:ivar Member: Specify the details of a member of a committee.
:vartype Member: MLPotential._Member
"""
[docs] class _FAIRChem(FixedBlock):
r"""
Options for the FAIRChem machine learning potential backend.
:ivar ModelTask: Model task to use if a custom UMA/eSEN model is supplied via a parameter file. Ignored if a specific FAIRChem model is selected.
:vartype ModelTask: Literal["None", "OC20", "OMat", "OMol", "ODAC", "OMC"]
"""
def __post_init__(self):
self.ModelTask: Literal["None", "OC20", "OMat", "OMol", "ODAC", "OMC"] = MultipleChoiceKey(name='ModelTask', comment='Model task to use if a custom UMA/eSEN model is supplied via a parameter file. Ignored if a specific FAIRChem model is selected.', default='None', choices=['None', 'OC20', 'OMat', 'OMol', 'ODAC', 'OMC'])
[docs] class _MACE(FixedBlock):
r"""
Options for the MACE machine learning potential backend.
:ivar DataType: Using ``float32`` is faster but less accurate, and generally recommended for MD. Conversely using ``float64`` is slower but more accurate, and recommended for geometry optimization.
:vartype DataType: Literal["float32", "float64"]
:ivar EnableCuEquivariance: Enable CUDA-accelerated cuEquivariance library for equivariant neural networks, if CUDA available.
:vartype EnableCuEquivariance: BoolType | BoolKey
"""
def __post_init__(self):
self.DataType: Literal["float32", "float64"] = MultipleChoiceKey(name='DataType', comment='Using ``float32`` is faster but less accurate, and generally recommended for MD. Conversely using ``float64`` is slower but more accurate, and recommended for geometry optimization.', default='float32', choices=['float32', 'float64'])
self.EnableCuEquivariance: BoolType | BoolKey = BoolKey(name='EnableCuEquivariance', comment='Enable CUDA-accelerated cuEquivariance library for equivariant neural networks, if CUDA available.', default=True)
[docs] class _Member(FixedBlock):
r"""
Specify the details of a member of a committee.
:ivar ParameterDir: Path to a set of parameters for the backend, if it expects to read from a directory.
:vartype ParameterDir: str | Path | StringKey
:ivar ParameterFile: Path to a set of parameters for the backend, if it expects to read from a file.
:vartype ParameterFile: str | Path | StringKey
"""
def __post_init__(self):
self.ParameterDir: str | Path | StringKey = PathStringKey(name='ParameterDir', comment='Path to a set of parameters for the backend, if it expects to read from a directory.', hidden=True, gui_name='Parameter directory:', default='', ispath=True, gui_type='directory')
self.ParameterFile: str | Path | StringKey = PathStringKey(name='ParameterFile', comment='Path to a set of parameters for the backend, if it expects to read from a file.', hidden=True, default='', ispath=True)
def __post_init__(self):
self.Backend: Literal["FAIRChem", "M3GNet", "MACE", "NequIP", "TorchANI"] = MultipleChoiceKey(name='Backend', comment='The machine learning potential backend.', choices=['FAIRChem', 'M3GNet', 'MACE', 'NequIP', 'TorchANI'], gui_type='literal choices')
self.Device: Literal["", "cpu", "cuda:0", "cuda:1", "mps"] = MultipleChoiceKey(name='Device', comment='Device on which to run the calculation (e.g. cpu, cuda:0).\n\nIf empty, the device can be controlled using environment variables for TensorFlow or PyTorch.', default='', choices=['', 'cpu', 'cuda:0', 'cuda:1', 'mps'], hiddenchoices=['mps'])
self.MLDistanceUnit: Literal["Auto", "angstrom", "bohr"] = MultipleChoiceKey(name='MLDistanceUnit', comment='Unit of distances expected by the ML backend (not the ASE calculator). The ASE calculator may require this information.', gui_name='Internal distance unit:', default='Auto', choices=['Auto', 'angstrom', 'bohr'])
self.MLEnergyUnit: Literal["Auto", "Hartree", "eV", "kcal/mol", "kJ/mol"] = MultipleChoiceKey(name='MLEnergyUnit', comment='Unit of energy output by the ML backend (not the unit output by the ASE calculator). The ASE calculator may require this information.', gui_name='Internal energy unit:', default='Auto', choices=['Auto', 'Hartree', 'eV', 'kcal/mol', 'kJ/mol'], gui_type='literal choices')
self.Model: Literal["Custom", "AIMNet2-B973c", "AIMNet2-wB97MD3", "ANI-1ccx", "ANI-1x", "ANI-2x", "eSEN-S-Con-OMol", "M3GNet-UP-2022", "MACE-MP-0-Large", "MACE-MP-0-Medium", "MACE-MP-0-Small", "MACE-MPA-0", "UMA-S-1.1-OC20", "UMA-S-1.1-ODAC", "UMA-S-1.1-OMat", "UMA-S-1.1-OMC", "UMA-S-1.1-OMol"] = MultipleChoiceKey(name='Model', comment='Select a pre-parameterized or custom model.\n\nAIMNet2-(wB97MD3/B973c): best for fast calculations of small, drug-like molecules; limited to aperiodic systems of 14 elements (H, B, C, N, O, F, Si, P, S, Cl, As, Se, Br, I).\n\nANI-(1x/1ccx/2x): best for very fast calculations of organic molecules; limited to elements H, C, N, O (ANI-1x/1ccx), F, S, Cl (ANI-2x).\n\neSEN-S-Con-OMol: best for highly accurate calculations of diverse organic and bio-relevant molecules; not intended for calculations on periodic inorganic bulk materials.\n\nM3GNet-UP-2022: best for fast calculations of inorganic crystalline materials; not designed for accurately modeling small organic molecules or biomolecules.\n\nMACE-MP-0-(Small/Medium/Large): best for accurate periodic calculations of inorganic materials; size trades speed/accuracy; not designed for accurately modeling small organic molecules or biomolecules. MACE-MPA-0 has improved accuracy vs MP-0.\n\nUMA-S-1.1 variants: best for high accuracy calculations on a broad range of systems; choose from OC20 (adsorption and surface chemistry), ODAC (adsorption in porous frameworks), OMat (inorganic materials), OMC (organic molecular crystals), OMol (molecules, biomolecules, metal complexes, electrolytes); can be computationally expensive compared to other, more targeted models.\n\nSet Custom to choose a backend and provide your own parameters.', default='ANI-2x', choices=['Custom', 'AIMNet2-B973c', 'AIMNet2-wB97MD3', 'ANI-1ccx', 'ANI-1x', 'ANI-2x', 'eSEN-S-Con-OMol', 'M3GNet-UP-2022', 'MACE-MP-0-Large', 'MACE-MP-0-Medium', 'MACE-MP-0-Small', 'MACE-MPA-0', 'UMA-S-1.1-OC20', 'UMA-S-1.1-ODAC', 'UMA-S-1.1-OMat', 'UMA-S-1.1-OMC', 'UMA-S-1.1-OMol'], gui_type='literal choices; submenus AIMNet2 AIMNet2-* ANI ANI-* eSEN eSEN-* M3GNet M3GNet-* MACE MACE-* UMA UMA-*')
self.NumThreads: str | StringKey = StringKey(name='NumThreads', comment='Number of threads.\n\nIf not empty, OMP_NUM_THREADS will be set to this number; for PyTorch-engines, torch.set_num_threads() will be called.', gui_name='Number of threads:', default='')
self.ParameterDir: str | Path | StringKey = PathStringKey(name='ParameterDir', comment='Path to a set of parameters for the backend, if it expects to read from a directory.', gui_name='Parameter directory:', default='', ispath=True, gui_type='directory')
self.ParameterFile: str | Path | StringKey = PathStringKey(name='ParameterFile', comment='Path to a set of parameters for the backend, if it expects to read from a file.', default='', ispath=True)
self.UnpairedElectrons: int | IntKey = IntKey(name='UnpairedElectrons', comment='The number of unpaired electrons in the system for a spin unrestricted calculation. The spin multiplicity is taken as this value plus one.', gui_name='Spin polarization:', default=0, range='value >= 0')
self.Unrestricted: BoolType | BoolKey = BoolKey(name='Unrestricted', comment="Enables spin unrestricted calculations, passing spin information to the machine learning model. Only applicable to 'UMA-S-1.1-OMol' and custom FAIRChem models.", default=False)
self.FAIRChem: MLPotential._FAIRChem = self._FAIRChem(name='FAIRChem', comment='Options for the FAIRChem machine learning potential backend.')
self.MACE: MLPotential._MACE = self._MACE(name='MACE', comment='Options for the MACE machine learning potential backend.')
self.Member: MLPotential._Member = self._Member(name='Member', comment='Specify the details of a member of a committee.', hidden=True, unique=False)