Add PhACE architecture

.github/workflows/architecture-tests.yml

@@ -15,6 +15,7 @@ jobs:
           - architecture-name: soap-bpnn
           - architecture-name: pet
           - architecture-name: nanopet
+          - architecture-name: phace
 
     runs-on: ubuntu-22.04
 

pyproject.toml

@@ -71,6 +71,10 @@ gap = [
     "metatensor-learn",
     "scipy",
 ]
+phace = [
+    "physical_basis",
+    "wigners",
+]
 
 [tool.setuptools.packages.find]
 where = ["src"]

src/metatrain/cli/train.py

@@ -395,18 +395,20 @@ def train_model(
     # TRAIN MODEL #############
     ###########################
 
-    logger.info("Calling trainer")
-    try:
-        trainer.train(
-            model=model,
-            dtype=dtype,
-            devices=devices,
-            train_datasets=train_datasets,
-            val_datasets=val_datasets,
-            checkpoint_dir=str(checkpoint_dir),
-        )
-    except Exception as e:
-        raise ArchitectureError(e)
+    # logger.info("Calling trainer")
+    # from torch.profiler import profile, ProfilerActivity
+    # with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof:
+    # if True:
+    trainer.train(
+        model=model,
+        dtype=dtype,
+        devices=devices,
+        train_datasets=train_datasets,
+        val_datasets=val_datasets,
+        checkpoint_dir=str(checkpoint_dir),
+    )
+    # print(prof.key_averages().table(sort_by="self_cpu_time_total", row_limit=20))
+    # exit()
 
     if not is_main_process():
         return  # only save and evaluate on the main process

src/metatrain/experimental/nanopet/trainer.py

@@ -67,7 +67,7 @@ def train(
             if len(devices) > 1:
                 raise ValueError(
                     "Requested distributed training with the `multi-gpu` device. "
-                    " If you want to run distributed training with SOAP-BPNN, please "
+                    " If you want to run distributed training with nanoPET, please "
                     "set `device` to cuda."
                 )
             # the calculation of the device number works both when GPUs on different
@@ -98,7 +98,7 @@ def train(
 
         # Move the model to the device and dtype:
         model.to(device=device, dtype=dtype)
-        # The additive models of the SOAP-BPNN are always in float64 (to avoid
+        # The additive models of nanoPET are always in float64 (to avoid
         # numerical errors in the composition weights, which can be very large).
         for additive_model in model.additive_models:
             additive_model.to(dtype=torch.float64)

src/metatrain/experimental/phace/__init__.py

@@ -0,0 +1,13 @@
+from .model import PhACE
+from .trainer import Trainer
+
+__model__ = PhACE
+__trainer__ = Trainer
+
+__authors__ = [
+    ("Filippo Bigi <[email protected]>", "@frostedoyster"),
+]
+
+__maintainers__ = [
+    ("Filippo Bigi <[email protected]>", "@frostedoyster"),
+]

src/metatrain/experimental/phace/default-hypers.yaml

@@ -0,0 +1,41 @@
+architecture:
+  name: experimental.phace 
+
+  model:
+    nu_max: 3
+    n_message_passing_layers: 2
+    cutoff: 5.0
+    n_element_channels: 8
+    radial_basis:
+      mlp: true
+      type: physical
+      E_max: 70.0
+      scale: 0.7
+    nu_scaling: 1.0
+    mp_scaling: 1.0
+    overall_scaling: 1.0
+    disable_nu_0: True
+    use_sphericart: False
+    use_mops: False
+    heads: {}
+    zbl: False
+
+  training:
+    distributed: False
+    distributed_port: 39591
+    batch_size: 8
+    num_epochs: 100
+    learning_rate: 0.001
+    early_stopping_patience: 100
+    scheduler_patience: 10
+    scheduler_factor: 0.8
+    log_interval: 10
+    checkpoint_interval: 25
+    scale_targets: true
+    fixed_composition_weights: {}
+    per_structure_targets: []
+    log_mae: False
+    loss:
+      type: mse
+      weights: {}
+      reduction: sum