Add grace to train

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:
 
@@ -699,6 +699,37 @@ Configuration of Nequip training is outlined in the `Nequip user guide <https://
 
 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.
 
+.. note::
+    Different architectures may have different restrictions or features. For example Nequip requires YAML files to be written as ``.yaml`` rather than ``.yml``. See the sections below for specific archictecture guidance.
+
+Foundational models can also be fine-tuned, by passing the ``--fine-tune`` option:
+
+.. code-block:: bash
+
+    janus train mace --mlip-config /path/to/mace/fine/tuning/config.yml --fine-tune
+
+By default the output of training or fine-tuning will be in the ``./janus_results`` directory. This directory's contents varies in structure depending on the architecture being trained and whether fine-tuning is being conducted. However as with other commands a log file, ``train-log.yml``, and summary file, ``train-summary.yml``, will be generated in ``./janus_results`` by default.
+
+Training MACE MLIPs
++++++++++++++++++++
+
+For MACE, training will create ``logs``, ``checkpoints`` and ``results`` sub-directories, as well as saving the trained model, and a compiled version of the model.
+
+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
++++++++++++++++++++++
+
+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,45 @@ 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 "grace":
+            if "potential" not in config:
+                raise ValueError("No potential is specified in you config.")
+
+            if fine_tune:
+                model = config["potential"].get("finetune_foundation_model")
+                if model is None:
+                    raise ValueError(
+                        """Fine-tuning was requested but your conifg
+                        does not contains a finetune_foundation_model"""
+                    )
+            elif "finetune_foundation_model" in config["potential"]:
+                raise ValueError(
+                    """Fine-tuning not requested but finetune_foundation_model
+                    is in your config. Please use --fine-tune.
+                    """
+                )
+
         case _:
             raise ValueError(f"Unsupported Architecture ({arch})")
 

janus_core/training/train.py

@@ -2,6 +2,8 @@
 
 from __future__ import annotations
 
+from argparse import ArgumentParser
+from pathlib import Path
 from typing import Any
 
 import yaml
@@ -93,6 +95,20 @@ 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 "grace":
+            from tensorpotential.cli.gracemaker import main as run
+
+            mlip_args = [str(mlip_config)]
+
         case _:
             raise ValueError(f"{arch} is currently unsupported in train.")
 
@@ -120,6 +136,11 @@ def train(
 
     run(mlip_args)
 
+    if arch == "grace" and (Path.cwd() / "seed").exists():
+        # Gracemaker always works in ./seed.
+        file_prefix.mkdir(parents=True, exist_ok=True)
+        (Path.cwd() / "seed").rename(file_prefix.resolve() / "seed")
+
     if logger:
         logger.info("Training complete")
     if tracker:

tests/data/grace_fine_tune.yml

@@ -0,0 +1,35 @@
+seed: 42
+cutoff: 6
+
+data:
+  filename: "tests/data/mlip_train.pkl.gz"
+  reference_energy: 0
+
+potential:
+  finetune_foundation_model: "GRACE-1L-OAM"
+
+fit:
+  loss:
+    energy:
+      type: huber
+      weight: 17
+      delta: 0.01
+    forces:
+      type: huber
+      weight: 32.
+      delta: 0.01
+
+  maxiter: 1 # Max number of optimization epochs
+  optimizer: Adam
+  opt_params: { learning_rate: 0.008, use_ema: True, ema_momentum: 0.99,  weight_decay: 1.e-20, clipnorm: 1.0}
+  scheduler: cosine_decay # scheduler for learning-rate reduction during training
+  scheduler_params: {"minimal_learning_rate": 0.0001}
+
+  batch_size: 32 # Important hyperparameter for Adam and irrelevant (but must be) for L-BFGS-B/BFGS
+  test_batch_size: 200 # test batch size (optional)
+
+  jit_compile: True  # for XLA compilation, must be used in almost all cases
+  train_max_n_buckets: 10  ## max number of buckets in train set
+  test_max_n_buckets: 3  ## same for test
+
+  checkpoint_freq: 10 # frequency for **REGULAR** checkpoints.

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

@@ -14,7 +14,13 @@
     args = parser.parse_args()
 
     args.path.mkdir(parents=True, exist_ok=True)
+
     urlretrieve(
         "https://zenodo.org/records/16980200/files/NequIP-MP-L-0.1.nequip.zip",
         filename=args.path / "NequIP-MP-L-0.1.nequip.zip",
     )
+
+    urlretrieve(
+        "https://github.com/MDIL-SNU/SevenNet/raw/dff008ac9c53d368b5bee30a27fa4bdfd73f19b2/sevenn/pretrained_potentials/SevenNet_l3i5/checkpoint_l3i5.pth",
+        filename=args.path / "SevenNet_l3i5.pth",
+    )

tests/models/extra_models.sh

@@ -0,0 +1,6 @@
+if [ ! -d tests/models/extra ]
+then
+    mkdir tests/models/extra
+fi
+
+(cd tests/models/extra; curl --output NequIP-MP-L-0.1.nequip.zip https://zenodo.org/records/16980200/files/NequIP-MP-L-0.1.nequip.zip)