Finetuning

docs/src/software/metatensor.sec

@@ -14,6 +14,7 @@ training and evaluating ML models.
 - examples/learn-tensors-with-mcov/learn-tensors-with-mcov
 - examples/pet-mad/pet-mad
 - examples/pet-mad-nc/pet-mad-nc
+- examples/pet-finetuning/pet-ft
 - examples/flashmd/flashmd-demo
 - examples/shiftml/shiftml-example
 - examples/hamiltonian-qm7/hamiltonian-qm7

docs/src/topics/ml-models.sec

@@ -13,6 +13,7 @@ data.
 - examples/learn-tensors-with-mcov/learn-tensors-with-mcov
 - examples/pet-mad/pet-mad
 - examples/pet-mad-nc/pet-mad-nc
+- examples/pet-finetuning/pet-ft
 - examples/flashmd/flashmd-demo
 - examples/shiftml/shiftml-example
 - examples/hamiltonian-qm7/hamiltonian-qm7

examples/pet-finetuning/.gitignore

@@ -0,0 +1,14 @@
+*.data
+*.pt
+*.ckpt
+log.lammps
+trajectory.xyz
+tmp*
+
+extensions/*
+
+outputs/
+data/ethanol_train.xyz
+data/ethanol_val.xyz
+data/ethanol_test.xyz
+output.xyz

examples/pet-finetuning/README.rst

@@ -0,0 +1,4 @@
+PET-MAD finetuning tutorial
+===========================
+
+An example of finetuning the PET-MAD universal machine-learning potential.

examples/pet-finetuning/basic_ft_options.yaml

@@ -0,0 +1,38 @@
+seed: 42
+
+architecture:
+  name: "pet"
+  training:
+    num_epochs: 10  # very short period for demostration
+    num_epochs_warmup: 1
+    learning_rate: 1e-5  # small learning rate to stabilize training
+    finetune:
+      method: "full"  # use fine-tuning strategy
+      read_from: pet-mad-latest.ckpt  # path to the pretrained checkpoint to start from
+
+training_set:
+  systems: 
+    read_from: "data/ethanol_train.xyz"  # path to the finetuning dataset
+    length_unit: angstrom
+  targets:
+    energy:
+      key: "energy-corrected"  # name of the target value
+      unit: "eV"
+
+validation_set:
+  systems: 
+    read_from: "data/ethanol_val.xyz"
+    length_unit: angstrom
+  targets:
+    energy:
+      key: "energy-corrected"
+      unit: "eV"
+
+test_set:
+  systems: 
+    read_from: "data/ethanol_test.xyz"
+    length_unit: angstrom
+  targets:
+    energy:
+      key: "energy-corrected"
+      unit: "eV"

examples/pet-finetuning/c_ft_options.yaml

@@ -0,0 +1,57 @@
+seed: 42
+
+architecture:
+  name: pet
+  training:
+    batch_size: 16
+    num_epochs: 5
+    num_epochs_warmup: 1
+    learning_rate: 1e-5
+
+training_set:
+  systems:
+    read_from: data/ethanol_train.xyz
+    length_unit: angstrom
+  targets:
+    energy:
+      key: energy
+      unit: eV
+      forces: on
+    non_conservative_forces:
+      key: forces
+      type:
+        cartesian:
+          rank: 1
+      per_atom: true
+
+validation_set:
+  systems: 
+    read_from: data/ethanol_val.xyz
+    length_unit: angstrom
+  targets:
+    energy:
+      key: energy
+      unit: eV
+      forces: on
+    non_conservative_forces:
+      key: forces
+      type:
+        cartesian:
+          rank: 1
+      per_atom: true
+
+test_set:
+  systems: 
+    read_from: data/ethanol_test.xyz
+    length_unit: angstrom
+  targets:
+    energy:
+      key: energy
+      unit: eV
+      forces: on
+    non_conservative_forces:
+      key: forces
+      type:
+        cartesian:
+          rank: 1
+      per_atom: true

examples/pet-finetuning/environment.yml

@@ -0,0 +1,11 @@
+channels:
+  - metatensor
+  - conda-forge
+dependencies:
+  - python=3.12
+  - pip
+  - pip:
+    - ase>=3.23
+    - metatrain[pet]>=2025.8,<2026
+    - matplotlib
+    - scikit-learn

examples/pet-finetuning/eval_ex1.yaml

@@ -0,0 +1,6 @@
+systems: 
+  read_from: data/ethanol_test.xyz
+targets: 
+  energy:
+    key: energy-corrected
+    unit: eV

examples/pet-finetuning/eval_ex2.yaml

@@ -0,0 +1,12 @@
+systems: 
+  read_from: data/ethanol_test.xyz
+targets: 
+  energy:
+    key: energy-corrected
+    unit: eV
+  non_conservative_forces:
+    key: forces
+    type:
+      cartesian:
+        rank: 1
+    per_atom: true

examples/pet-finetuning/nc_train_options.yaml

@@ -0,0 +1,56 @@
+seed: 42
+
+architecture:
+  name: pet
+  training:
+    batch_size: 16
+    num_epochs: 10
+    num_epochs_warmup: 1
+    learning_rate: 3e-4
+
+training_set:
+  systems:
+    read_from: data/ethanol_train.xyz
+    length_unit: angstrom
+  targets:
+    energy:
+      unit: eV
+      forces: off
+    non_conservative_forces:
+      key: forces
+      type:
+        cartesian:
+          rank: 1
+      per_atom: true
+
+validation_set:
+  systems: 
+    read_from: data/ethanol_val.xyz
+    length_unit: angstrom
+  targets:
+    energy:
+      key: energy
+      unit: eV
+      forces: off
+    non_conservative_forces:
+      key: forces
+      type:
+        cartesian:
+          rank: 1
+      per_atom: true
+
+test_set:
+  systems: 
+    read_from: data/ethanol_test.xyz
+    length_unit: angstrom
+  targets:
+    energy:
+      key: energy
+      unit: eV
+      forces: off
+    non_conservative_forces:
+      key: forces
+      type:
+        cartesian:
+          rank: 1
+      per_atom: true