Updated to latest tensorial 0.6.0

Author: muhrin

configs/listeners/default.yaml

@@ -2,7 +2,8 @@ defaults:
   - model_checkpoint
   - early_stopping
   #  - model_summary
-  - rich_progress_bar
+  - metrics_printer
+#  - rich_progress_bar
   - _self_
 
 model_checkpoint:

configs/listeners/metrics_printer.yaml

@@ -0,0 +1,2 @@
+metrics_printer:
+  _target_: tensorial.reaxkit.MetricsPrinter

configs/model/mace.yaml

@@ -15,12 +15,14 @@ loss_fn:
   weights: [ 1000., 1. ]
   loss_fns:
     - _target_: tensorial.gcnn.Loss
-      field: nodes.born_charges_predicted
-      target_field: nodes.born_charges
+      loss_fn: squared_error
+      targets: nodes.born_charges
+      predictions: nodes.born_charges_predicted
 
     - _target_: tensorial.gcnn.Loss
-      field: nodes.raman_tensors_predicted
-      target_field: nodes.raman_tensors
+      loss_fn: squared_error
+      predictions: nodes.raman_tensors_predicted
+      targets: nodes.raman_tensors
 
 
 model:

configs/model/nequip.yaml

@@ -15,12 +15,14 @@ loss_fn:
   weights: [ 1000., 1. ]
   loss_fns:
     - _target_: tensorial.gcnn.Loss
-      field: nodes.born_charges_predicted
-      target_field: nodes.born_charges
+      loss_fn: squared_error
+      predictions: nodes.born_charges_predicted
+      targets: nodes.born_charges
 
     - _target_: tensorial.gcnn.Loss
-      field: nodes.raman_tensors_predicted
-      target_field: nodes.raman_tensors
+      loss_fn: squared_error
+      predictions: nodes.raman_tensors_predicted
+      targets: nodes.raman_tensors
 
 model:
   _target_: tensorial.nn.Sequential

configs/model/nequip_nmr.yaml

@@ -15,12 +15,14 @@ loss_fn:
   weights: [ 1000., 1. ]
   loss_fns:
     - _target_: tensorial.gcnn.Loss
-      field: nodes.born_charges_predicted
-      target_field: nodes.born_charges
+      loss_fn: squared_error
+      targets: nodes.born_charges
+      predictions: nodes.born_charges_predicted
 
     - _target_: tensorial.gcnn.Loss
-      field: nodes.raman_tensors_predicted
-      target_field: nodes.raman_tensors
+      loss_fn: squared_error
+      targets: nodes.raman_tensors
+      predictions: nodes.raman_tensors_predicted
 
 model:
   _target_: tensorial.nn.Sequential

pyproject.toml

@@ -33,7 +33,7 @@ dependencies = [
     'e3nn-jax',
     "equinox",
     "reax>=0.2,<0.7",
-    "tensorial>=0.5.1",
+    "tensorial~=0.6",
     "pymatgen",
 ]
 

src/e3response/losses.py

@@ -2,6 +2,7 @@
 
 import jax
 import jraph
+import optax
 from tensorial import gcnn
 from tensorial.gcnn import atomic
 from tensorial.gcnn.keys import predicted
@@ -23,37 +24,61 @@ def response_loss(
     if energy:
         weights.append(1.0 if isinstance(energy, bool) else energy)
         loss_terms.append(
-            gcnn.Loss(f"globals.{predicted(atomic.TOTAL_ENERGY)}", f"globals.{atomic.TOTAL_ENERGY}")
+            gcnn.Loss(
+                optax.squared_error,
+                f"globals.{atomic.TOTAL_ENERGY}",
+                f"globals.{predicted(atomic.TOTAL_ENERGY)}",
+            )
         )
 
     if forces:
         weights.append(1.0 if isinstance(forces, bool) else forces)
-        loss_terms.append(gcnn.Loss(f"nodes.{predicted(atomic.FORCES)}", f"nodes.{atomic.FORCES}"))
+        loss_terms.append(
+            gcnn.Loss(
+                optax.squared_error,
+                f"nodes.{atomic.FORCES}",
+                f"nodes.{predicted(atomic.FORCES)}",
+            )
+        )
 
     if born_charges:
         weights.append(1.0 if isinstance(born_charges, bool) else born_charges)
         loss_terms.append(
-            gcnn.Loss(f"nodes.{predicted(keys.BORN_CHARGES)}", f"nodes.{keys.BORN_CHARGES}")
+            gcnn.Loss(
+                optax.squared_error,
+                f"nodes.{keys.BORN_CHARGES}",
+                f"nodes.{predicted(keys.BORN_CHARGES)}",
+            )
         )
 
     if polarization_tensors:
         weights.append(1.0 if isinstance(polarization_tensors, bool) else polarization_tensors)
         loss_terms.append(
-            gcnn.Loss(f"globals.{predicted(keys.POLARIZATION)}", f"globals.{keys.POLARIZATION}")
+            gcnn.Loss(
+                optax.squared_error,
+                f"globals.{keys.POLARIZATION}",
+                f"globals.{predicted(keys.POLARIZATION)}",
+            )
         )
 
     if dielectric_tensor:
         weights.append(1.0 if isinstance(dielectric_tensor, bool) else dielectric_tensor)
         loss_terms.append(
             gcnn.Loss(
-                f"globals.{predicted(keys.DIELECTRIC_TENSOR)}", f"globals.{keys.DIELECTRIC_TENSOR}"
+                optax.squared_error,
+                f"globals.{keys.DIELECTRIC_TENSOR}",
+                f"globals.{predicted(keys.DIELECTRIC_TENSOR)}",
             )
         )
 
     if raman_tensors:
         weights.append(1.0 if isinstance(raman_tensors, bool) else raman_tensors)
         loss_terms.append(
-            gcnn.Loss(f"nodes.{predicted(keys.RAMAN_TENSORS)}", f"nodes.{keys.RAMAN_TENSORS}")
+            gcnn.Loss(
+                optax.squared_error,
+                f"nodes.{keys.RAMAN_TENSORS}",
+                f"nodes.{predicted(keys.RAMAN_TENSORS)}",
+            )
         )
 
     if not loss_terms: