6.6. Task: MachineLearning
**************************

Set "Task MachineLearning" to fit machine learning potentials (ML
potentials). In ParAMS, all supported types of ML potentials can be
trained as committee models that provide an estimate of the
uncertainty of predicted energies and forces during production
simulations.

Training ML Potentials through ParAMS requires

* the job collection, and

* training and validation sets

You can construct these using the results importers, just as for
ReaxFF and DFTB parametrization.

Note:

  Unlike ReaxFF and DFTB parametrization, no Parameter Interface is
  needed. This is because ML potentials usually contain many thousands
  of parameters. It is typically not useful to manually control the
  values and ranges for all of those parameters.

You also need to specify which Backend to use, for example M3GNet.


6.6.1. Requirements for job collection and data sets
====================================================

The machine learning potentials are trained in a quite different way
from how ParAMS trains ReaxFF and DFTB.


6.6.1.1. Only singlepoint calculations in the job collection
------------------------------------------------------------

For ML potentials, **only singlepoint calculations may enter the job
collection**. The original reference job can still be of any type
(geometry optimization, PES Scan, ...).

Example: if you import a DFT-calculated bond scan (PES Scan), you must
import it using the "Add PESScan Singlepoints" option, not "add
singlejob with Task='PESScan'".

Any jobs in the job collection with Task different from "SinglePoint"
will be ignored.


6.6.1.2. Only single extractors in the training and validation sets
-------------------------------------------------------------------

Similarly, for the training and validation sets, the expressions can
only contain **one** extractor acting on a single job. This means that
you cannot train reaction energies. Instead, you can (and should)
train the total energy. As a result it is extra important that all
reference data was calculated using a single level of theory.

When training forces, you must extract all force components from the
job. However, depending on the backend, you may be able to set the
force weights.

For task MachineLearning, only a small set of extractors (that act on
singlepoint jobs) are supported:

* "energy"

* "forces"

Examples:

+----------------------------------------------------------------------------------+------------+------------+
| Expression                                                                       | Task Opti  | Task Mach  |
|                                                                                  | mization   | ineLearni  |
|                                                                                  |            | ng         |
|==================================================================================|============|============|
| "energy("job")"                                                                  | OK         | OK         |
+----------------------------------------------------------------------------------+------------+------------+
| "forces("job")"                                                                  | OK         | OK         |
+----------------------------------------------------------------------------------+------------+------------+
| "energy("job1")-energy("job2")"                                                  | OK         | Not OK     |
+----------------------------------------------------------------------------------+------------+------------+
| "forces("job", 3, 2)"                                                            | OK         | Not OK     |
+----------------------------------------------------------------------------------+------------+------------+
| "cell_volume("job")"                                                             | OK         | Not OK     |
+----------------------------------------------------------------------------------+------------+------------+

Expressions that do not follow the above requirements will be ignored
during the ML training, but they will still be stored on disk. This
means that if you after training your ML potential switch to the
ParAMS SinglePoint Task, you can use any expressions and job tasks to
test/validate/benchmark your trained potential.


6.6.1.3. The engine settings must be the same for all jobs
----------------------------------------------------------

When you train for example DFTB, you can have different engine
settings for different jobs. For example, you might want the k-space
sampling to be different depending on the system.

However, when training machine learning potentials, you cannot set any
job-dependent (structure-dependent) engine settings. Every job
(structure) will use the same settings.


6.6.2. Machine Learning Input Structure
=======================================

The input for the ParAMS "Task MachineLearning" is structured as
follows:

"MachineLearning" has multiple backends that can be selected through
the "Backend" key. Each backend has a corresponding block (the same
name as the value of the "Backend" key) with settings specific to that
backend. Additionally, there are several shared keywords, such as
"MaxEpochs" that modify the behavior of the backends in the same way.

Each backend might support multiple models and has a corresponding
block (the same name as the value of the "Model" key) with settings
specific to that model. For example the number of layers or how to
initialize the parameters.

Some models consist of only a single key rather than a block. For
example when a backend supports loading some file that contains model
settings and parameters.

Any number of settings may exist at the top level of a backend block
that are appropriate for all models.

The "MachineLearning%LoadModel" loads a previously fitted model from a
ParAMS results directory.  The ParAMS results directory must contain
the two subdirectories "optimization" and "settings_and_initial_data".
Enabling "MachineLearning%LoadModel" enforces the same "Backend" and
"CommitteeSize" as in the previous job and will ignore the model keys.
Instead it reads them from the previous ParAMS calculation.  Any
settings in the model blocks are ignored.  If any settings in the
backend blocks are incompatible with the loaded model then ParAMS will
crash or behave unexpectedly.

The exact same backend and model settings are used for every committee
member no matter the "CommitteeSize", although the models can still be
different due to stochastic effects (e.g. random parameters or a
stochastic optimization algorithm). When using "LoadModel" the
committee from the previous calculation is used.

Set "RunAMSAtEnd" to run the job collection with the newly trained
model once training is completed. This will provide additional output
such as scatter plots of prediction against reference values.

Tip:

  Learn using the ParAMS input for "Task MachineLearning" from the
  tutorials.

"MachineLearning"
   Type:
      Block

   Description:
      Options for Task MachineLearning.

   "Backend"
      Type:
         Multiple Choice

      Default value:
         M3GNet

      Options:
         [M3GNet, MACE, NequIP]

      Description:
         The backend to use. You must separately install the backend
         before running a training job.

   "MaxEpochs"
      Type:
         Integer

      Default value:
         1000

      Description:
         Set the maximum number of epochs a backend should perform.

   "LossCoeffs"
      Type:
         Block

      Description:
         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.

      "AverageForcePerAtom"
         Type:
            Bool

         Default value:
            No

         Description:
            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.

      "Energy"
         Type:
            Float

         Default value:
            10.0

         GUI name:
            Energy coefficient

         Description:
            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.

      "Forces"
         Type:
            Float

         Default value:
            1.0

         GUI name:
            Forces coefficient

         Description:
            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.

   "Target"
      Type:
         Block

      Description:
         Target values for stopping training. If both the training and
         validation metrics are smaller than the specified values, the
         training will stop early. Only supported by the M3GNet
         backend.

      "Forces"
         Type:
            Block

         Description:
            Forces (as reported by the backend)

         "Enabled"
            Type:
               Bool

            Default value:
               Yes

            Description:
               Whether to use target values for forces.

         "MAE"
            Type:
               Float

            Default value:
               0.05

            Unit:
               eV/angstrom

            Description:
               MAE for forces (as reported by the backend).

   "LoadModel"
      Type:
         String

      Description:
         Load a previously fitted model from a ParAMS results
         directory. A ParAMS results directory should contain two
         subdirectories "optimization" and
         "settings_and_initial_data". This option ignores all settings
         inside model blocks.

   "CommitteeSize"
      Type:
         Integer

      Default value:
         1

      Description:
         The number of independently trained ML potentials.

   "RunAMSAtEnd"
      Type:
         Bool

      Default value:
         Yes

      GUI name:
         Run AMS at end

      Description:
         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.


6.6.3. Backends: M3GNet, NequIP, ...
====================================


6.6.3.1. Installation
---------------------

The ML backends are not included by default with AMS or ParAMS, as
they can be quite large. Before you can train an ML potential, you
need to install the corresponding backend either through the AMS
package manager or manually.

Tip:

  Before training a custom model with ParAMS, we recommend that you
  first test the ML backend in a production (for example, molecular
  dynamics or geometry optimization) simulation with some already
  created parameters. For example, follow the M3GNet GUI tutorial to
  make sure that the M3GNet backend has been installed correctly.


6.6.3.2. M3GNet
---------------

"MachineLearning"
   Type:
      Block

   Description:
      Options for Task MachineLearning.

   "M3GNet"
      Type:
         Block

      Description:
         Options for M3GNet fitting.

      "Custom"
         Type:
            Block

         Description:
            Specify a custom M3GNet model.

         "Cutoff"
            Type:
               Float

            Default value:
               5.0

            Unit:
               angstrom

            Description:
               Cutoff radius of the graph

         "MaxL"
            Type:
               Integer

            Default value:
               3

            Description:
               Include spherical components up to order MaxL. Higher
               gives a better angular resolution, but increases
               computational cost substantially.

         "MaxN"
            Type:
               Integer

            Default value:
               3

            Description:
               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.

         "NumBlocks"
            Type:
               Integer

            Default value:
               3

            GUI name:
               Number of convolution blocks:

            Description:
               Number of convolution blocks.

         "NumNeurons"
            Type:
               Integer

            Default value:
               64

            GUI name:
               Number of neurons per layer

            Description:
               Number of neurons in each layer.

         "ThreebodyCutoff"
            Type:
               Float

            Default value:
               4.0

            Unit:
               angstrom

            Description:
               Cutoff radius of the three-body interaction.

      "LearningRate"
         Type:
            Float

         Default value:
            0.001

         Description:
            Learning rate for the M3GNet weight optimization.

      "Model"
         Type:
            Multiple Choice

         Default value:
            UniversalPotential

         Options:
            [UniversalPotential, Custom, ModelDir]

         Description:
            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.

      "ModelDir"
         Type:
            String

         Description:
            Path to the directory defining the model. This folder
            should contain the files: 'checkpoint',
            'm3gnet.data-00000-of-00001', ' m3gnet.index' and
            'm3gnet.json'

      "UniversalPotential"
         Type:
            Block

         Description:
            Settings for (transfer) learning with the M3GNet Universal
            Potential.

         "Featurizer"
            Type:
               Bool

            Default value:
               No

            GUI name:
               Train featurizer

            Description:
               Train the Featurizer layer of the M3GNet universal
               potential.

         "Final"
            Type:
               Bool

            Default value:
               Yes

            GUI name:
               Train final layer

            Description:
               Train the Final layer of the M3GNet universal
               potential.

         "GraphLayer1"
            Type:
               Bool

            Default value:
               No

            GUI name:
               Train layer 1 - graph

            Description:
               Train the first Graph layer of the M3GNet universal
               potential.

         "GraphLayer2"
            Type:
               Bool

            Default value:
               No

            GUI name:
               Train layer 2 - graph

            Description:
               Train the second Graph layer of the M3GNet universal
               potential.

         "GraphLayer3"
            Type:
               Bool

            Default value:
               Yes

            GUI name:
               Train layer 3 - graph

            Description:
               Train the third Graph layer of the M3GNet universal
               potential.

         "ThreeDInteractions1"
            Type:
               Bool

            Default value:
               No

            GUI name:
               Train layer 1 - 3D interactions

            Description:
               Train the first ThreeDInteractions (three-body terms)
               layer of the M3GNet universal potential.

         "ThreeDInteractions2"
            Type:
               Bool

            Default value:
               No

            GUI name:
               Train layer 2 - 3D interactions

            Description:
               Train the second ThreeDInteractions (three-body terms)
               layer of the M3GNet universal potential.

         "ThreeDInteractions3"
            Type:
               Bool

            Default value:
               Yes

            GUI name:
               Train layer 3 - 3D interactions

            Description:
               Train the third ThreeDInteractions (three-body terms)
               layer of the M3GNet universal potential.

M3GNet produces the parameter directory *<calculation
name>.results/optimization/m3gnet/results/model* which contains the
parameterized model and can be used with the MLPotential engine. Set
*Backend M3GNet* and *ParameterDir* to the path of the deployed model.

The M3GNet universal potential has the following
architecture/structure:

+----------------------------------------------------+----------------------------------------------------+
| Layer (type)                                       | Param #                                            |
|====================================================|====================================================|
| radius_cutoff_graph_converter                      | 0 (unused)                                         |
| (RadiusCutoffGraphConverter)                       |                                                    |
+----------------------------------------------------+----------------------------------------------------+
| graph_featurizer (GraphFeaturizer)                 | 6080                                               |
+----------------------------------------------------+----------------------------------------------------+
| graph_update_func (GraphUpdateFunc)                | 192                                                |
+----------------------------------------------------+----------------------------------------------------+
| spherical_bessel_with_harmonics                    | 0                                                  |
| (SphericalBesselWithHarmonics)                     |                                                    |
+----------------------------------------------------+----------------------------------------------------+
| three_d_interaction (ThreeDInteraction)            | 1737                                               |
+----------------------------------------------------+----------------------------------------------------+
| three_d_interaction_1 (ThreeDInteraction)          | 1737                                               |
+----------------------------------------------------+----------------------------------------------------+
| three_d_interaction_2 (ThreeDInteraction)          | 1737                                               |
+----------------------------------------------------+----------------------------------------------------+
| graph_network_layer (GraphNetworkLayer)            | 66432                                              |
+----------------------------------------------------+----------------------------------------------------+
| graph_network_layer_1 (GraphNetworkLayer)          | 66432                                              |
+----------------------------------------------------+----------------------------------------------------+
| graph_network_layer_2 (GraphNetworkLayer)          | 66432                                              |
+----------------------------------------------------+----------------------------------------------------+
| pipe_24 (Pipe)                                     | 16770                                              |
+----------------------------------------------------+----------------------------------------------------+
| atom_ref_2 (AtomRef)                               | 0                                                  |
+----------------------------------------------------+----------------------------------------------------+

Total params: 227,549


6.6.3.3. MACE
-------------

The options for MACE are:

"MachineLearning"
   Type:
      Block

   Description:
      Options for Task MachineLearning.

   "MACE"
      Type:
         Block

      Description:
         Options for MACE fitting.

      "Custom"
         Type:
            Block

         Description:
            Specify a custom MACE model.

         "DataType"
            Type:
               Multiple Choice

            Default value:
               float64

            Options:
               [float32, float64]

            Description:
               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.

         "LMax"
            Type:
               Integer

            Default value:
               1

            GUI name:
               Max spherical harmonic order

            Description:
               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.

         "NumChannels"
            Type:
               Integer

            Default value:
               128

            GUI name:
               Number of channels

            Description:
               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.

         "RMax"
            Type:
               Float

            Default value:
               5.0

            Unit:
               angstrom

            GUI name:
               Distance cutoff

            Description:
               Distance cutoff for interactions.

      "Foundation"
         Type:
            Block

         Description:
            Settings for (transfer) learning with the MACE foundation
            model.

         "Type"
            Type:
               Multiple Choice

            Default value:
               MACE-MPA-0

            Options:
               [MACE-MPA-0, MACE-MP-0-Large, MACE-MP-0-Medium, MACE-
               MP-0-Small]

            Description:
               Which MACE foundation model to use.

      "LearningRate"
         Type:
            Float

         Default value:
            0.01

         Description:
            Learning rate for the MACE weight optimization

      "LossCoeffs"
         Type:
            Block

         Description:
            Modify the coefficients for the machine learning loss
            function, by applying scale factors to the
            MachineLearning%LossCoeffs values.

         "EnergyScaleFactor"
            Type:
               Float

            Default value:
               0.1

            Description:
               Scale factor to apply to the energy loss coefficient,
               i.e. EnergyScaleFactor *
               MachineLearning%LossCoeffs%Energy.

         "ForcesScaleFactor"
            Type:
               Float

            Default value:
               100.0

            Description:
               Scale factor to apply to the forces loss coefficient
               i.e. ForcesScaleFactor *
               MachineLearning%LossCoeffs%Forces.

      "Model"
         Type:
            Multiple Choice

         Default value:
            Foundation

         Options:
            [Foundation, Custom, ModelFile]

         Description:
            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.

      "ModelFile"
         Type:
            String

         Description:
            Path to the .model file defining the model.

      "StageTwo"
         Type:
            Block

         Description:
            Settings for stage two of training.

         "Enabled"
            Type:
               Bool

            Default value:
               Yes

            GUI name:
               Stage two enabled

            Description:
               Whether to enable stage two training, defaults to
               "True".

         "LearningRate"
            Type:
               Float

            Default value:
               0.001

            GUI name:
               Stage two learning rate

            Description:
               Learning rate for the MACE weight optimization for
               stage two training

         "LossCoeffs"
            Type:
               Block

            Description:
               Modify the coefficients for the machine learning loss
               function for stage two training, by applying scale
               factors to the MachineLearning%LossCoeffs values.

            "EnergyScaleFactor"
               Type:
                  Float

               Default value:
                  100.0

               GUI name:
                  Stage two energy scale factor

               Description:
                  Scale factor to apply to the energy loss coefficient
                  for stage two training, i.e. EnergyScaleFactor *
                  MachineLearning%LossCoeffs%Energy.

            "ForcesScaleFactor"
               Type:
                  Float

               Default value:
                  100.0

               GUI name:
                  Stage two forces scale factor

               Description:
                  Scale factor to apply to the forces loss coefficient
                  for stage two training i.e. ForcesScaleFactor *
                  MachineLearning%LossCoeffs%Forces.

         "Start"
            Type:
               Float

            Default value:
               0.8

            GUI name:
               Stage two start

            Description:
               When to start stage two training, as a proportion of
               the MachineLearning%MaxEpochs


6.6.3.4. NequIP
---------------

Important:

  Training NequIP potentials with ParAMS is not a fully supported
  feature. To use NequIP with AMS, or to train NequIP with ParAMS, you
  need to manually install it into the AMS Python environment.SCM does
  not provide any packages for NequIP and cannot provide support for
  the installation. But we have compiled some helpful tips in the
  Engine ASE documentation that may help you with the installation.

The options for NequIP are:

"MachineLearning"
   Type:
      Block

   Description:
      Options for Task MachineLearning.

   "NequIP"
      Type:
         Block

      Description:
         Options for NequIP fitting.

      "Custom"
         Type:
            Block

         Description:
            Specify a custom NequIP model.

         "LMax"
            Type:
               Integer

            Default value:
               1

            Description:
               Maximum L value. 1 is probably high enough.

         "MetricsKey"
            Type:
               Multiple Choice

            Default value:
               validation_loss

            Options:
               [training_loss, validation_loss]

            Description:
               Which metric to use to generate the 'best' model.

         "NumLayers"
            Type:
               Integer

            Default value:
               4

            Description:
               Number of interaction layers in the NequIP neural
               network.

         "RMax"
            Type:
               Float

            Default value:
               3.5

            Unit:
               angstrom

            GUI name:
               Distance cutoff

            Description:
               Distance cutoff for interactions.

      "LearningRate"
         Type:
            Float

         Default value:
            0.005

         Description:
            Learning rate for the NequIP weight optimization

      "Model"
         Type:
            Multiple Choice

         Default value:
            Custom

         Options:
            [Custom, ModelFile]

         Description:
            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.

      "ModelFile"
         Type:
            String

         Description:
            Path to the model.pth file defining the model.

      "UseRescalingFromLoadedModel"
         Type:
            Bool

         Default value:
            Yes

         Description:
            When loading a model with LoadModel or NequiP%ModelFile do
            not recalculate the dataset rescaling but use the value
            from the loaded model.

NequIP produces the file *<calculation
name>.results/optimization/nequip/results/model.pth* which contains
the deployed model and can be used with the MLPotential engine. Set
*Backend NequIP* and *ParameterFile* to the path of the deployed
model.


6.6.4. ML Parallelization
=========================

Parallelization options can be set with "ParallelLevels". Note that
*Task MachineLearning* does not perform AMS jobs during optimization,
so the parallelization options are different.

Select the maximum number of parallel committee members with
*CommitteeMembers* or set it to zero to run all committee members in
parallel (up to the maximum number of cores or the *NSCM* environment
variable). Select the number of cores each committee is allowed to use
with *Cores* or set it to zero (default) to evenly distribute the
available cores over the committee members running in parallel.

Some backends may spawn additional threads for database management,
but they should not be using substantial CPU time. GPU offloading is
supported through TensorFlow or PyTorch depending on the backend.
Currently there are no settings available in ParAMS for GPU
offloading; the backends use GPU resources according to their
documentation.

"ParallelLevels"
   Type:
      Block

   GUI name:
      Parallelization distribution:

   Description:
      Distribution of threads/processes between the parallelization
      levels.

   "CommitteeMembers"
      Type:
         Integer

      Default value:
         1

      GUI name:
         Number of parallel committee members

      Description:
         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)

   "Cores"
      Type:
         Integer

      Default value:
         0

      GUI name:
         Processes (per Job)

      Description:
         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.
