Add hydra parsing

README.rst

@@ -0,0 +1,7 @@
+metatensor-models
+-----------------
+
+This is a repository for models using metatensor, in one shape or another. The only
+requirement is for these models to be able to take metatensor objects as inputs and
+outputs. The models do not need to live entirely in this repository: in the most extreme
+case, this repository can simply contain a wrapper to an external model.

docs/src/dev-docs/index.rst

@@ -2,7 +2,8 @@ Developer documentation
 =======================
 
 This is a collection of documentation for developers of the metatensor-models package.
-It includes documentation on how to add a new model, as well as the API of the utils module.
+It includes documentation on how to add a new model, as well as the API of the utils
+module.
 
 
 .. toctree::

docs/src/quickstart.rst

@@ -4,14 +4,14 @@ Quickstart
 You can install metatensor-models with pip:
 
 .. code-block:: bash
-    
+
     git clone https://github.com/lab-cosmo/metatensor-models
     cd metatensor-models
     pip install .
 
 
-In addition, specific models must be installed by specifying the model name.
-For example, to install the SOAP-BPNN model, you can run:
+In addition, specific models must be installed by specifying the model name. For
+example, to install the SOAP-BPNN model, you can run:
 
 .. code-block:: bash
 

pyproject.toml

@@ -11,11 +11,11 @@ authors = [{name = "metatensor-models developers"}]
 dependencies = [
     "ase",
     "torch",
+    "hydra-core",
     "rascaline-torch @ git+https://github.com/luthaf/rascaline#subdirectory=python/rascaline-torch",
     "metatensor-core",
     "metatensor-operations",
     "metatensor-torch",
-    "pyyaml",
 ]
 
 keywords = [] # TODO
@@ -83,3 +83,6 @@ indent = 4
 include_trailing_comma = true
 lines_after_imports = 2
 known_first_party = "metatensor-models"
+
+[tool.mypy]
+ignore_missing_imports = true

src/metatensor/models/__init__.py

@@ -1 +1,11 @@
+from pathlib import Path
+import torch
+
+PACKAGE_ROOT = Path(__file__).parent.resolve()
+
+CONFIG_PATH = PACKAGE_ROOT / "cli" / "conf"
+ARCHITECTURE_CONFIG_PATH = CONFIG_PATH / "architecture"
+
 __version__ = "2023.11.29"
+
+torch.set_default_dtype(torch.float64)

src/metatensor/models/__main__.py

@@ -1,9 +1,9 @@
-"""The main entry point for the metatensor model interface."""
+"""The main entry point for the metatensor models interface."""
 import argparse
 import sys
 
 from . import __version__
-from .scripts import evaluate, export, train
+from .cli import eval_model, export_model, train_model
 
 
 def main():
@@ -18,54 +18,49 @@ def main():
         version=f"metatensor-models {__version__}",
     )
 
-    ap.add_argument(
-        "--debug",
-        action="store_true",
-        help="Run with debug options.",
-    )
-
-    ap.add_argument(
-        "--logfile", dest="logfile", action="store", help="Logfile (optional)"
-    )
-
     subparser = ap.add_subparsers(help="sub-command help")
     evaluate_parser = subparser.add_parser(
-        "evaluate",
-        help=evaluate.__doc__,
-        description="evaluate",
+        "eval",
+        help=eval_model.__doc__,
+        description="eval model",
         formatter_class=argparse.RawDescriptionHelpFormatter,
     )
-    evaluate_parser.set_defaults(callable="evaluate")
+    evaluate_parser.set_defaults(callable="eval_model")
 
     export_parser = subparser.add_parser(
         "export",
-        help=export.__doc__,
-        description="export",
+        help=export_model.__doc__,
+        description="export model",
         formatter_class=argparse.RawDescriptionHelpFormatter,
     )
-    export_parser.set_defaults(callable="export")
+    export_parser.set_defaults(callable="export_model")
     train_parser = subparser.add_parser(
         "train",
-        help=train.__doc__,
-        description="train",
+        help=train_model.__doc__,
+        description="train model",
         formatter_class=argparse.RawDescriptionHelpFormatter,
     )
-    train_parser.set_defaults(callable="train")
+    train_parser.set_defaults(callable="train_model")
 
     if len(sys.argv) < 2:
-        ap.error("A subcommand is required.")
+        ap.error("You must specify a sub-command")
 
     # Be case insensitive for the subcommand
     sys.argv[1] = sys.argv[1].lower()
+    callable = sys.argv[1].lower()
 
-    args = ap.parse_args(sys.argv[1:])
+    # Workaround since we are using hydra for the train parsing
+    if callable == "train":
+        # Remove "metatensor_model" command to please hydra
+        sys.argv.pop(0)
+        train_model()
 
-    if args.callable == "evaluate":
-        evaluate()
-    elif args.callable == "export":
-        export()
-    elif args.callable == "train":
-        train()
+    else:
+        ap.parse_args([callable])
+        if callable == "eval":
+            eval_model()
+        elif callable == "export":
+            export_model()
 
 
 if __name__ == "__main__":

src/metatensor/models/cli/__init__.py

@@ -0,0 +1,3 @@
+from .eval_model import eval_model  # noqa
+from .export_model import export_model  # noqa
+from .train_model import train_model  # noqa

src/metatensor/models/cli/conf/architecture/soap_bpnn.yaml

@@ -0,0 +1,32 @@
+# default hyperparameters for the SOAP-BPNN model
+name: soap_bpnn
+
+model:
+  soap:
+    cutoff: 5.0
+    max_radial: 8
+    max_angular: 6
+    atomic_gaussian_width: 0.3
+    radial_basis:
+      Gto: {}
+    center_atom_weight: 1.0
+    cutoff_function:
+      ShiftedCosine:
+        width: 1.0
+    radial_scaling:
+      Willatt2018:
+        rate: 1.0
+        scale: 2.0
+        exponent: 7.0
+
+  bpnn:
+    num_hidden_layers: 2
+    num_neurons_per_layer: 32
+    activation_function: SiLU
+
+training:
+  batch_size: 8
+  num_epochs: 100
+  learning_rate: 0.001
+  log_interval: 10
+  checkpoint_interval: 25

src/metatensor/models/cli/conf/config.yaml

@@ -0,0 +1,2 @@
+defaults:
+  - architecture: ???

src/metatensor/models/scripts/evaluate.py → src/metatensor/models/cli/eval_model.py

@@ -1,3 +1,3 @@
-def evaluate():
+def eval_model():
     """evaluate a model"""
     print("Run evaluate...")

src/metatensor/models/scripts/export.py → src/metatensor/models/cli/export_model.py

@@ -1,3 +1,3 @@
-def export():
+def export_model():
     """export a model"""
     print("Run exort...")

src/metatensor/models/cli/train_model.py

@@ -0,0 +1,41 @@
+import importlib
+import logging
+
+import hydra
+from omegaconf import DictConfig, OmegaConf
+
+from metatensor.models.utils.data import Dataset
+from metatensor.models.utils.data.readers import read_structures, read_targets
+
+from .. import CONFIG_PATH
+
+
+logger = logging.getLogger(__name__)
+
+
+@hydra.main(config_path=str(CONFIG_PATH), config_name="config", version_base=None)
+def train_model(config: DictConfig) -> None:
+    """train a model."""
+
+    logger.info("Setting up dataset")
+    structures = read_structures(config["dataset"]["structure_path"])
+    targets = read_targets(
+        config["dataset"]["targets_path"],
+        target_value=config["dataset"]["target_value"],
+    )
+    dataset = Dataset(structures, targets)
+
+    logger.info("Setting up model")
+    architetcure_name = config["architecture"]["name"]
+    architecture = importlib.import_module(f"metatensor.models.{architetcure_name}")
+    model = architecture.Model(
+        all_species=dataset.all_species,
+        hypers=OmegaConf.to_container(config["architecture"]["model"]),
+    )
+
+    logger.info("Run training")
+    architecture.train(
+        model=model,
+        train_dataset=dataset,
+        hypers=OmegaConf.to_container(config["architecture"]["training"]),
+    )

src/metatensor/models/soap_bpnn/__init__.py

@@ -1,2 +1,2 @@
-from .model import SoapBPNN  # noqa: F401
+from .model import Model  # noqa: F401
 from .train import train  # noqa: F401

src/metatensor/models/soap_bpnn/model.py

@@ -80,7 +80,7 @@ def forward(self, features: TensorMap) -> TensorMap:
         return TensorMap(keys=features.keys, blocks=new_blocks)
 
 
-class SoapBPNN(torch.nn.Module):
+class Model(torch.nn.Module):
     def __init__(self, all_species: List[int], hypers: Dict) -> None:
         super().__init__()
         self.all_species = all_species

src/metatensor/models/soap_bpnn/tests/__init__.py

@@ -0,0 +1,13 @@
+from pathlib import Path
+
+from metatensor.models import ARCHITECTURE_CONFIG_PATH
+from omegaconf import OmegaConf
+
+
+DEAFAULT_HYPERS = OmegaConf.to_container(
+    OmegaConf.load(ARCHITECTURE_CONFIG_PATH / "soap_bpnn.yaml")
+)
+DATASET_PATH = str(
+    Path(__file__).parent.resolve()
+    / "../../../../../tests/resources/qm9_reduced_100.xyz"
+)

src/metatensor/models/soap_bpnn/tests/test_functionality.py

@@ -1,25 +1,18 @@
-import os
-
 import ase
 import rascaline.torch
 import torch
-import yaml
-
-from metatensor.models.soap_bpnn import SoapBPNN
 
+from metatensor.models.soap_bpnn import Model
 
-path = os.path.dirname(__file__)
-hypers_path = os.path.join(path, "../default.yml")
-dataset_path = os.path.join(path, "data/qm9_reduced_100.xyz")
+from . import DEAFAULT_HYPERS
 
 
 def test_prediction_subset():
     """Tests that the model can predict on a subset
     of the elements it was trained on."""
 
     all_species = [1, 6, 7, 8]
-    hypers = yaml.safe_load(open(hypers_path, "r"))
-    soap_bpnn = SoapBPNN(all_species, hypers).to(torch.float64)
+    soap_bpnn = Model(all_species, DEAFAULT_HYPERS["model"]).to(torch.float64)
 
     structure = ase.Atoms("O2", positions=[[0.0, 0.0, 0.0], [0.0, 0.0, 1.0]])
     soap_bpnn([rascaline.torch.systems_to_torch(structure)])

src/metatensor/models/soap_bpnn/tests/test_invariance.py

@@ -1,27 +1,21 @@
 import copy
-import os
 
 import ase.io
 import rascaline.torch
 import torch
-import yaml
 
-from metatensor.models.soap_bpnn import SoapBPNN
+from metatensor.models.soap_bpnn import Model
 
-
-path = os.path.dirname(__file__)
-hypers_path = os.path.join(path, "../default.yml")
-dataset_path = os.path.join(path, "data/qm9_reduced_100.xyz")
+from . import DATASET_PATH, DEAFAULT_HYPERS
 
 
 def test_rotational_invariance():
     """Tests that the model is rotationally invariant."""
 
     all_species = [1, 6, 7, 8]
-    hypers = yaml.safe_load(open(hypers_path, "r"))
-    soap_bpnn = SoapBPNN(all_species, hypers).to(torch.float64)
+    soap_bpnn = Model(all_species, DEAFAULT_HYPERS["model"]).to(torch.float64)
 
-    structure = ase.io.read(dataset_path)
+    structure = ase.io.read(DATASET_PATH)
     original_structure = copy.deepcopy(structure)
     structure.rotate(48, "y")