Add instructions for fine-tuning (#26)

Author: frostedoyster

README.md

@@ -105,6 +105,13 @@ Models
 See [here](docs/models.md) for the complete list of the models we provide. If you are
 new to FlashMD, we recommend starting with the ``pet-omatpes-v2`` models.
 
+Training/fine-tuning your own FlashMD models
+--------------------------------------------
+
+FlashMD models can be trained from the metatrain library. This
+[tutorial](https://docs.metatensor.org/metatrain/latest/generated_examples/1-advanced/04-flashmd.html)
+shows how to train your own FlashMD model, either from scratch or via fine-tuning of one of our universal models.
+
 Disclaimer
 ----------
 

src/flashmd/__init__.py

@@ -1,6 +1,7 @@
 import warnings
 
 from .models import get_pretrained as get_pretrained
+from .models import save_checkpoint as save_checkpoint
 
 
 warnings.filterwarnings("ignore", category=UserWarning, message="custom data")

src/flashmd/models.py

@@ -1,4 +1,5 @@
 import os
+import shutil
 import subprocess
 import time
 
@@ -81,3 +82,21 @@ def get_pretrained(mlip: str = "pet-omatpes-v2", time_step: int = 16) -> Atomist
         flashmd_model = load_atomistic_model(exported_flashmd_path)
 
     return mlip_model, flashmd_model
+
+
+def save_checkpoint(mlip: str = "pet-omatpes-v2", time_step: int = 16):
+    if time_step not in AVAILABLE_TIME_STEPS[mlip]:
+        raise ValueError(
+            f"Pre-trained FlashMD models based on the {mlip} MLIP are only available "
+            f"for time steps of {', '.join(map(str, AVAILABLE_TIME_STEPS[mlip]))} fs."
+        )
+
+    checkpoint_path = hf_hub_download(
+        repo_id="lab-cosmo/flashmd",
+        filename=f"flashmd_{mlip}_{time_step}fs.ckpt",
+        cache_dir=None,
+        revision="main",
+    )
+
+    # Copy it to the current directory
+    shutil.copyfile(checkpoint_path, f"flashmd_{mlip}_{time_step}fs.ckpt")

tests/test_models.py

@@ -37,3 +37,29 @@ def test_get_pretrained_invalid_time_step():
 
     with pytest.raises(ValueError, match="Pre-trained FlashMD models"):
         get_pretrained(mlip="pet-omatpes", time_step=999)
+
+
+def test_save_checkpoint_invalid_time_step():
+    """Test that save_checkpoint raises ValueError for invalid time step."""
+    from flashmd.models import save_checkpoint
+
+    with pytest.raises(ValueError, match="Pre-trained FlashMD models"):
+        save_checkpoint(mlip="pet-omatpes-v2", time_step=999)
+
+
+def test_save_checkpoint(monkeypatch):
+    """Test that save_checkpoint saves the checkpoint file."""
+    from flashmd.models import save_checkpoint
+
+    # Mock hf_hub_download and shutil.copyfile
+    def mock_hf_hub_download(repo_id, filename, cache_dir, revision):
+        return f"/path/to/{filename}"
+
+    def mock_copyfile(src, dst):
+        assert src == "/path/to/flashmd_pet-omatpes-v2_16fs.ckpt"
+        assert dst == "flashmd_pet-omatpes-v2_16fs.ckpt"
+
+    monkeypatch.setattr("flashmd.models.hf_hub_download", mock_hf_hub_download)
+    monkeypatch.setattr("flashmd.models.shutil.copyfile", mock_copyfile)
+
+    save_checkpoint(mlip="pet-omatpes-v2", time_step=16)