Merge pull request #2416 from NNPDF/fix_perhaps_nan_issue2322

Fix NaN accumulation in parallel runs

Author: scarlehoff

n3fit/src/n3fit/backends/keras_backend/MetaModel.py

@@ -50,7 +50,7 @@
 def _default_loss(y_true, y_pred):  # pylint: disable=unused-argument
     """Default loss to be used when the model is compiled with loss = Null
     (for instance if the prediction of the model is already the loss"""
-    return ops.sum(y_pred)
+    return ops.nansum(y_pred)
 
 
 class MetaModel(Model):
@@ -219,7 +219,7 @@ def losses_fun():
                 # If we only have one dataset the output changes
                 if len(out_names) == 2:
                     predictions = [predictions]
-                total_loss = ops.sum(predictions, axis=0)
+                total_loss = ops.nansum(predictions, axis=0)
                 ret = [total_loss] + predictions
                 return dict(zip(out_names, ret))
 

n3fit/src/n3fit/backends/keras_backend/operations.py

@@ -46,6 +46,7 @@
     expand_dims,
     leaky_relu,
     reshape,
+    nan_to_num,
     repeat,
     split,
     sum,
@@ -290,3 +291,8 @@ def tensor_splitter(ishape, split_sizes, axis=2, name="splitter"):
         lambda x: Kops.split(x, indices, axis=axis), output_shape=oshapes, name=name
     )
     return sp_layer
+
+
+def nansum(x, *args, **kwargs):
+    """Like np.nansum, returns the sum treating NaN as 0.0 (and inf as a very large number)."""
+    return sum(nan_to_num(x), *args, **kwargs)

n3fit/src/n3fit/tests/test_backend.py

@@ -1,6 +1,6 @@
 """
-    This module tests the mathematical functions in the n3fit backend
-    and ensures they do the same thing as their numpy counterparts
+This module tests the mathematical functions in the n3fit backend
+and ensures they do the same thing as their numpy counterparts
 """
 
 import operator
@@ -153,3 +153,14 @@ def test_tensor_product():
 
 def test_sum():
     numpy_check(op.sum, np.sum, mode='single')
+
+
+def test_nansum():
+    """Tests that sums with NaN in the arrays work as expected"""
+    arr_nonan = np.array([1.0, 2.0])
+    arr_nan = np.array([2.0, np.nan, 2.0])
+    arr_axis_nan = np.array([[3.0, np.nan], [2.0, 6.0]])
+
+    np.testing.assert_allclose(op.nansum(arr_nonan), np.nansum(arr_nonan))
+    np.testing.assert_allclose(op.nansum(arr_nan), np.nansum(arr_nan))
+    np.testing.assert_allclose(op.nansum(arr_axis_nan, axis=0), np.nansum(arr_axis_nan, axis=0))