Fix Nan gradients in Force model with padded_disjoint representaiton

Author: PatReis

kgcnn/__init__.py

@@ -8,3 +8,7 @@
 
 # Behaviour for backend functions.
 __safe_scatter_max_min_to_zero__ = True
+
+# Geometry
+__geom_euclidean_norm_add_eps__ = False
+__geom_euclidean_norm_no_nan__ = True  # Only used for inverse norm.

kgcnn/layers/casting.py

@@ -7,7 +7,6 @@
 
 
 def _pad_left(t):
-    # return ops.concatenate([ops.zeros_like(t[:1]), t], axis=0)
     return ops.pad(t, [[1, 0]] + [[0, 0] for _ in range(len(ops.shape(t)) - 1)])
 
 

kgcnn/layers/geom.py

@@ -10,6 +10,8 @@
 from kgcnn.layers.polynom import SphericalBesselJnExplicit, SphericalHarmonicsYl
 from kgcnn.ops.axis import get_positive_axis
 from kgcnn.ops.core import cross as kgcnn_cross
+from kgcnn import __geom_euclidean_norm_add_eps__ as global_geom_euclidean_norm_add_eps
+from kgcnn import __geom_euclidean_norm_no_nan__ as global_geom_euclidean_norm_no_nan
 
 
 class NodePosition(Layer):
@@ -142,8 +144,11 @@ class EuclideanNorm(Layer):
     with :obj:`invert_norm` layer arguments.
     """
 
-    def __init__(self, axis: int = -1, keepdims: bool = False, invert_norm: bool = False, add_eps: bool = False,
-                 no_nan: bool = True, square_norm: bool = False, **kwargs):
+    def __init__(self, axis: int = -1, keepdims: bool = False,
+                 invert_norm: bool = False,
+                 add_eps: bool = global_geom_euclidean_norm_add_eps,
+                 no_nan: bool = global_geom_euclidean_norm_no_nan,
+                 square_norm: bool = False, **kwargs):
         """Initialize layer.
 
         Args:
@@ -177,7 +182,7 @@ def compute_output_shape(self, input_shape):
 
     @staticmethod
     def _compute_euclidean_norm(inputs, axis: int = -1, keepdims: bool = False, invert_norm: bool = False,
-                                add_eps: bool = False, no_nan: bool = True, square_norm: bool = False):
+                                add_eps: bool = False, no_nan: bool = False, square_norm: bool = False):
         """Function to compute euclidean norm for inputs.
 
         Args:
@@ -306,7 +311,10 @@ class NodeDistanceEuclidean(Layer):
     the output of :obj:`NodePosition`.
     """
 
-    def __init__(self, add_eps: bool = False, no_nan: bool = True, **kwargs):
+    def __init__(self,
+                 add_eps: bool = global_geom_euclidean_norm_add_eps,
+                 no_nan: bool = global_geom_euclidean_norm_no_nan,
+                 **kwargs):
         r"""Initialize layer instance of :obj:`NodeDistanceEuclidean`. """
         super(NodeDistanceEuclidean, self).__init__(**kwargs)
         self.layer_subtract = Subtract()
@@ -354,7 +362,9 @@ class EdgeDirectionNormalized(Layer):
     As the first index defines the incoming edge.
     """
 
-    def __init__(self, add_eps: bool = False, no_nan: bool = True, **kwargs):
+    def __init__(self, add_eps: bool = global_geom_euclidean_norm_add_eps,
+                 no_nan: bool = global_geom_euclidean_norm_no_nan,
+                 **kwargs):
         """Initialize layer."""
         super(EdgeDirectionNormalized, self).__init__(**kwargs)
         self.layer_subtract = Subtract()

training/train_force.py

@@ -4,6 +4,7 @@
 import argparse
 import keras as ks
 from datetime import timedelta
+import kgcnn
 import kgcnn.training.schedule
 import kgcnn.training.scheduler
 from kgcnn.data.utils import save_pickle_file
@@ -18,6 +19,9 @@
 from kgcnn.metrics.metrics import ScaledMeanAbsoluteError, ScaledForceMeanAbsoluteError
 from kgcnn.data.transform.scaler.force import EnergyForceExtensiveLabelScaler
 
+# For force gradients
+kgcnn.__geom_euclidean_norm_add_eps__ = True
+
 # Input arguments from command line.
 parser = argparse.ArgumentParser(description='Train a GNN on an Energy-Force Dataset.')
 parser.add_argument("--hyper", required=False, help="Filepath to hyper-parameter config file (.py or .json).",