Add SevenNet to training (#662)

Co-authored-by: Jacob Wilkins <46597752+oerc0122@users.noreply.github.com>
Co-authored-by: Elliott Kasoar <45317199+ElliottKasoar@users.noreply.github.com>

Author: 3 people

docs/source/user_guide/command_line.rst

@@ -659,7 +659,7 @@ Training and fine-tuning MLIPs
 ------------------------------
 
 .. note::
-    Currently only MACE and Nequip models are supported.
+    Currently MACE, Nequip, and SevenNet models are supported.
 
 Models can be trained by passing an archictecture and an archictecture specific configuration file as options to the ``janus train`` command. The configuration file will be passed to the corresponding MLIPs command line interface. For example to train a MACE MLIP:
 
@@ -692,14 +692,19 @@ For MACE, training will create ``logs``, ``checkpoints`` and ``results`` sub-dir
 Instructions for writing a MACE ``config.yml`` file can be found in the `MACE Readme <https://github.com/ACEsuit/mace?tab=readme-ov-file#training>`_ and the `MACE run_train CLI <https://github.com/ACEsuit/mace/blob/main/mace/cli/run_train.py>`_.
 
 
-Training Nequip MLIPS
+Training Nequip MLIPs
 +++++++++++++++++++++
 
 Configuration of Nequip training is outlined in the `Nequip user guide <https://nequip.readthedocs.io/en/latest/guide/guide.html>`_. In particular note that the configuration file must have a ``.yaml`` extension.
 
 The results directory contents depends on the options selected in the configuration file, but may typically contain model checkpoint, ``.ckpt``, files and a metrics directory.
 
 
+Training SevenNet MLIPs
++++++++++++++++++++++++
+
+The `SevenNet documentation <https://sevennet.readthedocs.io/en/latest/>`_ contains information on training SevenNet MLIPs. The SevenNet `tutorial repository <https://github.com/MDIL-SNU/sevennet_tutorial/tree/main>`_ also contains some example ```.yaml``` configuration files for training and fine-tuning.
+
 Preprocessing training data
 ----------------------------
 

janus_core/cli/train.py

@@ -122,6 +122,25 @@ def train(
                     """Fine-tuning requested but there is no checkpoint or
                     package specified in your config."""
                 )
+        case "sevennet":
+            continue_section = config["train"].get("continue")
+            if continue_section is None and fine_tune:
+                raise ValueError(
+                    """Fine-tuning requested but there is no continue
+                    section in your config."""
+                )
+                model = continue_section.get("checkpoint")
+                if model is None:
+                    raise ValueError(
+                        "No model specified as a checkpoint for fine-tuning."
+                    )
+            if not fine_tune and continue_section is not None:
+                raise ValueError(
+                    """Fine-tuning not requested but a continue
+                    section is in your config. Please use
+                    --fine-tune"""
+                )
+
         case _:
             raise ValueError(f"Unsupported Architecture ({arch})")
 

janus_core/training/train.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+from argparse import ArgumentParser
 from typing import Any
 
 import yaml
@@ -93,6 +94,15 @@ def train(
                     )
                 foundation_model = model["checkpoint_path"]
 
+        case "sevennet":
+            from sevenn.main.sevenn import cmd_parser_train, run
+
+            parser = ArgumentParser()
+            cmd_parser_train(parser)
+            mlip_args = parser.parse_args(
+                [str(mlip_config), "--working_dir", str(file_prefix), "-s"]
+            )
+
         case _:
             raise ValueError(f"{arch} is currently unsupported in train.")
 

tests/data/sevennet_fine_tune.yml

@@ -0,0 +1,92 @@
+model:
+  chemical_species: auto
+
+  cutoff: 2.0
+  irreps_manual:
+    - 128x0e
+    - 128x0e+64x1e+32x2e+32x3e
+    - 128x0e+64x1e+32x2e+32x3e
+    - 128x0e+64x1e+32x2e+32x3e
+    - 128x0e+64x1e+32x2e+32x3e
+    - 128x0e
+  channel: 128
+  lmax: 3
+  num_convolution_layer: 5
+  is_parity: false
+  radial_basis:
+    radial_basis_name: bessel
+    bessel_basis_num: 8
+  cutoff_function:
+    cutoff_function_name: poly_cut
+    poly_cut_p_value: 6
+
+  act_radial: silu
+  weight_nn_hidden_neurons:
+    - 64
+    - 64
+  act_scalar:
+    e: silu
+    o: tanh
+  act_gate:
+    e: silu
+    o: tanh
+
+  train_denominator: false
+  train_shift_scale: false
+  use_bias_in_linear: false
+
+  readout_as_fcn: false
+  self_connection_type: linear
+  interaction_type: nequip
+
+train:
+    random_seed: 1
+    is_train_stress: True
+    epoch: 1
+
+
+
+    optimizer: 'adam'
+    optim_param:
+        lr: 0.005
+    scheduler: 'exponentiallr'
+    scheduler_param:
+        gamma: 0.99
+
+    force_loss_weight: 0.1
+    stress_loss_weight: 1e-06
+
+    per_epoch: 1
+
+
+
+    error_record:
+        - ['Energy', 'RMSE']
+        - ['Force', 'RMSE']
+        - ['Stress', 'RMSE']
+        - ['TotalLoss', 'None']
+
+    continue:
+        reset_optimizer: True
+        reset_scheduler: True
+        reset_epoch: True
+        checkpoint: 'tests/models/extra/SevenNet_l3i5.pth'
+
+        use_statistic_values_of_checkpoint: True
+
+data:
+    batch_size: 4
+    data_divide_ratio: 0.1
+
+    shift: 'per_atom_energy_mean'
+    scale: 'force_rms'
+
+
+
+    data_format: 'ase'
+    data_format_args:
+        index: ':'
+
+
+
+    load_dataset_path: ['tests/data/mlip_train.xyz']

tests/data/sevennet_train.yml

@@ -0,0 +1,55 @@
+model:
+    chemical_species: 'Auto'
+    cutoff: 2.0
+    channel: 4
+    lmax: 1
+    num_convolution_layer: 1
+
+    weight_nn_hidden_neurons: [4, 4]
+    radial_basis:
+        radial_basis_name: 'bessel'
+        bessel_basis_num: 8
+    cutoff_function:
+        cutoff_function_name: 'poly_cut'
+        poly_cut_p_value: 6
+
+    act_gate: {'e': 'silu', 'o': 'tanh'}
+    act_scalar: {'e': 'silu', 'o': 'tanh'}
+
+    is_parity: False
+
+    self_connection_type: 'nequip'
+
+    conv_denominator: "avg_num_neigh"
+    train_denominator: False
+    train_shift_scale: False
+
+train:
+    random_seed: 1
+    is_train_stress: True
+    epoch: 2
+    optimizer: 'adam'
+    optim_param:
+        lr: 0.005
+    scheduler: 'exponentiallr'
+    scheduler_param:
+        gamma: 0.99
+    force_loss_weight: 0.1
+    stress_loss_weight: 1e-06
+    per_epoch: 1
+    error_record:
+        - ['Energy', 'RMSE']
+        - ['Force', 'RMSE']
+        - ['Stress', 'RMSE']
+        - ['TotalLoss', 'None']
+
+data:
+    batch_size: 4
+    data_divide_ratio: 0.1
+
+    shift: 'per_atom_energy_mean'
+    scale: 'force_rms'
+    data_format: 'ase'
+    data_format_args:
+        index: ':'
+    load_dataset_path: ['tests/data/mlip_train.xyz']

tests/models/extra_models.py

@@ -5,6 +5,24 @@
 from argparse import ArgumentParser
 from pathlib import Path
 from urllib.request import urlretrieve
+from warnings import warn
+
+
+def try_retrieve(url: str, filename: Path):
+    """Attempt to retrieve from url.
+
+    Patameters
+    ----------
+    url
+        The url to attempt to retrieve from.
+    filename
+        The local filename of the retrieved data.
+    """
+    try:
+        urlretrieve(url, filename)
+    except Exception as e:
+        warn(f"Unable to retrieve from {url} because of:\n  {e}", stacklevel=2)
+
 
 if __name__ == "__main__":
     parser = ArgumentParser(
@@ -14,7 +32,13 @@
     args = parser.parse_args()
 
     args.path.mkdir(parents=True, exist_ok=True)
-    urlretrieve(
+
+    try_retrieve(
         "https://zenodo.org/records/16980200/files/NequIP-MP-L-0.1.nequip.zip",
-        filename=args.path / "NequIP-MP-L-0.1.nequip.zip",
+        args.path / "NequIP-MP-L-0.1.nequip.zip",
+    )
+
+    try_retrieve(
+        "https://github.com/MDIL-SNU/SevenNet/raw/dff008ac9c53d368b5bee30a27fa4bdfd73f19b2/sevenn/pretrained_potentials/SevenNet_l3i5/checkpoint_l3i5.pth",
+        args.path / "SevenNet_l3i5.pth",
     )