option to ignore NaN in CombineResultsRule

Author: wschoonveld

decoimpact/business/entities/rules/combine_results_rule.py

@@ -38,13 +38,18 @@ def __init__(
     ):
         super().__init__(name, input_variable_names)
         self._operation_type: MultiArrayOperationType = operation_type
+        self._ignore_nan = ignore_nan
         self._operations = self._create_operations()
 
     @property
     def operation_type(self) -> MultiArrayOperationType:
         """Name of the rule"""
         return self._operation_type
 
+    @property
+    def ignore_nan(self) -> bool:
+        return self._ignore_nan
+
     def validate(self, logger: ILogger) -> bool:
         if self._operation_type not in self._operations:
 
@@ -90,17 +95,30 @@ def execute(
         return result_variable
 
     def _create_operations(self) -> dict[MultiArrayOperationType, Callable]:
-        return {
-            MultiArrayOperationType.MULTIPLY: lambda npa: _np.prod(npa, axis=0),
-            MultiArrayOperationType.MIN: lambda npa: _np.min(npa, axis=0),
-            MultiArrayOperationType.MAX: lambda npa: _np.max(npa, axis=0),
-            MultiArrayOperationType.AVERAGE: lambda npa: _np.average(npa, axis=0),
-            MultiArrayOperationType.MEDIAN: lambda npa: _np.median(npa, axis=0),
-            MultiArrayOperationType.ADD: lambda npa: _np.sum(npa, axis=0),
-            MultiArrayOperationType.SUBTRACT: lambda npa: _np.subtract(
-                npa[0], _np.sum(npa[1:], axis=0)
-            ),
-        }
+        if self.ignore_nan:
+            return {
+                MultiArrayOperationType.MULTIPLY: lambda npa: _np.prod(npa, axis=0),
+                MultiArrayOperationType.MIN: lambda npa: _np.nanmin(npa, axis=0),
+                MultiArrayOperationType.MAX: lambda npa: _np.nanmax(npa, axis=0),
+                MultiArrayOperationType.AVERAGE: lambda npa: _np.nanmean(npa, axis=0),
+                MultiArrayOperationType.MEDIAN: lambda npa: _np.nanmedian(npa, axis=0),
+                MultiArrayOperationType.ADD: lambda npa: _np.nansum(npa, axis=0),
+                MultiArrayOperationType.SUBTRACT: lambda npa: _np.subtract(
+                    npa[0], _np.nansum(npa[1:], axis=0)
+                ),
+            }
+        else:
+            return {
+                MultiArrayOperationType.MULTIPLY: lambda npa: _np.prod(npa, axis=0),
+                MultiArrayOperationType.MIN: lambda npa: _np.min(npa, axis=0),
+                MultiArrayOperationType.MAX: lambda npa: _np.max(npa, axis=0),
+                MultiArrayOperationType.AVERAGE: lambda npa: _np.average(npa, axis=0),
+                MultiArrayOperationType.MEDIAN: lambda npa: _np.median(npa, axis=0),
+                MultiArrayOperationType.ADD: lambda npa: _np.sum(npa, axis=0),
+                MultiArrayOperationType.SUBTRACT: lambda npa: _np.subtract(
+                    npa[0], _np.sum(npa[1:], axis=0)
+                ),
+            }
 
     def _check_dimensions(self, np_arrays: List[_np.ndarray]) -> bool:
         """Brief check if all the arrays to be combined have the

tests/business/entities/rules/test_combine_results_rule.py

@@ -193,7 +193,7 @@ def test_all_operations_incl_nan(
         (MultiArrayOperationType.MAX, [20, 12, 24]),
         (MultiArrayOperationType.MULTIPLY, [_np.nan, 420, 432]),
         (MultiArrayOperationType.AVERAGE, [12, 8, 11]),
-        (MultiArrayOperationType.MEDIAN, [20, 7, 6]),
+        (MultiArrayOperationType.MEDIAN, [12, 7, 6]),
         (MultiArrayOperationType.ADD, [24, 24, 33]),
         (MultiArrayOperationType.SUBTRACT, [16, -10, -27]),
     ],