Merge pull request #35 from Aignosi/fix/SIENTIAPDE-2215

fix(metrics): ignore NaN target/prediction pairs in validation metrics (SIENTIAPDE-2215)

model_manager/utils/repository/data_manager_repository.py

@@ -357,7 +357,8 @@ class DataManagerRepository(SientiaMonitoring):
         This helper mirrors the previous TrainingRepository.after_train_calculation
         behavior, assuming that predictions (y_pred/y_train_pred) are already on the
         correct scale for metric calculation (any scaling is handled inside the
-        model wrapper).
+        model wrapper). Pairs with a NaN target or a NaN prediction are left out of
+        mse/mae/r2. With a single valid pair left, r2_val is NaN (R² needs 2 samples).
 
         Args:
             tmr: Training result containing:
@@ -368,6 +369,10 @@ class DataManagerRepository(SientiaMonitoring):
 
         Return:
             TrainModelResult: Same object with mse_val, mae_val and r2_val set.
+
+        Raises:
+            ValueError: If y_pred is not set, val_data and y_pred share no indices,
+                or no pair has both a valid target and a valid prediction.
         """
         if tmr.y_pred is None:
             raise ValueError('y_pred must be set before computing regression metrics')
@@ -405,6 +410,24 @@ class DataManagerRepository(SientiaMonitoring):
         y_true_val = y_true_val.loc[common_index]
         y_pred_val = y_pred_val.loc[common_index]
 
+        # NaN target, or NaN prediction for rows with NaN features, carries no error signal
+        # (SIENTIAPDE-2215).
+        valid = y_true_val.notna() & y_pred_val.notna()
+        if not valid.any():
+            raise ValueError(
+                'No validation rows with both target and prediction: all '
+                f'{len(valid)} overlapping rows have NaN'
+            )
+        if not valid.all():
+            n_valid = int(valid.sum())
+            message = (
+                f'Ignoring {len(valid) - n_valid} of {len(valid)} validation row(s) with NaN '
+                f'target or prediction; regression metrics use the remaining {n_valid}'
+            )
+            self.warning(message, metadata)
+        y_true_val = y_true_val[valid]
+        y_pred_val = y_pred_val[valid]
+
         # Metrics helpers already round to 2 decimals.
         tmr.mse_val = mse(y_true_val, y_pred_val)
         tmr.mae_val = mae(y_true_val, y_pred_val)

tests/utils/repository/test_data_manager_repository.py

@@ -309,6 +309,46 @@ def test_compute_regression_metrics_no_overlap():
         repo.compute_regression_metrics(tmr, MagicMock())
 
 
+def test_compute_regression_metrics_ignores_nan_pairs():
+    """NaN target or NaN prediction rows are left out of mse/mae/r2 (SIENTIAPDE-2215)."""
+    repo = dmr.DataManagerRepository(MagicMock())
+    p = TrainModelParams.from_dict(_minimal_dict_for_prepare())
+    idx = pd.Index([0, 1, 2, 3])
+    tmr = TrainModelResult(
+        params=p,
+        train_data=pd.DataFrame({'t': [1.0, 2.0, 3.0, 4.0]}, index=idx),
+        val_data=pd.DataFrame({'t': [1.0, np.nan, 3.0, 4.0]}, index=idx),
+        y_pred=pd.DataFrame({'p': [1.0, 2.0, np.nan, 4.0]}, index=idx),
+    )
+
+    wrapper = MagicMock()
+    wrapper.model = object()  # no sklearn .regr -> equation extraction is skipped
+    with patch.object(repo, 'warning') as warn:
+        out = repo.compute_regression_metrics(tmr, wrapper)
+
+    assert out.mse_val == pytest.approx(0.0)
+    assert out.mae_val == pytest.approx(0.0)
+    assert out.r2_val == pytest.approx(1.0)
+    warn.assert_called_once()
+    assert 'Ignoring 2 of 4' in warn.call_args.args[0]
+
+
+def test_compute_regression_metrics_all_nan_pairs_raises():
+    repo = dmr.DataManagerRepository(MagicMock())
+    p = TrainModelParams.from_dict(_minimal_dict_for_prepare())
+    idx = pd.Index([0, 1])
+    tmr = TrainModelResult(
+        params=p,
+        train_data=pd.DataFrame({'t': [1.0, 2.0]}, index=idx),
+        val_data=pd.DataFrame({'t': [1.0, 2.0]}, index=idx),
+        y_pred=pd.DataFrame({'p': [np.nan, np.nan]}, index=idx),
+    )
+    wrapper = MagicMock()
+
+    with pytest.raises(ValueError, match='No validation rows with both target and prediction'):
+        repo.compute_regression_metrics(tmr, wrapper)
+
+
 def test_compute_regression_metrics_linear_equation():
     repo = dmr.DataManagerRepository(MagicMock())
     p = TrainModelParams.from_dict(_minimal_dict_for_prepare())