Suggestions from review

Author: frostedoyster

src/metatrain/cli/train.py

@@ -422,14 +422,17 @@ def train_model(
     ###########################
 
     logger.info("Calling trainer")
-    trainer.train(
-        model=model,
-        dtype=dtype,
-        devices=devices,
-        train_datasets=train_datasets,
-        val_datasets=val_datasets,
-        checkpoint_dir=str(checkpoint_dir),
-    )
+    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)
 
     if not is_main_process():
         return  # only save and evaluate on the main process

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

@@ -2,14 +2,14 @@ architecture:
   name: experimental.phace 
 
   model:
-    nu_max: 3
+    max_correlation_order_per_layer: 3
     num_message_passing_layers: 2
     cutoff: 5.0
     cutoff_width: 1.0
     num_element_channels: 64
     radial_basis:
       mlp: true
-      E_max: 50.0
+      max_eigenvalue: 50.0
       scale: 0.7
       optimizable_lengthscales: false
     nu_scaling: 0.1

src/metatrain/experimental/phace/model.py

@@ -23,7 +23,7 @@
 from .modules.cg import get_cg_coefficients
 from .modules.cg_iterator import CGIterator
 from .modules.initial_features import get_initial_features
-from .modules.layers import EquivariantLastLayer, InvariantMLP, NothingLayer
+from .modules.layers import EquivariantLastLayer, Identity, InvariantMLP
 from .modules.message_passing import EquivariantMessagePasser, InvariantMessagePasser
 from .modules.precomputations import Precomputer
 from .utils import systems_to_batch
@@ -74,7 +74,7 @@ def __init__(self, model_hypers: Dict, dataset_info: DatasetInfo) -> None:
         )
         self.register_buffer("species_to_species_index", species_to_species_index)
 
-        self.nu_max = model_hypers["nu_max"]
+        self.nu_max = model_hypers["max_correlation_order_per_layer"]
         self.num_message_passing_layers = model_hypers["num_message_passing_layers"]
         if self.num_message_passing_layers < 1:
             raise ValueError("Number of message-passing layers must be at least 1")
@@ -521,7 +521,7 @@ def _add_output(self, target_name: str, target_info: TargetInfo) -> None:
                 self.k_max_l[0], self.head_num_layers
             )
         else:
-            self.heads[target_name] = NothingLayer()
+            self.heads[target_name] = Identity()
 
         if target_info.is_scalar:
             self.last_layers[target_name] = EquivariantLastLayer(

src/metatrain/experimental/phace/modules/center_embedding.py

@@ -1,7 +1,7 @@
 from typing import List
 
 import torch
-from metatensor.torch import Labels, TensorBlock, TensorMap
+from metatensor.torch import TensorBlock, TensorMap
 
 
 def embed_centers(equivariants: TensorMap, center_embeddings: torch.Tensor):
@@ -29,9 +29,6 @@ def embed_centers(equivariants: TensorMap, center_embeddings: torch.Tensor):
         )
 
     return TensorMap(
-        keys=Labels(
-            names=["nu", "o3_lambda", "o3_sigma"],
-            values=torch.stack(keys).to(equivariants.keys.values.device),
-        ),
+        keys=equivariants.keys,
         blocks=blocks,
     )

src/metatrain/experimental/phace/modules/layers.py

@@ -179,7 +179,7 @@ def forward(self, features: TensorMap) -> TensorMap:
         return TensorMap(keys=new_keys, blocks=new_blocks)
 
 
-class NothingLayer(torch.nn.Module):
+class Identity(torch.nn.Module):
     # useful when the head for an output is a simple linear layer
 
     def forward(self, features: TensorMap) -> TensorMap:

src/metatrain/experimental/phace/modules/message_passing.py

@@ -15,10 +15,8 @@ class DummyAdder(torch.nn.Module):
     def __init__(self):
         super().__init__()
 
-    def forward(
-        self, tmap_1: metatensor.torch.TensorMap, tmap_2: metatensor.torch.TensorMap
-    ) -> metatensor.torch.TensorMap:
-        return metatensor.torch.TensorMap(
+    def forward(self, tmap_1: TensorMap, tmap_2: TensorMap) -> TensorMap:
+        return TensorMap(
             keys=Labels(names=["dummy"], values=torch.empty(1, 1)), blocks=[]
         )
 

src/metatrain/experimental/phace/modules/physical_basis.py

@@ -29,7 +29,7 @@ def get_physical_basis_spliner(E_max, r_cut, normalize):
     for l in range(l_max + 1):  # noqa: E741
         n_max_l.append(np.where(E_nl[:, l] <= E_max)[0][-1] + 1)
     if n_max_l[0] > n_max:
-        raise ValueError("n_max too large, try decreasing E_max")
+        raise ValueError("n_max too large, try decreasing max_eigenvalue")
 
     def function_for_splining(n, l, x):  # noqa: E741
         ret = physical_basis.compute(n, l, x)

src/metatrain/experimental/phace/modules/radial_basis.py

@@ -16,7 +16,7 @@ def __init__(self, hypers, all_species) -> None:
         for species in all_species:
             lengthscales[species] = np.log(hypers["scale"] * covalent_radii[species])
         self.n_max_l, self.spliner = get_physical_basis_spliner(
-            hypers["E_max"], hypers["cutoff"], normalize=True
+            hypers["max_eigenvalue"], hypers["cutoff"], normalize=True
         )
         if hypers["optimizable_lengthscales"]:
             self.lengthscales = torch.nn.Parameter(lengthscales)

src/metatrain/experimental/phace/schema-hypers.json

@@ -9,7 +9,7 @@
     "model": {
       "type": "object",
       "properties": {
-        "nu_max": {
+        "max_correlation_order_per_layer": {
           "type": "integer"
         },
         "num_message_passing_layers": {
@@ -30,7 +30,7 @@
             "mlp": {
               "type": "boolean"
             },
-            "E_max": {
+            "max_eigenvalue": {
               "type": "number"
             },
             "scale": {

tox.ini

@@ -137,11 +137,10 @@ commands =
     pytest {posargs}
 
 [testenv:phace-tests]
-description = Run NanoPET tests with pytest
+description = Run PhACE tests with pytest
 passenv = *
 deps =
     pytest
-    spherical  # for nanoPET spherical target
 extras = phace
 changedir = src/metatrain/experimental/phace/tests/
 commands =