Update filter_nonzero_weight calls in visualise_train.py to correct API

Author: bernstei

mace/cli/visualise_train.py

@@ -569,24 +569,22 @@ def update(self, batch, output):  # pylint: disable=arguments-differ
             self.pred_energies_per_atom.append(output["energy"] / atoms_per_config)
 
             self.n_energy += filter_nonzero_weight(
-                batch, self.ref_energies, batch.weight, batch.energy_weight, "config"
+                batch, self.ref_energies, batch.weight, batch.energy_weight,
             )
             filter_nonzero_weight(
-                batch, self.pred_energies, batch.weight, batch.energy_weight, "config"
+                batch, self.pred_energies, batch.weight, batch.energy_weight,
             )
             filter_nonzero_weight(
                 batch,
                 self.ref_energies_per_atom,
                 batch.weight,
                 batch.energy_weight,
-                "config",
             )
             filter_nonzero_weight(
                 batch,
                 self.pred_energies_per_atom,
                 batch.weight,
                 batch.energy_weight,
-                "config",
             )
 
         if output.get("interaction_energy") is not None and batch.energy is not None:
@@ -608,28 +606,24 @@ def update(self, batch, output):  # pylint: disable=arguments-differ
                 self.ref_interaction_energies,
                 batch.weight,
                 batch.energy_weight,
-                "config",
             )
             filter_nonzero_weight(
                 batch,
                 self.pred_interaction_energies,
                 batch.weight,
                 batch.energy_weight,
-                "config",
             )
             filter_nonzero_weight(
                 batch,
                 self.ref_interaction_energies_per_atom,
                 batch.weight,
                 batch.energy_weight,
-                "config",
             )
             filter_nonzero_weight(
                 batch,
                 self.pred_interaction_energies_per_atom,
                 batch.weight,
                 batch.energy_weight,
-                "config",
             )
 
         # Forces
@@ -638,10 +632,10 @@ def update(self, batch, output):  # pylint: disable=arguments-differ
             self.pred_forces.append(output["forces"])
 
             self.n_forces += filter_nonzero_weight(
-                batch, self.ref_forces, batch.weight, batch.forces_weight, "atom"
+                batch, self.ref_forces, batch.weight, batch.forces_weight, spread_atoms=True,
             )
             filter_nonzero_weight(
-                batch, self.pred_forces, batch.weight, batch.forces_weight, "atom"
+                batch, self.pred_forces, batch.weight, batch.forces_weight, spread_atoms=True,
             )
 
         # Stress
@@ -650,10 +644,10 @@ def update(self, batch, output):  # pylint: disable=arguments-differ
             self.pred_stress.append(output["stress"])
 
             self.n_stress += filter_nonzero_weight(
-                batch, self.ref_stress, batch.weight, batch.stress_weight, "config"
+                batch, self.ref_stress, batch.weight, batch.stress_weight,
             )
             filter_nonzero_weight(
-                batch, self.pred_stress, batch.weight, batch.stress_weight, "config"
+                batch, self.pred_stress, batch.weight, batch.stress_weight,
             )
 
         # Virials
@@ -666,24 +660,22 @@ def update(self, batch, output):  # pylint: disable=arguments-differ
             self.pred_virials_per_atom.append(output["virials"] / atoms_per_config_3d)
 
             self.n_virials += filter_nonzero_weight(
-                batch, self.ref_virials, batch.weight, batch.virials_weight, "config"
+                batch, self.ref_virials, batch.weight, batch.virials_weight,
             )
             filter_nonzero_weight(
-                batch, self.pred_virials, batch.weight, batch.virials_weight, "config"
+                batch, self.pred_virials, batch.weight, batch.virials_weight,
             )
             filter_nonzero_weight(
                 batch,
                 self.ref_virials_per_atom,
                 batch.weight,
                 batch.virials_weight,
-                "config",
             )
             filter_nonzero_weight(
                 batch,
                 self.pred_virials_per_atom,
                 batch.weight,
                 batch.virials_weight,
-                "config",
             )
 
         # Dipole
@@ -705,14 +697,14 @@ def update(self, batch, output):  # pylint: disable=arguments-differ
                 self.ref_dipole_per_atom,
                 batch.weight,
                 batch.dipole_weight,
-                "config",
+                spread_quantity_vector=False,
             )
             filter_nonzero_weight(
                 batch,
                 self.pred_dipole_per_atom,
                 batch.weight,
                 batch.dipole_weight,
-                "config",
+                spread_quantity_vector=False,
             )
 
     def _process_data(self, ref_list, pred_list):

mace/tools/train.py

@@ -611,7 +611,7 @@ def update(self, batch, output):  # pylint: disable=arguments-differ
             )
             self.E_computed += filter_nonzero_weight(
                 batch, self.delta_es, batch.weight, batch.energy_weight
-            )  # DEBUG , label="delta_es")
+            )
         if output.get("forces") is not None and batch.forces is not None:
             self.fs.append(batch.forces)
             self.delta_fs.append(batch.forces - output["forces"])
@@ -621,12 +621,12 @@ def update(self, batch, output):  # pylint: disable=arguments-differ
                 batch.weight,
                 batch.forces_weight,
                 spread_atoms=True,
-            )  # DEBUG , label="delta_fs")
+            )
         if output.get("stress") is not None and batch.stress is not None:
             self.delta_stress.append(batch.stress - output["stress"])
             self.stress_computed += filter_nonzero_weight(
                 batch, self.delta_stress, batch.weight, batch.stress_weight
-            )  # DEBUG , label="delta_stress")
+            )
         if output.get("virials") is not None and batch.virials is not None:
             self.delta_virials.append(batch.virials - output["virials"])
             self.delta_virials_per_atom.append(
@@ -635,7 +635,7 @@ def update(self, batch, output):  # pylint: disable=arguments-differ
             )
             self.virials_computed += filter_nonzero_weight(
                 batch, self.delta_virials, batch.weight, batch.virials_weight
-            )  # DEBUG , label="delta_virials")
+            )
         if output.get("dipole") is not None and batch.dipole is not None:
             self.mus.append(batch.dipole)
             self.delta_mus.append(batch.dipole - output["dipole"])
@@ -649,7 +649,7 @@ def update(self, batch, output):  # pylint: disable=arguments-differ
                 batch.weight,
                 batch.dipole_weight,
                 spread_quantity_vector=False,
-            )  # DEBUG , label="delta_mus")
+            )
         if (
             output.get("polarizability") is not None
             and batch.polarizability is not None