print NaN loss when labeled data is not found

Currently, when loss is defined and the labeled data is not found, `lcurve.out` shows wrong RMSE (assuming data is all zero).
In this case, printing NaN is better.

Signed-off-by: Jinzhe Zeng <[email protected]>

deepmd/loss/dos.py

@@ -143,16 +143,20 @@ def build(self, learning_rate, natoms, model_dict, label_dict, suffix):
         more_loss = {}
         if self.has_dos:
             l2_loss += atom_norm_ener * (pref_dos * l2_dos_loss)
-            more_loss["l2_dos_loss"] = l2_dos_loss
+            more_loss["l2_dos_loss"] = self.display_if_exist(l2_dos_loss, find_dos)
         if self.has_cdf:
             l2_loss += atom_norm_ener * (pref_cdf * l2_cdf_loss)
-            more_loss["l2_cdf_loss"] = l2_cdf_loss
+            more_loss["l2_cdf_loss"] = self.display_if_exist(l2_cdf_loss, find_dos)
         if self.has_ados:
             l2_loss += global_cvt_2_ener_float(pref_ados * l2_atom_dos_loss)
-            more_loss["l2_atom_dos_loss"] = l2_atom_dos_loss
+            more_loss["l2_atom_dos_loss"] = self.display_if_exist(
+                l2_atom_dos_loss, find_atom_dos
+            )
         if self.has_acdf:
             l2_loss += global_cvt_2_ener_float(pref_acdf * l2_atom_cdf_loss)
-            more_loss["l2_atom_cdf_loss"] = l2_atom_cdf_loss
+            more_loss["l2_atom_cdf_loss"] = self.display_if_exist(
+                l2_atom_cdf_loss, find_atom_dos
+            )
 
         # only used when tensorboard was set as true
         self.l2_loss_summary = tf.summary.scalar("l2_loss_" + suffix, tf.sqrt(l2_loss))

deepmd/loss/ener.py

@@ -291,22 +291,32 @@ def build(self, learning_rate, natoms, model_dict, label_dict, suffix):
         more_loss = {}
         if self.has_e:
             l2_loss += atom_norm_ener * (pref_e * l2_ener_loss)
-            more_loss["l2_ener_loss"] = l2_ener_loss
+            more_loss["l2_ener_loss"] = self.display_if_exist(l2_ener_loss, find_energy)
         if self.has_f:
             l2_loss += global_cvt_2_ener_float(pref_f * l2_force_loss)
-            more_loss["l2_force_loss"] = l2_force_loss
+            more_loss["l2_force_loss"] = self.display_if_exist(
+                l2_force_loss, find_force
+            )
         if self.has_v:
             l2_loss += global_cvt_2_ener_float(atom_norm * (pref_v * l2_virial_loss))
-            more_loss["l2_virial_loss"] = l2_virial_loss
+            more_loss["l2_virial_loss"] = self.display_if_exist(
+                l2_virial_loss, find_virial
+            )
         if self.has_ae:
             l2_loss += global_cvt_2_ener_float(pref_ae * l2_atom_ener_loss)
-            more_loss["l2_atom_ener_loss"] = l2_atom_ener_loss
+            more_loss["l2_atom_ener_loss"] = self.display_if_exist(
+                l2_atom_ener_loss, find_atom_ener
+            )
         if self.has_pf:
             l2_loss += global_cvt_2_ener_float(pref_pf * l2_pref_force_loss)
-            more_loss["l2_pref_force_loss"] = l2_pref_force_loss
+            more_loss["l2_pref_force_loss"] = self.display_if_exist(
+                l2_pref_force_loss, find_atom_pref
+            )
         if self.has_gf:
             l2_loss += global_cvt_2_ener_float(pref_gf * l2_gen_force_loss)
-            more_loss["l2_gen_force_loss"] = l2_gen_force_loss
+            more_loss["l2_gen_force_loss"] = self.display_if_exist(
+                l2_gen_force_loss, find_drdq
+            )
 
         # only used when tensorboard was set as true
         self.l2_loss_summary = tf.summary.scalar("l2_loss_" + suffix, tf.sqrt(l2_loss))
@@ -553,19 +563,25 @@ def build(self, learning_rate, natoms, model_dict, label_dict, suffix):
         more_loss = {}
         if self.has_e:
             l2_loss += atom_norm_ener * (pref_e * l2_ener_loss)
-        more_loss["l2_ener_loss"] = l2_ener_loss
+        more_loss["l2_ener_loss"] = self.display_if_exist(l2_ener_loss, find_energy)
         if self.has_fr:
             l2_loss += global_cvt_2_ener_float(pref_fr * l2_force_r_loss)
-        more_loss["l2_force_r_loss"] = l2_force_r_loss
+        more_loss["l2_force_r_loss"] = self.display_if_exist(
+            l2_force_r_loss, find_force
+        )
         if self.has_fm:
             l2_loss += global_cvt_2_ener_float(pref_fm * l2_force_m_loss)
-        more_loss["l2_force_m_loss"] = l2_force_m_loss
+        more_loss["l2_force_m_loss"] = self.display_if_exist(
+            l2_force_m_loss, find_force
+        )
         if self.has_v:
             l2_loss += global_cvt_2_ener_float(atom_norm * (pref_v * l2_virial_loss))
-        more_loss["l2_virial_loss"] = l2_virial_loss
+        more_loss["l2_virial_loss"] = self.display_if_exist(l2_virial_loss, find_virial)
         if self.has_ae:
             l2_loss += global_cvt_2_ener_float(pref_ae * l2_atom_ener_loss)
-        more_loss["l2_atom_ener_loss"] = l2_atom_ener_loss
+        more_loss["l2_atom_ener_loss"] = self.display_if_exist(
+            l2_atom_ener_loss, find_atom_ener
+        )
 
         # only used when tensorboard was set as true
         self.l2_loss_summary = tf.summary.scalar("l2_loss", tf.sqrt(l2_loss))
@@ -785,8 +801,10 @@ def build(self, learning_rate, natoms, model_dict, label_dict, suffix):
         more_loss = {}
         l2_loss += atom_norm_ener * (pref_e * l2_ener_loss)
         l2_loss += global_cvt_2_ener_float(pref_ed * l2_ener_dipole_loss)
-        more_loss["l2_ener_loss"] = l2_ener_loss
-        more_loss["l2_ener_dipole_loss"] = l2_ener_dipole_loss
+        more_loss["l2_ener_loss"] = self.display_if_exist(l2_ener_loss, find_energy)
+        more_loss["l2_ener_dipole_loss"] = self.display_if_exist(
+            l2_ener_dipole_loss, find_ener_dipole
+        )
 
         self.l2_loss_summary = tf.summary.scalar("l2_loss_" + suffix, tf.sqrt(l2_loss))
         self.l2_loss_ener_summary = tf.summary.scalar(

deepmd/loss/loss.py

@@ -8,6 +8,8 @@
     Tuple,
 )
 
+import numpy as np
+
 from deepmd.env import (
     tf,
 )
@@ -72,3 +74,20 @@ def eval(
             A dictionary that maps keys to values. It
             should contain key `natoms`
         """
+
+    @staticmethod
+    def display_if_exist(loss: tf.Tensor, find_property: float) -> tf.Tensor:
+        """Display NaN if labeled property is not found.
+
+        Parameters
+        ----------
+        loss : tf.Tensor
+            the loss tensor
+        find_property : float
+            whether the property is found
+        """
+        return tf.cond(
+            tf.cast(find_property, tf.bool),
+            lambda: loss,
+            lambda: tf.cast(np.nan, dtype=loss.dtype),
+        )

deepmd/loss/tensor.py

@@ -87,7 +87,7 @@ def build(self, learning_rate, natoms, model_dict, label_dict, suffix):
             local_loss = global_cvt_2_tf_float(find_atomic) * tf.reduce_mean(
                 tf.square(self.scale * (polar - atomic_polar_hat)), name="l2_" + suffix
             )
-            more_loss["local_loss"] = local_loss
+            more_loss["local_loss"] = self.display_if_exist(local_loss, find_atomic)
             l2_loss += self.local_weight * local_loss
             self.l2_loss_local_summary = tf.summary.scalar(
                 "l2_local_loss_" + suffix, tf.sqrt(more_loss["local_loss"])
@@ -118,7 +118,7 @@ def build(self, learning_rate, natoms, model_dict, label_dict, suffix):
                 tf.square(self.scale * (global_polar - polar_hat)), name="l2_" + suffix
             )
 
-            more_loss["global_loss"] = global_loss
+            more_loss["global_loss"] = self.display_if_exist(global_loss, find_global)
             self.l2_loss_global_summary = tf.summary.scalar(
                 "l2_global_loss_" + suffix,
                 tf.sqrt(more_loss["global_loss"]) / global_cvt_2_tf_float(atoms),