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 ParAMSMachineLearning(DriverBlock):
r"""
:ivar EngineCollection: Path to (optional) JobCollection Engines YAML file.
:vartype EngineCollection: str | StringKey
:ivar JobCollection: Path to JobCollection YAML file.
:vartype JobCollection: str | StringKey
:ivar ResultsDirectory: Directory in which output files will be created.
:vartype ResultsDirectory: str | Path | StringKey
:ivar Task: Task to run.
Available options:
•MachineLearning: Optimization for machine learning models.
•Optimization: Global optimization powered by GloMPO
•Generate Reference: Run jobs with reference engine to get reference values
•Single Point: Evaluate the current configuration of jobs, training data, and parameters
•Sensitivity: Measure the sensitivity of the loss function to each of the active parameters
:vartype Task: Literal["Optimization", "GenerateReference", "SinglePoint", "Sensitivity", "MachineLearning"]
:ivar DataSet: Configuration settings for each data set in the optimization.
:vartype DataSet: ParAMSMachineLearning._DataSet
:ivar MachineLearning: Options for Task MachineLearning.
:vartype MachineLearning: ParAMSMachineLearning._MachineLearning
:ivar ParallelLevels: Distribution of threads/processes between the parallelization levels.
:vartype ParallelLevels: ParAMSMachineLearning._ParallelLevels
"""
[docs] class _DataSet(FixedBlock):
r"""
Configuration settings for each data set in the optimization.
:ivar BatchSize: Number of data set entries to be evaluated per epoch. Default 0 means all entries.
:vartype BatchSize: int | IntKey
:ivar EvaluateEvery: This data set is evaluated every n evaluations of the training set.
This will always be set to 1 for the training set. For other data sets it will be adjusted to the closest multiple of LoggingInterval%General, i.e., you cannot evaluate an extra data set more frequently than you log it.
:vartype EvaluateEvery: int | IntKey
:ivar LossFunction: Loss function used to quantify the error between model and reference values. This becomes the minimization task.
Available options:
• mae: Mean absolute error
• rmse: Root mean squared error
• sse: Sum of squared errors
• sae: Sum of absolute errors
:vartype LossFunction: Literal["mae", "rmse", "sse", "sae"]
:ivar MaxJobs: Limit each evaluation to a subset of n jobs. Default 0 meaning all jobs are used.
:vartype MaxJobs: int | IntKey
:ivar MaxJobsShuffle: Use a different job subset every for every evaluation.
:vartype MaxJobsShuffle: BoolType | BoolKey
:ivar Name: Unique data set identifier.
The first occurrence of DataSet will always be called training_set.
The second will always be called validation_set.
These cannot be overwritten.
Later occurrences will default to data_set_xx where xx starts at 03 and increments from there. This field can be used to customize the latter names.
:vartype Name: str | StringKey
:ivar Path: Path to DataSet YAML file.
:vartype Path: str | StringKey
:ivar UsePipe: Use AMS Pipe for suitable jobs to speed-up evaluation.
:vartype UsePipe: BoolType | BoolKey
"""
def __post_init__(self):
self.BatchSize: int | IntKey = IntKey(name='BatchSize', comment='Number of data set entries to be evaluated per epoch. Default 0 means all entries.', default=0)
self.EvaluateEvery: int | IntKey = IntKey(name='EvaluateEvery', comment='This data set is evaluated every n evaluations of the training set.\n\nThis will always be set to 1 for the training set. For other data sets it will be adjusted to the closest multiple of LoggingInterval%General, i.e., you cannot evaluate an extra data set more frequently than you log it.', default=1)
self.LossFunction: Literal["mae", "rmse", "sse", "sae"] = MultipleChoiceKey(name='LossFunction', comment='Loss function used to quantify the error between model and reference values. This becomes the minimization task.\n\nAvailable options:\n• mae: Mean absolute error\n• rmse: Root mean squared error\n• sse: Sum of squared errors\n• sae: Sum of absolute errors', default='sse', choices=['mae', 'rmse', 'sse', 'sae'])
self.MaxJobs: int | IntKey = IntKey(name='MaxJobs', comment='Limit each evaluation to a subset of n jobs. Default 0 meaning all jobs are used.', default=0)
self.MaxJobsShuffle: BoolType | BoolKey = BoolKey(name='MaxJobsShuffle', comment='Use a different job subset every for every evaluation.', default=False)
self.Name: str | StringKey = StringKey(name='Name', comment='Unique data set identifier.\n\nThe first occurrence of DataSet will always be called training_set.\nThe second will always be called validation_set.\nThese cannot be overwritten.\n\nLater occurrences will default to data_set_xx where xx starts at 03 and increments from there. This field can be used to customize the latter names.', default='')
self.Path: str | StringKey = StringKey(name='Path', comment='Path to DataSet YAML file.')
self.UsePipe: BoolType | BoolKey = BoolKey(name='UsePipe', comment='Use AMS Pipe for suitable jobs to speed-up evaluation.', default=True)
[docs] class _MachineLearning(FixedBlock):
r"""
Options for Task MachineLearning.
:ivar Backend: The backend to use. You must separately install the backend before running a training job.
:vartype Backend: Literal["Custom", "MatGL", "M3GNet", "MACE", "NEP", "NequIP", "Test"]
:ivar CommitteeSize: The number of independently trained ML potentials.
:vartype CommitteeSize: int | IntKey
:ivar LoadModel: Load a previously fitted model from a ParAMS results directory. A ParAMS results directory should contain two subdirectories ``optimization`` and ``settings_and_initial_data``. The loaded model defines the model source and architecture; fitting controls such as the learning rate and trainable layers still apply when supported by the backend.
:vartype LoadModel: str | Path | StringKey
:ivar MaxEpochs: Maximum number of epochs during the training.
:vartype MaxEpochs: int | IntKey
:ivar RunAMSAtEnd: Whether to run the (committee) ML potential through AMS at the end. This will create the energy/forces scatter plots for the final trained model.
:vartype RunAMSAtEnd: BoolType | BoolKey
:ivar Custom: Set up a custom fitting program within ParAMS
:vartype Custom: ParAMSMachineLearning._MachineLearning._Custom
:ivar EarlyStopping: Stop training when the validation loss has not improved sufficiently. Supported by the MatGL and M3GNet backends.
:vartype EarlyStopping: ParAMSMachineLearning._MachineLearning._EarlyStopping
:ivar LossCoeffs: Modify the coefficients for the machine learning loss function. For backends that support weights, this is on top of the supplied dataset weights and sigmas.
:vartype LossCoeffs: ParAMSMachineLearning._MachineLearning._LossCoeffs
:ivar M3GNet: Options for M3GNet fitting.
:vartype M3GNet: ParAMSMachineLearning._MachineLearning._M3GNet
:ivar MACE: Options for MACE fitting.
:vartype MACE: ParAMSMachineLearning._MachineLearning._MACE
:ivar MatGL: Options for fitting MatGL QET and TensorNet potentials.
:vartype MatGL: ParAMSMachineLearning._MachineLearning._MatGL
:ivar NEP: Experimental GPUMD NEP4 fitting. Training uses AMS_NEP_TRAIN_EXECUTABLE; production calculations use the native NEP worker installed by AMS. The default Foundation model fine-tunes the installed NEP89 potential for neutral, fully 3D-periodic systems.
:vartype NEP: ParAMSMachineLearning._MachineLearning._NEP
:ivar NequIP: Options for NequIP fitting.
:vartype NequIP: ParAMSMachineLearning._MachineLearning._NequIP
:ivar Target: Target values for stopping training. If both the training and validation metrics are smaller than the specified values, the training will stop early. Supported by the MatGL and M3GNet backends.
:vartype Target: ParAMSMachineLearning._MachineLearning._Target
"""
[docs] class _Custom(FixedBlock):
r"""
Set up a custom fitting program within ParAMS
:ivar File: Python file containing a function called 'get_fit_job' that returns a subclass of 'FitJob'
:vartype File: str | Path | StringKey
:ivar Arguments: Pass on keyword arguments to the 'get_fit_job' function.
:vartype Arguments: str | Sequence[str] | FreeBlock
"""
[docs] class _Arguments(FreeBlock):
r"""
Pass on keyword arguments to the 'get_fit_job' function.
"""
def __post_init__(self):
pass
def __post_init__(self):
self.File: str | Path | StringKey = PathStringKey(name='File', comment="Python file containing a function called 'get_fit_job' that returns a subclass of 'FitJob'", ispath=True)
self.Arguments: str | Sequence[str] | FreeBlock = self._Arguments(name='Arguments', comment="Pass on keyword arguments to the 'get_fit_job' function.")
[docs] class _EarlyStopping(FixedBlock):
r"""
Stop training when the validation loss has not improved sufficiently. Supported by the MatGL and M3GNet backends.
:ivar Enabled: Whether to stop training after the validation loss stops improving.
:vartype Enabled: BoolType | BoolKey
:ivar MinDelta: Minimum absolute decrease in validation loss that counts as an improvement. Must not be negative.
:vartype MinDelta: float | FloatKey
:ivar Patience: Number of consecutive validation checks without a sufficient improvement before stopping. Must be greater than zero.
:vartype Patience: int | IntKey
"""
def __post_init__(self):
self.Enabled: BoolType | BoolKey = BoolKey(name='Enabled', comment='Whether to stop training after the validation loss stops improving.', default=True)
self.MinDelta: float | FloatKey = FloatKey(name='MinDelta', comment='Minimum absolute decrease in validation loss that counts as an improvement. Must not be negative.', default=1e-05, isenabled='glompo_machinelearning.machinelearning.earlystopping.enabled # hide')
self.Patience: int | IntKey = IntKey(name='Patience', comment='Number of consecutive validation checks without a sufficient improvement before stopping. Must be greater than zero.', default=100, isenabled='glompo_machinelearning.machinelearning.earlystopping.enabled # hide')
[docs] class _LossCoeffs(FixedBlock):
r"""
Modify the coefficients for the machine learning loss function. For backends that support weights, this is on top of the supplied dataset weights and sigmas.
:ivar AverageForcePerAtom: For each force data entry, divide the loss contribution by the number of concomitant atoms. This is the same as the behavior for ParAMS Optimization, but it is turned off by default in Task MachineLearning. For machine learning, setting this to 'No' can be better since larger molecules will contribute more to the loss. For backends that support weights, this is on top of the supplied dataset weights and sigmas.
:vartype AverageForcePerAtom: BoolType | BoolKey
:ivar Energy: Coefficient for the contribution of loss due to the energy. For backends that support weights, this is on top of the supplied dataset weights and sigmas.
:vartype Energy: float | FloatKey
:ivar Forces: Coefficient for the contribution of loss due to the forces. For backends that support weights, this is on top of the supplied dataset weights and sigmas.
:vartype Forces: float | FloatKey
"""
def __post_init__(self):
self.AverageForcePerAtom: BoolType | BoolKey = BoolKey(name='AverageForcePerAtom', comment="For each force data entry, divide the loss contribution by the number of concomitant atoms. This is the same as the behavior for ParAMS Optimization, but it is turned off by default in Task MachineLearning. For machine learning, setting this to 'No' can be better since larger molecules will contribute more to the loss. For backends that support weights, this is on top of the supplied dataset weights and sigmas.", default=False)
self.Energy: float | FloatKey = FloatKey(name='Energy', comment='Coefficient for the contribution of loss due to the energy. For backends that support weights, this is on top of the supplied dataset weights and sigmas.', gui_name='Energy coefficient:', default=10.0)
self.Forces: float | FloatKey = FloatKey(name='Forces', comment='Coefficient for the contribution of loss due to the forces. For backends that support weights, this is on top of the supplied dataset weights and sigmas.', gui_name='Forces coefficient:', default=1.0)
[docs] class _M3GNet(FixedBlock):
r"""
Options for M3GNet fitting.
:ivar LearningRate: Learning rate for the M3GNet weight optimization.
:vartype LearningRate: float | FloatKey
:ivar Model: How to specify the model for the M3GNet backend. Either a Custom model can be made from scratch or an existing model directory can be loaded to obtain the model settings.
:vartype Model: Literal["UniversalPotential", "Custom", "ModelDir"]
:ivar ModelDir: Path to the directory defining the model. This folder should contain the files: 'checkpoint', 'm3gnet.data-00000-of-00001', ' m3gnet.index' and 'm3gnet.json'
:vartype ModelDir: str | Path | StringKey
:ivar Custom: Specify a custom M3GNet model.
:vartype Custom: ParAMSMachineLearning._MachineLearning._M3GNet._Custom
:ivar UniversalPotential: Settings for (transfer) learning with the M3GNet Universal Potential.
:vartype UniversalPotential: ParAMSMachineLearning._MachineLearning._M3GNet._UniversalPotential
"""
[docs] class _Custom(FixedBlock):
r"""
Specify a custom M3GNet model.
:ivar Cutoff: Cutoff radius of the graph
:vartype Cutoff: float | FloatKey
:ivar MaxL: Include spherical components up to order MaxL. Higher gives a better angular resolution, but increases computational cost substantially.
:vartype MaxL: int | IntKey
:ivar MaxN: Include radial components up to the MaxN'th root of the spherical Bessel function. Higher gives a better radial resolution, but increases computational cost substantially.
:vartype MaxN: int | IntKey
:ivar NumBlocks: Number of convolution blocks.
:vartype NumBlocks: int | IntKey
:ivar NumNeurons: Number of neurons in each layer.
:vartype NumNeurons: int | IntKey
:ivar ThreebodyCutoff: Cutoff radius of the three-body interaction.
:vartype ThreebodyCutoff: float | FloatKey
"""
def __post_init__(self):
self.Cutoff: float | FloatKey = FloatKey(name='Cutoff', comment='Cutoff radius of the graph', default=5.0, unit='angstrom')
self.MaxL: int | IntKey = IntKey(name='MaxL', comment='Include spherical components up to order MaxL. Higher gives a better angular resolution, but increases computational cost substantially.', default=3)
self.MaxN: int | IntKey = IntKey(name='MaxN', comment="Include radial components up to the MaxN'th root of the spherical Bessel function. Higher gives a better radial resolution, but increases computational cost substantially.", default=3)
self.NumBlocks: int | IntKey = IntKey(name='NumBlocks', comment='Number of convolution blocks.', gui_name='Number of convolution blocks: ', default=3)
self.NumNeurons: int | IntKey = IntKey(name='NumNeurons', comment='Number of neurons in each layer.', gui_name='Number of neurons per layer:', default=64)
self.ThreebodyCutoff: float | FloatKey = FloatKey(name='ThreebodyCutoff', comment='Cutoff radius of the three-body interaction.', default=4.0, unit='angstrom')
[docs] class _UniversalPotential(FixedBlock):
r"""
Settings for (transfer) learning with the M3GNet Universal Potential.
:ivar Featurizer: Train the Featurizer layer of the M3GNet universal potential.
:vartype Featurizer: BoolType | BoolKey
:ivar Final: Train the Final layer of the M3GNet universal potential.
:vartype Final: BoolType | BoolKey
:ivar GraphLayer1: Train the first Graph layer of the M3GNet universal potential.
:vartype GraphLayer1: BoolType | BoolKey
:ivar GraphLayer2: Train the second Graph layer of the M3GNet universal potential.
:vartype GraphLayer2: BoolType | BoolKey
:ivar GraphLayer3: Train the third Graph layer of the M3GNet universal potential.
:vartype GraphLayer3: BoolType | BoolKey
:ivar ThreeDInteractions1: Train the first ThreeDInteractions (three-body terms) layer of the M3GNet universal potential.
:vartype ThreeDInteractions1: BoolType | BoolKey
:ivar ThreeDInteractions2: Train the second ThreeDInteractions (three-body terms) layer of the M3GNet universal potential.
:vartype ThreeDInteractions2: BoolType | BoolKey
:ivar ThreeDInteractions3: Train the third ThreeDInteractions (three-body terms) layer of the M3GNet universal potential.
:vartype ThreeDInteractions3: BoolType | BoolKey
:ivar Version: Which version of the M3GNet Universal Potential to use.
:vartype Version: Literal["2022"]
"""
def __post_init__(self):
self.Featurizer: BoolType | BoolKey = BoolKey(name='Featurizer', comment='Train the Featurizer layer of the M3GNet universal potential.', gui_name='Train featurizer:', default=False, isenabled="glompo_machinelearning.machinelearning.m3gnet.model == 'Universal Potential' # hide")
self.Final: BoolType | BoolKey = BoolKey(name='Final', comment='Train the Final layer of the M3GNet universal potential.', gui_name='Train final layer:', default=True, isenabled="glompo_machinelearning.machinelearning.m3gnet.model == 'Universal Potential' # hide")
self.GraphLayer1: BoolType | BoolKey = BoolKey(name='GraphLayer1', comment='Train the first Graph layer of the M3GNet universal potential.', gui_name='Train layer 1 - graph:', default=False, isenabled="glompo_machinelearning.machinelearning.m3gnet.model == 'Universal Potential' # hide")
self.GraphLayer2: BoolType | BoolKey = BoolKey(name='GraphLayer2', comment='Train the second Graph layer of the M3GNet universal potential.', gui_name='Train layer 2 - graph:', default=False, isenabled="glompo_machinelearning.machinelearning.m3gnet.model == 'Universal Potential' # hide")
self.GraphLayer3: BoolType | BoolKey = BoolKey(name='GraphLayer3', comment='Train the third Graph layer of the M3GNet universal potential.', gui_name='Train layer 3 - graph:', default=True, isenabled="glompo_machinelearning.machinelearning.m3gnet.model == 'Universal Potential' # hide")
self.ThreeDInteractions1: BoolType | BoolKey = BoolKey(name='ThreeDInteractions1', comment='Train the first ThreeDInteractions (three-body terms) layer of the M3GNet universal potential.', gui_name='Train layer 1 - 3D interactions:', default=False, isenabled="glompo_machinelearning.machinelearning.m3gnet.model == 'Universal Potential' # hide")
self.ThreeDInteractions2: BoolType | BoolKey = BoolKey(name='ThreeDInteractions2', comment='Train the second ThreeDInteractions (three-body terms) layer of the M3GNet universal potential.', gui_name='Train layer 2 - 3D interactions:', default=False, isenabled="glompo_machinelearning.machinelearning.m3gnet.model == 'Universal Potential' # hide")
self.ThreeDInteractions3: BoolType | BoolKey = BoolKey(name='ThreeDInteractions3', comment='Train the third ThreeDInteractions (three-body terms) layer of the M3GNet universal potential.', gui_name='Train layer 3 - 3D interactions:', default=True, isenabled="glompo_machinelearning.machinelearning.m3gnet.model == 'Universal Potential' # hide")
self.Version: Literal["2022"] = MultipleChoiceKey(name='Version', comment='Which version of the M3GNet Universal Potential to use.', hidden=True, default='2022', choices=['2022'])
def __post_init__(self):
self.LearningRate: float | FloatKey = FloatKey(name='LearningRate', comment='Learning rate for the M3GNet weight optimization.', default=0.001, isenabled="glompo_machinelearning.machinelearning.backend == 'M3GNet' # hide")
self.Model: Literal["UniversalPotential", "Custom", "ModelDir"] = MultipleChoiceKey(name='Model', comment='How to specify the model for the M3GNet backend. Either a Custom model can be made from scratch or an existing model directory can be loaded to obtain the model settings.', default='UniversalPotential', choices=['UniversalPotential', 'Custom', 'ModelDir'])
self.ModelDir: str | Path | StringKey = PathStringKey(name='ModelDir', comment="Path to the directory defining the model. This folder should contain the files: 'checkpoint', 'm3gnet.data-00000-of-00001', ' m3gnet.index' and 'm3gnet.json'", ispath=True, gui_type='directory', isenabled="glompo_machinelearning.machinelearning.m3gnet.model == 'Model Dir' # hide")
self.Custom: ParAMSMachineLearning._MachineLearning._M3GNet._Custom = self._Custom(name='Custom', comment='Specify a custom M3GNet model.', isenabled="glompo_machinelearning.machinelearning.m3gnet.model == 'Custom' # hide")
self.UniversalPotential: ParAMSMachineLearning._MachineLearning._M3GNet._UniversalPotential = self._UniversalPotential(name='UniversalPotential', comment='Settings for (transfer) learning with the M3GNet Universal Potential.')
[docs] class _MACE(FixedBlock):
r"""
Options for MACE fitting.
:ivar LearningRate: Learning rate for the MACE weight optimization
:vartype LearningRate: float | FloatKey
:ivar Model: How to specify the model for the MACE backend. A foundation model can be used, a custom model can be made from scratch or an existing model file can be loaded to obtain the model settings.
:vartype Model: Literal["Foundation", "Custom", "ModelFile"]
:ivar ModelFile: Path to the .model file defining the model.
:vartype ModelFile: str | Path | StringKey
:ivar Custom: Specify a custom MACE model.
:vartype Custom: ParAMSMachineLearning._MachineLearning._MACE._Custom
:ivar Foundation: Settings for (transfer) learning with the MACE foundation model.
:vartype Foundation: ParAMSMachineLearning._MachineLearning._MACE._Foundation
:ivar LoRA: Settings for LoRA (Low-Rank Adaptation) fine-tuning. LoRA freezes the base model weights and trains only small low-rank adapters, reducing overfitting and catastrophic forgetting when fine-tuning a foundation model on a small dataset. Only applied when fine-tuning a foundation/restart model.
:vartype LoRA: ParAMSMachineLearning._MachineLearning._MACE._LoRA
:ivar LossCoeffs: Modify the coefficients for the machine learning loss function, by applying scale factors to the MachineLearning%LossCoeffs values.
:vartype LossCoeffs: ParAMSMachineLearning._MachineLearning._MACE._LossCoeffs
:ivar StageTwo: Settings for stage two of training.
:vartype StageTwo: ParAMSMachineLearning._MachineLearning._MACE._StageTwo
"""
[docs] class _Custom(FixedBlock):
r"""
Specify a custom MACE model.
: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 LMax: Maximum spherical harmonic order of the messages in the message passing step; 0 is invariant, 1 and 2 are higher order (equivariant). Defaults to ``1``. A higher value increases accuracy, but increases model size and computational cost.
:vartype LMax: int | IntKey
:ivar NumChannels: Number of channels per angular momentum order. Defaults to ``128``. Set to ``64`` for a smaller model or ``256`` for a larger model. A higher value increases accuracy, but increases model size and computational cost.
:vartype NumChannels: int | IntKey
:ivar RMax: Distance cutoff for interactions.
:vartype RMax: float | FloatKey
"""
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='float64', choices=['float32', 'float64'], isenabled="glompo_machinelearning.machinelearning.mace.model == 'Custom' # hide")
self.LMax: int | IntKey = IntKey(name='LMax', comment='Maximum spherical harmonic order of the messages in the message passing step; 0 is invariant, 1 and 2 are higher order (equivariant). Defaults to ``1``. A higher value increases accuracy, but increases model size and computational cost.', gui_name='Max spherical harmonic order:', default=1, isenabled="glompo_machinelearning.machinelearning.mace.model == 'Custom' # hide")
self.NumChannels: int | IntKey = IntKey(name='NumChannels', comment='Number of channels per angular momentum order. Defaults to ``128``. Set to ``64`` for a smaller model or ``256`` for a larger model. A higher value increases accuracy, but increases model size and computational cost.', gui_name='Number of channels:', default=128, isenabled="glompo_machinelearning.machinelearning.mace.model == 'Custom' # hide")
self.RMax: float | FloatKey = FloatKey(name='RMax', comment='Distance cutoff for interactions.', gui_name='Distance cutoff:', default=5.0, unit='angstrom', isenabled="glompo_machinelearning.machinelearning.mace.model == 'Custom' # hide")
[docs] class _Foundation(FixedBlock):
r"""
Settings for (transfer) learning with the MACE foundation model.
:ivar Type: Which MACE foundation model to use.
:vartype Type: Literal["MACE-MPA-0", "MACE-MP-0-Large", "MACE-MP-0-Medium", "MACE-MP-0-Small"]
"""
def __post_init__(self):
self.Type: Literal["MACE-MPA-0", "MACE-MP-0-Large", "MACE-MP-0-Medium", "MACE-MP-0-Small"] = MultipleChoiceKey(name='Type', comment='Which MACE foundation model to use.', default='MACE-MPA-0', choices=['MACE-MPA-0', 'MACE-MP-0-Large', 'MACE-MP-0-Medium', 'MACE-MP-0-Small'], isenabled="glompo_machinelearning.machinelearning.mace.model == 'Foundation' # hide")
[docs] class _LoRA(FixedBlock):
r"""
Settings for LoRA (Low-Rank Adaptation) fine-tuning. LoRA freezes the base model weights and trains only small low-rank adapters, reducing overfitting and catastrophic forgetting when fine-tuning a foundation model on a small dataset. Only applied when fine-tuning a foundation/restart model.
:ivar Alpha: Scaling factor for the LoRA update. The effective scaling applied to the low-rank path is Alpha / Rank.
:vartype Alpha: float | FloatKey
:ivar Enabled: Whether to enable LoRA fine-tuning, defaults to ``True``.
:vartype Enabled: BoolType | BoolKey
:ivar Rank: Rank of the LoRA matrices. Higher rank increases capacity and the number of trainable parameters.
:vartype Rank: int | IntKey
"""
def __post_init__(self):
self.Alpha: float | FloatKey = FloatKey(name='Alpha', comment='Scaling factor for the LoRA update. The effective scaling applied to the low-rank path is Alpha / Rank.', gui_name='LoRA alpha:', default=1.0, isenabled="glompo_machinelearning.machinelearning.backend == 'MACE' and glompo_machinelearning.machinelearning.mace.lora.enabled")
self.Enabled: BoolType | BoolKey = BoolKey(name='Enabled', comment='Whether to enable LoRA fine-tuning, defaults to ``True``.', gui_name='LoRA enabled:', default=True, isenabled="glompo_machinelearning.machinelearning.backend == 'MACE' # hide")
self.Rank: int | IntKey = IntKey(name='Rank', comment='Rank of the LoRA matrices. Higher rank increases capacity and the number of trainable parameters.', gui_name='LoRA rank:', default=4, isenabled="glompo_machinelearning.machinelearning.backend == 'MACE' and glompo_machinelearning.machinelearning.mace.lora.enabled")
[docs] class _LossCoeffs(FixedBlock):
r"""
Modify the coefficients for the machine learning loss function, by applying scale factors to the MachineLearning%LossCoeffs values.
:ivar EnergyScaleFactor: Scale factor to apply to the energy loss coefficient, i.e. EnergyScaleFactor * MachineLearning%LossCoeffs%Energy.
:vartype EnergyScaleFactor: float | FloatKey
:ivar ForcesScaleFactor: Scale factor to apply to the forces loss coefficient i.e. ForcesScaleFactor * MachineLearning%LossCoeffs%Forces.
:vartype ForcesScaleFactor: float | FloatKey
"""
def __post_init__(self):
self.EnergyScaleFactor: float | FloatKey = FloatKey(name='EnergyScaleFactor', comment='Scale factor to apply to the energy loss coefficient, i.e. EnergyScaleFactor * MachineLearning%LossCoeffs%Energy.', default=0.1, isenabled="glompo_machinelearning.machinelearning.backend == 'MACE' # hide")
self.ForcesScaleFactor: float | FloatKey = FloatKey(name='ForcesScaleFactor', comment='Scale factor to apply to the forces loss coefficient i.e. ForcesScaleFactor * MachineLearning%LossCoeffs%Forces.', default=100.0, isenabled="glompo_machinelearning.machinelearning.backend == 'MACE' # hide")
[docs] class _StageTwo(FixedBlock):
r"""
Settings for stage two of training.
:ivar Enabled: Whether to enable stage two training, defaults to ``True``.
:vartype Enabled: BoolType | BoolKey
:ivar LearningRate: Learning rate for the MACE weight optimization for stage two training
:vartype LearningRate: float | FloatKey
:ivar Start: When to start stage two training, as a proportion of the MachineLearning%MaxEpochs
:vartype Start: float | FloatKey
:ivar LossCoeffs: Modify the coefficients for the machine learning loss function for stage two training, by applying scale factors to the MachineLearning%LossCoeffs values.
:vartype LossCoeffs: ParAMSMachineLearning._MachineLearning._MACE._StageTwo._LossCoeffs
"""
[docs] class _LossCoeffs(FixedBlock):
r"""
Modify the coefficients for the machine learning loss function for stage two training, by applying scale factors to the MachineLearning%LossCoeffs values.
:ivar EnergyScaleFactor: Scale factor to apply to the energy loss coefficient for stage two training, i.e. EnergyScaleFactor * MachineLearning%LossCoeffs%Energy.
:vartype EnergyScaleFactor: float | FloatKey
:ivar ForcesScaleFactor: Scale factor to apply to the forces loss coefficient for stage two training i.e. ForcesScaleFactor * MachineLearning%LossCoeffs%Forces.
:vartype ForcesScaleFactor: float | FloatKey
"""
def __post_init__(self):
self.EnergyScaleFactor: float | FloatKey = FloatKey(name='EnergyScaleFactor', comment='Scale factor to apply to the energy loss coefficient for stage two training, i.e. EnergyScaleFactor * MachineLearning%LossCoeffs%Energy.', gui_name='Stage two energy scale factor:', default=100.0)
self.ForcesScaleFactor: float | FloatKey = FloatKey(name='ForcesScaleFactor', comment='Scale factor to apply to the forces loss coefficient for stage two training i.e. ForcesScaleFactor * MachineLearning%LossCoeffs%Forces.', gui_name='Stage two forces scale factor:', default=100.0)
def __post_init__(self):
self.Enabled: BoolType | BoolKey = BoolKey(name='Enabled', comment='Whether to enable stage two training, defaults to ``True``.', gui_name='Stage two enabled:', default=True, isenabled="glompo_machinelearning.machinelearning.backend == 'MACE' # hide")
self.LearningRate: float | FloatKey = FloatKey(name='LearningRate', comment='Learning rate for the MACE weight optimization for stage two training', gui_name='Stage two learning rate:', default=0.001, isenabled="glompo_machinelearning.machinelearning.backend == 'MACE' and glompo_machinelearning.machinelearning.mace.stagetwo.enabled")
self.Start: float | FloatKey = FloatKey(name='Start', comment='When to start stage two training, as a proportion of the MachineLearning%MaxEpochs', gui_name='Stage two start:', default=0.8, isenabled="glompo_machinelearning.machinelearning.backend == 'MACE' and glompo_machinelearning.machinelearning.mace.stagetwo.enabled")
self.LossCoeffs: ParAMSMachineLearning._MachineLearning._MACE._StageTwo._LossCoeffs = self._LossCoeffs(name='LossCoeffs', comment='Modify the coefficients for the machine learning loss function for stage two training, by applying scale factors to the MachineLearning%LossCoeffs values.', isenabled="glompo_machinelearning.machinelearning.backend == 'MACE' and glompo_machinelearning.machinelearning.mace.stagetwo.enabled")
def __post_init__(self):
self.LearningRate: float | FloatKey = FloatKey(name='LearningRate', comment='Learning rate for the MACE weight optimization', default=0.01, isenabled="glompo_machinelearning.machinelearning.backend == 'MACE' # hide")
self.Model: Literal["Foundation", "Custom", "ModelFile"] = MultipleChoiceKey(name='Model', comment='How to specify the model for the MACE backend. A foundation model can be used, a custom model can be made from scratch or an existing model file can be loaded to obtain the model settings.', default='Foundation', choices=['Foundation', 'Custom', 'ModelFile'], isenabled="glompo_machinelearning.machinelearning.backend == 'MACE' and glompo_machinelearning.machinelearning.loadmodel == '' # hide")
self.ModelFile: str | Path | StringKey = PathStringKey(name='ModelFile', comment='Path to the .model file defining the model.', ispath=True, gui_type='file', isenabled="glompo_machinelearning.machinelearning.backend == 'MACE' and glompo_machinelearning.machinelearning.loadmodel == '' and glompo_machinelearning.machinelearning.mace.model == 'Model File' # hide")
self.Custom: ParAMSMachineLearning._MachineLearning._MACE._Custom = self._Custom(name='Custom', comment='Specify a custom MACE model.', isenabled="glompo_machinelearning.machinelearning.backend == 'MACE' and glompo_machinelearning.machinelearning.loadmodel == '' and glompo_machinelearning.machinelearning.mace.model == 'Custom' # hide")
self.Foundation: ParAMSMachineLearning._MachineLearning._MACE._Foundation = self._Foundation(name='Foundation', comment='Settings for (transfer) learning with the MACE foundation model.', isenabled="glompo_machinelearning.machinelearning.backend == 'MACE' and glompo_machinelearning.machinelearning.loadmodel == '' and glompo_machinelearning.machinelearning.mace.model == 'Foundation' # hide")
self.LoRA: ParAMSMachineLearning._MachineLearning._MACE._LoRA = self._LoRA(name='LoRA', comment='Settings for LoRA (Low-Rank Adaptation) fine-tuning. LoRA freezes the base model weights and trains only small low-rank adapters, reducing overfitting and catastrophic forgetting when fine-tuning a foundation model on a small dataset. Only applied when fine-tuning a foundation/restart model.', isenabled="glompo_machinelearning.machinelearning.backend == 'MACE'")
self.LossCoeffs: ParAMSMachineLearning._MachineLearning._MACE._LossCoeffs = self._LossCoeffs(name='LossCoeffs', comment='Modify the coefficients for the machine learning loss function, by applying scale factors to the MachineLearning%LossCoeffs values.')
self.StageTwo: ParAMSMachineLearning._MachineLearning._MACE._StageTwo = self._StageTwo(name='StageTwo', comment='Settings for stage two of training.', isenabled="glompo_machinelearning.machinelearning.backend == 'MACE'")
[docs] class _MatGL(FixedBlock):
r"""
Options for fitting MatGL QET and TensorNet potentials.
:ivar LearningRate: Initial learning rate for MatGL weight optimization. The LearningRateSchedule block controls how this rate changes during training.
:vartype LearningRate: float | FloatKey
:ivar Model: Select a packaged foundation model to fine-tune, create a custom architecture, or load a MatGL potential directory.
:vartype Model: Literal["M3GNet-PBE-2025", "M3GNet-r2SCAN-2025", "QET-PBE-2025", "QET-r2SCAN-2025", "TensorNet-PBE-M-2025", "TensorNet-r2SCAN-M-2025", "Custom", "ModelDir"]
:ivar ModelDir: Path to a MatGL potential directory containing model.json, model.pt, and state.pt.
:vartype ModelDir: str | Path | StringKey
:ivar TrainChargeHeads: Whether QET electronegativity and hardness heads are optimized. The default, No, retains those heads from the loaded model while training the layers selected by TrainableLayers. Set this to Yes for a custom QET model trained from scratch so that its randomly initialized charge-head weights are trained. This option only affects QET models.
:vartype TrainChargeHeads: BoolType | BoolKey
:ivar TrainableLayers: Select which neural-network parameters to optimize. All is the recommended default, trains the complete model, and should normally be used for a custom model trained from scratch. ReadoutOnly trains only the final energy-related output heads. LastBlockAndReadout additionally trains the final interaction block and its output projection. The analytically fitted elemental energy references are updated independently of this setting. This setting is reapplied when loading or restarting a model.
:vartype TrainableLayers: Literal["All", "ReadoutOnly", "LastBlockAndReadout"]
:ivar Custom: Create a model from scratch with MatGL 4.0.2 constructor defaults.
:vartype Custom: ParAMSMachineLearning._MachineLearning._MatGL._Custom
:ivar LearningRateSchedule: Control how the MatGL learning rate changes during training and after a full-checkpoint restart.
:vartype LearningRateSchedule: ParAMSMachineLearning._MachineLearning._MatGL._LearningRateSchedule
"""
[docs] class _Custom(FixedBlock):
r"""
Create a model from scratch with MatGL 4.0.2 constructor defaults.
:ivar Architecture: MatGL architecture to train.
:vartype Architecture: Literal["M3GNet", "QET", "TensorNet"]
:ivar M3GNet: M3GNet constructor settings.
:vartype M3GNet: ParAMSMachineLearning._MachineLearning._MatGL._Custom._M3GNet
:ivar QET: QET constructor settings.
:vartype QET: ParAMSMachineLearning._MachineLearning._MatGL._Custom._QET
:ivar TensorNet: TensorNet constructor settings.
:vartype TensorNet: ParAMSMachineLearning._MachineLearning._MatGL._Custom._TensorNet
"""
[docs] class _M3GNet(FixedBlock):
r"""
M3GNet constructor settings.
:ivar Activation: Activation function.
:vartype Activation: Literal["swish", "tanh", "sigmoid", "softplus2", "softexp"]
:ivar Cutoff: Pair interaction cutoff.
:vartype Cutoff: float | FloatKey
:ivar DimEdgeEmbedding: Edge embedding dimension.
:vartype DimEdgeEmbedding: int | IntKey
:ivar DimNodeEmbedding: Node embedding dimension.
:vartype DimNodeEmbedding: int | IntKey
:ivar MaxL: Maximum angular basis order.
:vartype MaxL: int | IntKey
:ivar MaxN: Maximum radial basis index.
:vartype MaxN: int | IntKey
:ivar NumBlocks: Number of graph convolution blocks.
:vartype NumBlocks: int | IntKey
:ivar NumNeurons: Number of neurons in each update layer.
:vartype NumNeurons: int | IntKey
:ivar ThreebodyCutoff: Three-body interaction cutoff.
:vartype ThreebodyCutoff: float | FloatKey
:ivar UsePhi: Include the azimuthal angle in the three-body basis.
:vartype UsePhi: BoolType | BoolKey
:ivar UseSmooth: Use the smooth spherical Bessel basis.
:vartype UseSmooth: BoolType | BoolKey
"""
def __post_init__(self):
self.Activation: Literal["swish", "tanh", "sigmoid", "softplus2", "softexp"] = MultipleChoiceKey(name='Activation', comment='Activation function.', default='swish', choices=['swish', 'tanh', 'sigmoid', 'softplus2', 'softexp'])
self.Cutoff: float | FloatKey = FloatKey(name='Cutoff', comment='Pair interaction cutoff.', default=5.0, unit='angstrom')
self.DimEdgeEmbedding: int | IntKey = IntKey(name='DimEdgeEmbedding', comment='Edge embedding dimension.', default=64)
self.DimNodeEmbedding: int | IntKey = IntKey(name='DimNodeEmbedding', comment='Node embedding dimension.', default=64)
self.MaxL: int | IntKey = IntKey(name='MaxL', comment='Maximum angular basis order.', default=3)
self.MaxN: int | IntKey = IntKey(name='MaxN', comment='Maximum radial basis index.', default=3)
self.NumBlocks: int | IntKey = IntKey(name='NumBlocks', comment='Number of graph convolution blocks.', default=3)
self.NumNeurons: int | IntKey = IntKey(name='NumNeurons', comment='Number of neurons in each update layer.', default=64)
self.ThreebodyCutoff: float | FloatKey = FloatKey(name='ThreebodyCutoff', comment='Three-body interaction cutoff.', default=4.0, unit='angstrom')
self.UsePhi: BoolType | BoolKey = BoolKey(name='UsePhi', comment='Include the azimuthal angle in the three-body basis.', default=False)
self.UseSmooth: BoolType | BoolKey = BoolKey(name='UseSmooth', comment='Use the smooth spherical Bessel basis.', default=False)
[docs] class _QET(FixedBlock):
r"""
QET constructor settings.
:ivar Activation: Activation function.
:vartype Activation: Literal["swish", "tanh", "sigmoid", "softplus2", "softexp"]
:ivar Cutoff: Pair interaction cutoff.
:vartype Cutoff: float | FloatKey
:ivar EnvironmentDependentHardness: Predict QET hardness from the local environment.
:vartype EnvironmentDependentHardness: BoolType | BoolKey
:ivar EquivarianceGroup: Equivariance and invariance group.
:vartype EquivarianceGroup: Literal["O(3)", "SO(3)"]
:ivar MaxL: Maximum angular basis order.
:vartype MaxL: int | IntKey
:ivar MaxN: Maximum radial basis index.
:vartype MaxN: int | IntKey
:ivar NumBlocks: Number of interaction blocks.
:vartype NumBlocks: int | IntKey
:ivar NumNeurons: Hidden embedding size.
:vartype NumNeurons: int | IntKey
:ivar NumRBF: Number of radial basis functions.
:vartype NumRBF: int | IntKey
:ivar RBFType: Radial basis type.
:vartype RBFType: Literal["Gaussian", "SphericalBessel"]
:ivar TrainSigma: Train QET Gaussian charge widths.
:vartype TrainSigma: BoolType | BoolKey
:ivar UseSmooth: Use a smooth spherical Bessel basis.
:vartype UseSmooth: BoolType | BoolKey
:ivar Width: Gaussian radial basis width.
:vartype Width: float | FloatKey
"""
def __post_init__(self):
self.Activation: Literal["swish", "tanh", "sigmoid", "softplus2", "softexp"] = MultipleChoiceKey(name='Activation', comment='Activation function.', default='swish', choices=['swish', 'tanh', 'sigmoid', 'softplus2', 'softexp'])
self.Cutoff: float | FloatKey = FloatKey(name='Cutoff', comment='Pair interaction cutoff.', default=5.0, unit='angstrom')
self.EnvironmentDependentHardness: BoolType | BoolKey = BoolKey(name='EnvironmentDependentHardness', comment='Predict QET hardness from the local environment.', default=False)
self.EquivarianceGroup: Literal["O(3)", "SO(3)"] = MultipleChoiceKey(name='EquivarianceGroup', comment='Equivariance and invariance group.', default='O(3)', choices=['O(3)', 'SO(3)'])
self.MaxL: int | IntKey = IntKey(name='MaxL', comment='Maximum angular basis order.', default=3)
self.MaxN: int | IntKey = IntKey(name='MaxN', comment='Maximum radial basis index.', default=3)
self.NumBlocks: int | IntKey = IntKey(name='NumBlocks', comment='Number of interaction blocks.', default=2)
self.NumNeurons: int | IntKey = IntKey(name='NumNeurons', comment='Hidden embedding size.', default=64)
self.NumRBF: int | IntKey = IntKey(name='NumRBF', comment='Number of radial basis functions.', default=32)
self.RBFType: Literal["Gaussian", "SphericalBessel"] = MultipleChoiceKey(name='RBFType', comment='Radial basis type.', default='Gaussian', choices=['Gaussian', 'SphericalBessel'])
self.TrainSigma: BoolType | BoolKey = BoolKey(name='TrainSigma', comment='Train QET Gaussian charge widths.', default=False)
self.UseSmooth: BoolType | BoolKey = BoolKey(name='UseSmooth', comment='Use a smooth spherical Bessel basis.', default=False)
self.Width: float | FloatKey = FloatKey(name='Width', comment='Gaussian radial basis width.', default=0.5)
[docs] class _TensorNet(FixedBlock):
r"""
TensorNet constructor settings.
:ivar Activation: Activation function.
:vartype Activation: Literal["swish", "tanh", "sigmoid", "softplus2", "softexp"]
:ivar Cutoff: Pair interaction cutoff.
:vartype Cutoff: float | FloatKey
:ivar EquivarianceGroup: Equivariance and invariance group.
:vartype EquivarianceGroup: Literal["O(3)", "SO(3)"]
:ivar MaxL: Maximum angular basis order.
:vartype MaxL: int | IntKey
:ivar MaxN: Maximum radial basis index.
:vartype MaxN: int | IntKey
:ivar NumBlocks: Number of interaction blocks.
:vartype NumBlocks: int | IntKey
:ivar NumNeurons: Hidden embedding size.
:vartype NumNeurons: int | IntKey
:ivar NumRBF: Number of radial basis functions.
:vartype NumRBF: int | IntKey
:ivar RBFType: Radial basis type.
:vartype RBFType: Literal["Gaussian", "SphericalBessel"]
:ivar UseSmooth: Use a smooth spherical Bessel basis.
:vartype UseSmooth: BoolType | BoolKey
:ivar Width: Gaussian radial basis width.
:vartype Width: float | FloatKey
"""
def __post_init__(self):
self.Activation: Literal["swish", "tanh", "sigmoid", "softplus2", "softexp"] = MultipleChoiceKey(name='Activation', comment='Activation function.', default='swish', choices=['swish', 'tanh', 'sigmoid', 'softplus2', 'softexp'])
self.Cutoff: float | FloatKey = FloatKey(name='Cutoff', comment='Pair interaction cutoff.', default=5.0, unit='angstrom')
self.EquivarianceGroup: Literal["O(3)", "SO(3)"] = MultipleChoiceKey(name='EquivarianceGroup', comment='Equivariance and invariance group.', default='O(3)', choices=['O(3)', 'SO(3)'])
self.MaxL: int | IntKey = IntKey(name='MaxL', comment='Maximum angular basis order.', default=3)
self.MaxN: int | IntKey = IntKey(name='MaxN', comment='Maximum radial basis index.', default=3)
self.NumBlocks: int | IntKey = IntKey(name='NumBlocks', comment='Number of interaction blocks.', default=2)
self.NumNeurons: int | IntKey = IntKey(name='NumNeurons', comment='Hidden embedding size.', default=64)
self.NumRBF: int | IntKey = IntKey(name='NumRBF', comment='Number of radial basis functions.', default=32)
self.RBFType: Literal["Gaussian", "SphericalBessel"] = MultipleChoiceKey(name='RBFType', comment='Radial basis type.', default='Gaussian', choices=['Gaussian', 'SphericalBessel'])
self.UseSmooth: BoolType | BoolKey = BoolKey(name='UseSmooth', comment='Use a smooth spherical Bessel basis.', default=False)
self.Width: float | FloatKey = FloatKey(name='Width', comment='Gaussian radial basis width.', default=0.5)
def __post_init__(self):
self.Architecture: Literal["M3GNet", "QET", "TensorNet"] = MultipleChoiceKey(name='Architecture', comment='MatGL architecture to train.', default='QET', choices=['M3GNet', 'QET', 'TensorNet'], hiddenchoices=['M3GNet'])
self.M3GNet: ParAMSMachineLearning._MachineLearning._MatGL._Custom._M3GNet = self._M3GNet(name='M3GNet', comment='M3GNet constructor settings.', isenabled="glompo_machinelearning.machinelearning.matgl.custom.architecture == 'M3GNet' # hide")
self.QET: ParAMSMachineLearning._MachineLearning._MatGL._Custom._QET = self._QET(name='QET', comment='QET constructor settings.', isenabled="glompo_machinelearning.machinelearning.matgl.custom.architecture == 'QET' # hide")
self.TensorNet: ParAMSMachineLearning._MachineLearning._MatGL._Custom._TensorNet = self._TensorNet(name='TensorNet', comment='TensorNet constructor settings.', gui_name='TensorNet', isenabled="glompo_machinelearning.machinelearning.matgl.custom.architecture == 'TensorNet' # hide")
[docs] class _LearningRateSchedule(FixedBlock):
r"""
Control how the MatGL learning rate changes during training and after a full-checkpoint restart.
:ivar FinalFactor: Final cosine learning rate as a fraction of LearningRate. Must be between 0 and 1. Example: if LearningRate is 0.01 and FinalFactor is 0.1, then the learning rate will decrease from 0.01 to 0.001 during the training.
:vartype FinalFactor: float | FloatKey
:ivar Restart: For a full-checkpoint restart, Continue starts a new schedule from the checkpoint learning rate without increasing it. Reset starts from LearningRate. In both cases, MaxEpochs is the length of the new schedule. This option has no effect for a weights-only load.
:vartype Restart: Literal["Reset", "Continue"]
:ivar Type: Cosine decreases the learning rate over MaxEpochs. Constant keeps it fixed.
:vartype Type: Literal["Cosine", "Constant"]
"""
def __post_init__(self):
self.FinalFactor: float | FloatKey = FloatKey(name='FinalFactor', comment='Final cosine learning rate as a fraction of LearningRate. Must be between 0 and 1. Example: if LearningRate is 0.01 and FinalFactor is 0.1, then the learning rate will decrease from 0.01 to 0.001 during the training.', default=0.01, isenabled="glompo_machinelearning.machinelearning.matgl.learningrateschedule.type == 'Cosine'")
self.Restart: Literal["Reset", "Continue"] = MultipleChoiceKey(name='Restart', comment='For a full-checkpoint restart, Continue starts a new schedule from the checkpoint learning rate without increasing it. Reset starts from LearningRate. In both cases, MaxEpochs is the length of the new schedule. This option has no effect for a weights-only load.', default='Reset', choices=['Reset', 'Continue'])
self.Type: Literal["Cosine", "Constant"] = MultipleChoiceKey(name='Type', comment='Cosine decreases the learning rate over MaxEpochs. Constant keeps it fixed.', default='Cosine', choices=['Cosine', 'Constant'])
def __post_init__(self):
self.LearningRate: float | FloatKey = FloatKey(name='LearningRate', comment='Initial learning rate for MatGL weight optimization. The LearningRateSchedule block controls how this rate changes during training.', default=0.001)
self.Model: Literal["M3GNet-PBE-2025", "M3GNet-r2SCAN-2025", "QET-PBE-2025", "QET-r2SCAN-2025", "TensorNet-PBE-M-2025", "TensorNet-r2SCAN-M-2025", "Custom", "ModelDir"] = MultipleChoiceKey(name='Model', comment='Select a packaged foundation model to fine-tune, create a custom architecture, or load a MatGL potential directory.', default='QET-PBE-2025', choices=['M3GNet-PBE-2025', 'M3GNet-r2SCAN-2025', 'QET-PBE-2025', 'QET-r2SCAN-2025', 'TensorNet-PBE-M-2025', 'TensorNet-r2SCAN-M-2025', 'Custom', 'ModelDir'], hiddenchoices=['M3GNet-PBE-2025', 'M3GNet-r2SCAN-2025'], isenabled="glompo_machinelearning.machinelearning.loadmodel == '' # hide")
self.ModelDir: str | Path | StringKey = PathStringKey(name='ModelDir', comment='Path to a MatGL potential directory containing model.json, model.pt, and state.pt.', ispath=True, gui_type='directory', isenabled="glompo_machinelearning.machinelearning.loadmodel == '' and glompo_machinelearning.machinelearning.matgl.model == 'ModelDir' # hide")
self.TrainChargeHeads: BoolType | BoolKey = BoolKey(name='TrainChargeHeads', comment='Whether QET electronegativity and hardness heads are optimized. The default, No, retains those heads from the loaded model while training the layers selected by TrainableLayers. Set this to Yes for a custom QET model trained from scratch so that its randomly initialized charge-head weights are trained. This option only affects QET models.', gui_name='Train QET charge heads:', default=False)
self.TrainableLayers: Literal["All", "ReadoutOnly", "LastBlockAndReadout"] = MultipleChoiceKey(name='TrainableLayers', comment='Select which neural-network parameters to optimize. All is the recommended default, trains the complete model, and should normally be used for a custom model trained from scratch. ReadoutOnly trains only the final energy-related output heads. LastBlockAndReadout additionally trains the final interaction block and its output projection. The analytically fitted elemental energy references are updated independently of this setting. This setting is reapplied when loading or restarting a model.', gui_name='Trainable layers:', default='All', choices=['All', 'ReadoutOnly', 'LastBlockAndReadout'])
self.Custom: ParAMSMachineLearning._MachineLearning._MatGL._Custom = self._Custom(name='Custom', comment='Create a model from scratch with MatGL 4.0.2 constructor defaults.', isenabled="glompo_machinelearning.machinelearning.loadmodel == '' and glompo_machinelearning.machinelearning.matgl.model == 'Custom' # hide")
self.LearningRateSchedule: ParAMSMachineLearning._MachineLearning._MatGL._LearningRateSchedule = self._LearningRateSchedule(name='LearningRateSchedule', comment='Control how the MatGL learning rate changes during training and after a full-checkpoint restart.')
[docs] class _NEP(FixedBlock):
r"""
Experimental GPUMD NEP4 fitting. Training uses AMS_NEP_TRAIN_EXECUTABLE; production calculations use the native NEP worker installed by AMS. The default Foundation model fine-tunes the installed NEP89 potential for neutral, fully 3D-periodic systems.
:ivar Model: Fine-tune the included NEP89 foundation potential, or train a scalar NEP4 model from scratch.
:vartype Model: Literal["Foundation", "Custom"]
:ivar Advanced: One supported GPUMD directive per line. Generated architecture, loss, SNES and checkpoint directives are rejected to keep input unambiguous.
:vartype Advanced: ParAMSMachineLearning._MachineLearning._NEP._Advanced
:ivar Custom: Architecture for a new scalar NEP4 model. These settings are locked when fine-tuning NEP89.
:vartype Custom: ParAMSMachineLearning._MachineLearning._NEP._Custom
:ivar Optimization: GPUMD SNES and checkpoint settings. MachineLearning%MaxEpochs maps to generation; the common energy and force coefficients map to lambda_e and lambda_f.
:vartype Optimization: ParAMSMachineLearning._MachineLearning._NEP._Optimization
"""
[docs] class _Advanced(FixedBlock):
r"""
One supported GPUMD directive per line. Generated architecture, loss, SNES and checkpoint directives are rejected to keep input unambiguous.
:ivar Raw: Additional supported GPUMD nep.in directives, one per line.
:vartype Raw: str | Sequence[str] | FreeBlock
"""
[docs] class _Raw(FreeBlock):
r"""
Additional supported GPUMD nep.in directives, one per line.
"""
def __post_init__(self):
pass
def __post_init__(self):
self.Raw: str | Sequence[str] | FreeBlock = self._Raw(name='Raw', comment='Additional supported GPUMD nep.in directives, one per line.')
[docs] class _Custom(FixedBlock):
r"""
Architecture for a new scalar NEP4 model. These settings are locked when fine-tuning NEP89.
:ivar BasisSizeAngular: Angular basis size. GPUMD default: 6; NEP89 foundation: 8.
:vartype BasisSizeAngular: int | IntKey
:ivar BasisSizeRadial: Radial basis size. GPUMD default: 6; NEP89 foundation: 8.
:vartype BasisSizeRadial: int | IntKey
:ivar CutoffAngular: Angular descriptor cutoff. GPUMD default: 4.0 angstrom; NEP89 foundation: 5.0 angstrom.
:vartype CutoffAngular: float | FloatKey
:ivar CutoffRadial: Radial descriptor cutoff. GPUMD default: 8.0 angstrom; NEP89 foundation: 6.0 angstrom.
:vartype CutoffRadial: float | FloatKey
:ivar LMax3Body: Three-body angular expansion. GPUMD default: 4; NEP89 foundation: 4.
:vartype LMax3Body: int | IntKey
:ivar LMax4Body: Four-body angular expansion. GPUMD default: 1; NEP89 foundation: 2.
:vartype LMax4Body: int | IntKey
:ivar LMax5Body: Five-body angular expansion. GPUMD default: 0; NEP89 foundation: 1.
:vartype LMax5Body: int | IntKey
:ivar NMaxAngular: Maximum angular basis order. GPUMD default: 6; NEP89 foundation: 4.
:vartype NMaxAngular: int | IntKey
:ivar NMaxRadial: Maximum radial basis order. GPUMD default: 6; NEP89 foundation: 4.
:vartype NMaxRadial: int | IntKey
:ivar NumNeurons: Number of neurons in the hidden layer. GPUMD default: 30; NEP89 foundation: 80.
:vartype NumNeurons: int | IntKey
:ivar Version: NEP model version. GPUMD default: 4; NEP89 foundation: 4.
:vartype Version: int | IntKey
:ivar ZBL: Outer cutoff radius in angstrom for the universal Ziegler-Biersack-Littmark (ZBL) repulsive potential. Set to 0 to disable ZBL. GPUMD default: 0 (disabled); NEP89 foundation: 2 angstrom.
:vartype ZBL: int | IntKey
"""
def __post_init__(self):
self.BasisSizeAngular: int | IntKey = IntKey(name='BasisSizeAngular', comment='Angular basis size. GPUMD default: 6; NEP89 foundation: 8.', default=8)
self.BasisSizeRadial: int | IntKey = IntKey(name='BasisSizeRadial', comment='Radial basis size. GPUMD default: 6; NEP89 foundation: 8.', default=8)
self.CutoffAngular: float | FloatKey = FloatKey(name='CutoffAngular', comment='Angular descriptor cutoff. GPUMD default: 4.0 angstrom; NEP89 foundation: 5.0 angstrom.', default=5.0, unit='angstrom')
self.CutoffRadial: float | FloatKey = FloatKey(name='CutoffRadial', comment='Radial descriptor cutoff. GPUMD default: 8.0 angstrom; NEP89 foundation: 6.0 angstrom.', default=6.0, unit='angstrom')
self.LMax3Body: int | IntKey = IntKey(name='LMax3Body', comment='Three-body angular expansion. GPUMD default: 4; NEP89 foundation: 4.', default=4)
self.LMax4Body: int | IntKey = IntKey(name='LMax4Body', comment='Four-body angular expansion. GPUMD default: 1; NEP89 foundation: 2.', default=2)
self.LMax5Body: int | IntKey = IntKey(name='LMax5Body', comment='Five-body angular expansion. GPUMD default: 0; NEP89 foundation: 1.', default=1)
self.NMaxAngular: int | IntKey = IntKey(name='NMaxAngular', comment='Maximum angular basis order. GPUMD default: 6; NEP89 foundation: 4.', default=4)
self.NMaxRadial: int | IntKey = IntKey(name='NMaxRadial', comment='Maximum radial basis order. GPUMD default: 6; NEP89 foundation: 4.', default=4)
self.NumNeurons: int | IntKey = IntKey(name='NumNeurons', comment='Number of neurons in the hidden layer. GPUMD default: 30; NEP89 foundation: 80.', default=80)
self.Version: int | IntKey = IntKey(name='Version', comment='NEP model version. GPUMD default: 4; NEP89 foundation: 4.', default=4)
self.ZBL: int | IntKey = IntKey(name='ZBL', comment='Outer cutoff radius in angstrom for the universal Ziegler-Biersack-Littmark (ZBL) repulsive potential. Set to 0 to disable ZBL. GPUMD default: 0 (disabled); NEP89 foundation: 2 angstrom.', default=0)
[docs] class _Optimization(FixedBlock):
r"""
GPUMD SNES and checkpoint settings. MachineLearning%MaxEpochs maps to generation; the common energy and force coefficients map to lambda_e and lambda_f.
:ivar Batch: Number of structures in each GPUMD training batch. GPUMD default: 1000; NEP89 foundation: 5000.
:vartype Batch: int | IntKey
:ivar Lambda1: L1 regularization coefficient. GPUMD default: -1 (automatic regularization); NEP89 foundation: 0.
:vartype Lambda1: float | FloatKey
:ivar LambdaV: Virial loss coefficient. Virial data are not currently supplied by ParAMS. GPUMD default: 0.1; NEP89 foundation: 1.
:vartype LambdaV: float | FloatKey
:ivar OutputInterval: Number of generations between loss.out rows and updates to nep.txt. GPUMD default: 100; the NEP89 template leaves this at the GPUMD default.
:vartype OutputInterval: int | IntKey
:ivar Population: Population size for GPUMD's separable natural evolution strategy optimizer. GPUMD default: 50; NEP89 foundation: 50.
:vartype Population: int | IntKey
:ivar SavePotential: Arguments to GPUMD save_potential: interval, filename format, and whether to save a restart artifact. GPUMD default interval: 100000 (timestamped filenames); NEP89 foundation: 1000 0 1.
:vartype SavePotential: str | StringKey
"""
def __post_init__(self):
self.Batch: int | IntKey = IntKey(name='Batch', comment='Number of structures in each GPUMD training batch. GPUMD default: 1000; NEP89 foundation: 5000.', default=5000)
self.Lambda1: float | FloatKey = FloatKey(name='Lambda1', comment='L1 regularization coefficient. GPUMD default: -1 (automatic regularization); NEP89 foundation: 0.', default=0.0)
self.LambdaV: float | FloatKey = FloatKey(name='LambdaV', comment='Virial loss coefficient. Virial data are not currently supplied by ParAMS. GPUMD default: 0.1; NEP89 foundation: 1.', default=1.0)
self.OutputInterval: int | IntKey = IntKey(name='OutputInterval', comment='Number of generations between loss.out rows and updates to nep.txt. GPUMD default: 100; the NEP89 template leaves this at the GPUMD default.', default=100)
self.Population: int | IntKey = IntKey(name='Population', comment="Population size for GPUMD's separable natural evolution strategy optimizer. GPUMD default: 50; NEP89 foundation: 50.", default=50)
self.SavePotential: str | StringKey = StringKey(name='SavePotential', comment='Arguments to GPUMD save_potential: interval, filename format, and whether to save a restart artifact. GPUMD default interval: 100000 (timestamped filenames); NEP89 foundation: 1000 0 1.', default='1000 0 1')
def __post_init__(self):
self.Model: Literal["Foundation", "Custom"] = MultipleChoiceKey(name='Model', comment='Fine-tune the included NEP89 foundation potential, or train a scalar NEP4 model from scratch.', default='Foundation', choices=['Foundation', 'Custom'])
self.Advanced: ParAMSMachineLearning._MachineLearning._NEP._Advanced = self._Advanced(name='Advanced', comment='One supported GPUMD directive per line. Generated architecture, loss, SNES and checkpoint directives are rejected to keep input unambiguous.')
self.Custom: ParAMSMachineLearning._MachineLearning._NEP._Custom = self._Custom(name='Custom', comment='Architecture for a new scalar NEP4 model. These settings are locked when fine-tuning NEP89.', isenabled="glompo_machinelearning.machinelearning.nep.model == 'Custom' # hide")
self.Optimization: ParAMSMachineLearning._MachineLearning._NEP._Optimization = self._Optimization(name='Optimization', comment='GPUMD SNES and checkpoint settings. MachineLearning%MaxEpochs maps to generation; the common energy and force coefficients map to lambda_e and lambda_f.')
[docs] class _NequIP(FixedBlock):
r"""
Options for NequIP fitting.
:ivar LearningRate: Learning rate for the NequIP weight optimization
:vartype LearningRate: float | FloatKey
:ivar Model: How to specify the model for the NequIP backend. Either a Custom model can be made from scratch or an existing 'model.pth' file can be loaded to obtain the model settings.
:vartype Model: Literal["Custom", "ModelFile"]
:ivar ModelFile: Path to the model.pth file defining the model.
:vartype ModelFile: str | Path | StringKey
:ivar UseRescalingFromLoadedModel: When loading a model with LoadModel or NequiP%ModelFile do not recalculate the dataset rescaling but use the value from the loaded model.
:vartype UseRescalingFromLoadedModel: BoolType | BoolKey
:ivar Custom: Specify a custom NequIP model.
:vartype Custom: ParAMSMachineLearning._MachineLearning._NequIP._Custom
"""
[docs] class _Custom(FixedBlock):
r"""
Specify a custom NequIP model.
:ivar LMax: Maximum L value. 1 is probably high enough.
:vartype LMax: int | IntKey
:ivar MetricsKey: Which metric to use to generate the 'best' model.
:vartype MetricsKey: Literal["training_loss", "validation_loss"]
:ivar NumLayers: Number of interaction layers in the NequIP neural network.
:vartype NumLayers: int | IntKey
:ivar RMax: Distance cutoff for interactions.
:vartype RMax: float | FloatKey
"""
def __post_init__(self):
self.LMax: int | IntKey = IntKey(name='LMax', comment='Maximum L value. 1 is probably high enough.', default=1)
self.MetricsKey: Literal["training_loss", "validation_loss"] = MultipleChoiceKey(name='MetricsKey', comment="Which metric to use to generate the 'best' model.", default='validation_loss', choices=['training_loss', 'validation_loss'])
self.NumLayers: int | IntKey = IntKey(name='NumLayers', comment='Number of interaction layers in the NequIP neural network.', default=4)
self.RMax: float | FloatKey = FloatKey(name='RMax', comment='Distance cutoff for interactions.', gui_name='Distance cutoff:', default=3.5, unit='angstrom')
def __post_init__(self):
self.LearningRate: float | FloatKey = FloatKey(name='LearningRate', comment='Learning rate for the NequIP weight optimization', default=0.005, isenabled="glompo_machinelearning.machinelearning.backend == 'NequIP' # hide")
self.Model: Literal["Custom", "ModelFile"] = MultipleChoiceKey(name='Model', comment="How to specify the model for the NequIP backend. Either a Custom model can be made from scratch or an existing 'model.pth' file can be loaded to obtain the model settings.", default='Custom', choices=['Custom', 'ModelFile'])
self.ModelFile: str | Path | StringKey = PathStringKey(name='ModelFile', comment='Path to the model.pth file defining the model.', ispath=True, gui_type='file', isenabled="glompo_machinelearning.machinelearning.nequip.model == 'Model File' # hide")
self.UseRescalingFromLoadedModel: BoolType | BoolKey = BoolKey(name='UseRescalingFromLoadedModel', comment='When loading a model with LoadModel or NequiP%ModelFile do not recalculate the dataset rescaling but use the value from the loaded model.', default=True, isenabled="glompo_machinelearning.machinelearning.backend == 'NequIP' # hide")
self.Custom: ParAMSMachineLearning._MachineLearning._NequIP._Custom = self._Custom(name='Custom', comment='Specify a custom NequIP model.', isenabled="glompo_machinelearning.machinelearning.nequip.model == 'Custom' # hide")
[docs] class _Target(FixedBlock):
r"""
Target values for stopping training. If both the training and validation metrics are smaller than the specified values, the training will stop early. Supported by the MatGL and M3GNet backends.
:ivar Forces: Forces (as reported by the backend)
:vartype Forces: ParAMSMachineLearning._MachineLearning._Target._Forces
"""
[docs] class _Forces(FixedBlock):
r"""
Forces (as reported by the backend)
:ivar Enabled: Whether to use target values for forces.
:vartype Enabled: BoolType | BoolKey
:ivar MAE: MAE for forces (as reported by the backend).
:vartype MAE: float | FloatKey
"""
def __post_init__(self):
self.Enabled: BoolType | BoolKey = BoolKey(name='Enabled', comment='Whether to use target values for forces.', default=True)
self.MAE: float | FloatKey = FloatKey(name='MAE', comment='MAE for forces (as reported by the backend).', default=0.05, unit='eV/angstrom')
def __post_init__(self):
self.Forces: ParAMSMachineLearning._MachineLearning._Target._Forces = self._Forces(name='Forces', comment='Forces (as reported by the backend)')
def __post_init__(self):
self.Backend: Literal["Custom", "MatGL", "M3GNet", "MACE", "NEP", "NequIP", "Test"] = MultipleChoiceKey(name='Backend', comment='The backend to use. You must separately install the backend before running a training job.', default='M3GNet', choices=['Custom', 'MatGL', 'M3GNet', 'MACE', 'NEP', 'NequIP', 'Test'], hiddenchoices=['Custom', 'NEP', 'Test'], gui_type='literal choices')
self.CommitteeSize: int | IntKey = IntKey(name='CommitteeSize', comment='The number of independently trained ML potentials.', default=1)
self.LoadModel: str | Path | StringKey = PathStringKey(name='LoadModel', comment='Load a previously fitted model from a ParAMS results directory. A ParAMS results directory should contain two subdirectories ``optimization`` and ``settings_and_initial_data``. The loaded model defines the model source and architecture; fitting controls such as the learning rate and trainable layers still apply when supported by the backend.', ispath=True, gui_type='directory')
self.MaxEpochs: int | IntKey = IntKey(name='MaxEpochs', comment='Maximum number of epochs during the training.', default=1000)
self.RunAMSAtEnd: BoolType | BoolKey = BoolKey(name='RunAMSAtEnd', comment='Whether to run the (committee) ML potential through AMS at the end. This will create the energy/forces scatter plots for the final trained model.', gui_name='Run AMS at end:', default=True)
self.Custom: ParAMSMachineLearning._MachineLearning._Custom = self._Custom(name='Custom', comment='Set up a custom fitting program within ParAMS', hidden=True)
self.EarlyStopping: ParAMSMachineLearning._MachineLearning._EarlyStopping = self._EarlyStopping(name='EarlyStopping', comment='Stop training when the validation loss has not improved sufficiently. Supported by the MatGL and M3GNet backends.')
self.LossCoeffs: ParAMSMachineLearning._MachineLearning._LossCoeffs = self._LossCoeffs(name='LossCoeffs', comment='Modify the coefficients for the machine learning loss function. For backends that support weights, this is on top of the supplied dataset weights and sigmas.')
self.M3GNet: ParAMSMachineLearning._MachineLearning._M3GNet = self._M3GNet(name='M3GNet', comment='Options for M3GNet fitting.', isenabled="glompo_machinelearning.machinelearning.backend == 'M3GNet' and glompo_machinelearning.machinelearning.loadmodel == '' # hide")
self.MACE: ParAMSMachineLearning._MachineLearning._MACE = self._MACE(name='MACE', comment='Options for MACE fitting.', isenabled="glompo_machinelearning.machinelearning.backend == 'MACE' # hide")
self.MatGL: ParAMSMachineLearning._MachineLearning._MatGL = self._MatGL(name='MatGL', comment='Options for fitting MatGL QET and TensorNet potentials.', isenabled="glompo_machinelearning.machinelearning.backend == 'MatGL' # hide")
self.NEP: ParAMSMachineLearning._MachineLearning._NEP = self._NEP(name='NEP', comment='Experimental GPUMD NEP4 fitting. Training uses AMS_NEP_TRAIN_EXECUTABLE; production calculations use the native NEP worker installed by AMS. The default Foundation model fine-tunes the installed NEP89 potential for neutral, fully 3D-periodic systems.', hidden=True, isenabled="glompo_machinelearning.machinelearning.backend == 'NEP' and glompo_machinelearning.machinelearning.loadmodel == '' # hide")
self.NequIP: ParAMSMachineLearning._MachineLearning._NequIP = self._NequIP(name='NequIP', comment='Options for NequIP fitting.', isenabled="glompo_machinelearning.machinelearning.backend == 'NequIP' and glompo_machinelearning.machinelearning.loadmodel == '' # hide")
self.Target: ParAMSMachineLearning._MachineLearning._Target = self._Target(name='Target', comment='Target values for stopping training. If both the training and validation metrics are smaller than the specified values, the training will stop early. Supported by the MatGL and M3GNet backends.')
[docs] class _ParallelLevels(FixedBlock):
r"""
Distribution of threads/processes between the parallelization levels.
:ivar CommitteeMembers: Maximum number of committee member optimizations to run in parallel. If set to zero will take the minimum of MachineLearning%CommitteeSize and the number of available cores (NSCM)
:vartype CommitteeMembers: int | IntKey
:ivar Cores: Number of cores to use per committee member optimization. By default (0) the available cores (NSCM) divided equally among committee members. When using GPU offloading, consider setting this to 1.
:vartype Cores: int | IntKey
:ivar Jobs: Number of JobCollection jobs to run in parallel for each loss function evaluation.
:vartype Jobs: int | IntKey
:ivar Optimizations: Number of independent optimizers to run in parallel.
:vartype Optimizations: int | IntKey
:ivar ParameterVectors: Number of parameter vectors to try in parallel for each optimizer iteration. This level of parallelism can only be used with optimizers that support parallel optimization!
Default (0) will set this value to the number of cores on the system divided by the number of optimizers run in parallel, i.e., each optimizer will be given an equal share of the resources.
:vartype ParameterVectors: int | IntKey
:ivar Processes: Number of processes (MPI ranks) to spawn for each JobCollection job. This effectively sets the NSCM environment variable for each job.
A value of `-1` will disable explicit setting of related variables. We recommend a value of `1` in almost all cases. A value greater than 1 would only be useful if you parametrize DFTB with a serial optimizer and have very few jobs in the job collection.
:vartype Processes: int | IntKey
:ivar Threads: Number of threads to use for each of the processes. This effectively set the OMP_NUM_THREADS environment variable.
Note that the DFTB engine does not use threads, so the value of this variable would not have any effect. We recommend always leaving it at the default value of 1. Please consult the manual of the engine you are parameterizing.
A value of `-1` will disable explicit setting of related variables.
:vartype Threads: int | IntKey
"""
def __post_init__(self):
self.CommitteeMembers: int | IntKey = IntKey(name='CommitteeMembers', comment='Maximum number of committee member optimizations to run in parallel. If set to zero will take the minimum of MachineLearning%CommitteeSize and the number of available cores (NSCM)', gui_name='Number of parallel committee members:', default=1)
self.Cores: int | IntKey = IntKey(name='Cores', comment='Number of cores to use per committee member optimization. By default (0) the available cores (NSCM) divided equally among committee members. When using GPU offloading, consider setting this to 1.', gui_name='Processes (per Job):', default=0)
self.Jobs: int | IntKey = IntKey(name='Jobs', comment='Number of JobCollection jobs to run in parallel for each loss function evaluation.', gui_name='Jobs (per loss function evaluation):', default=0)
self.Optimizations: int | IntKey = IntKey(name='Optimizations', comment='Number of independent optimizers to run in parallel.', gui_name='Number of parallel optimizers:', default=1)
self.ParameterVectors: int | IntKey = IntKey(name='ParameterVectors', comment='Number of parameter vectors to try in parallel for each optimizer iteration. This level of parallelism can only be used with optimizers that support parallel optimization!\n\nDefault (0) will set this value to the number of cores on the system divided by the number of optimizers run in parallel, i.e., each optimizer will be given an equal share of the resources.', gui_name='Loss function evaluations (per optimizer):', default=0)
self.Processes: int | IntKey = IntKey(name='Processes', comment='Number of processes (MPI ranks) to spawn for each JobCollection job. This effectively sets the NSCM environment variable for each job.\n\nA value of `-1` will disable explicit setting of related variables. We recommend a value of `1` in almost all cases. A value greater than 1 would only be useful if you parametrize DFTB with a serial optimizer and have very few jobs in the job collection.', gui_name='Processes (per Job):', default=1)
self.Threads: int | IntKey = IntKey(name='Threads', comment='Number of threads to use for each of the processes. This effectively set the OMP_NUM_THREADS environment variable.\nNote that the DFTB engine does not use threads, so the value of this variable would not have any effect. We recommend always leaving it at the default value of 1. Please consult the manual of the engine you are parameterizing.\n\nA value of `-1` will disable explicit setting of related variables.', gui_name='Threads (per Process):', default=1)
def __post_init__(self):
self.EngineCollection: str | StringKey = StringKey(name='EngineCollection', comment='Path to (optional) JobCollection Engines YAML file.', default='job_collection_engines.yaml')
self.JobCollection: str | StringKey = StringKey(name='JobCollection', comment='Path to JobCollection YAML file.', default='job_collection.yaml')
self.ResultsDirectory: str | Path | StringKey = PathStringKey(name='ResultsDirectory', comment='Directory in which output files will be created.', gui_name='Working directory: ', default='results', ispath=True)
self.Task: Literal["Optimization", "GenerateReference", "SinglePoint", "Sensitivity", "MachineLearning"] = MultipleChoiceKey(name='Task', comment='Task to run.\n\nAvailable options:\n•MachineLearning: Optimization for machine learning models.\n•Optimization: Global optimization powered by GloMPO\n•Generate Reference: Run jobs with reference engine to get reference values\n•Single Point: Evaluate the current configuration of jobs, training data, and parameters\n•Sensitivity: Measure the sensitivity of the loss function to each of the active parameters', default='Optimization', choices=['Optimization', 'GenerateReference', 'SinglePoint', 'Sensitivity', 'MachineLearning'])
self.DataSet: ParAMSMachineLearning._DataSet = self._DataSet(name='DataSet', comment='Configuration settings for each data set in the optimization.', unique=False, gui_type='Repeat at least once')
self.MachineLearning: ParAMSMachineLearning._MachineLearning = self._MachineLearning(name='MachineLearning', comment='Options for Task MachineLearning.')
self.ParallelLevels: ParAMSMachineLearning._ParallelLevels = self._ParallelLevels(name='ParallelLevels', comment='Distribution of threads/processes between the parallelization levels.', gui_name='Parallelization distribution: ')
