Filter NaN parameter rows in EnIF update

Surface and field parameters commonly have NaN for inactive grid cells.
When EnIF stacks all parameters into a joint matrix, these NaN values
corrupt boost_linear_regression (np.argmax treats NaN as maximum) and
propagate through StandardScaler, contaminating whichever parameter
sits at index 0 in the joint matrix.

Author: larsevj

src/ert/analysis/_enif_update.py

@@ -3,6 +3,7 @@
 import time
 from collections.abc import Callable, Iterable
 
+import networkx as nx
 import numpy as np
 import polars as pl
 import scipy as sp
@@ -147,12 +148,41 @@ def analysis_EnIF(
         iens_active_index=iens_active_index,
     )
 
-    X_full_scaler = StandardScaler()
-    X_full_scaled = X_full_scaler.fit_transform(X_full.T)
+    # np.argmax in boost regression treats NaN as maximum, corrupting H
+    nan_row_mask = np.any(np.isnan(X_full), axis=1)
+    if nan_row_mask.any():
+        n_nan = int(nan_row_mask.sum())
+        n_all_nan = int(np.all(np.isnan(X_full), axis=1).sum())
+        n_partial_nan = n_nan - n_all_nan
+        log_msg = (
+            f"EnIF: Excluding {n_nan}/{len(nan_row_mask)} parameter rows "
+            f"containing NaN ({n_all_nan} fully inactive"
+        )
+        if n_partial_nan > 0:
+            log_msg += (
+                f", {n_partial_nan} partially active — "
+                f"these will not be updated for any realization"
+            )
+        log_msg += ")"
+        logger.info(log_msg)
+        progress_callback(AnalysisStatusEvent(msg=log_msg))
+
+    X_clean = X_full[~nan_row_mask]
+    if X_clean.shape[0] == 0:
+        msg = "All parameter rows contain NaN — cannot run EnIF update"
+        data = DataSection(
+            header=smoother_snapshot.header,
+            data=smoother_snapshot.csv,
+            extra=smoother_snapshot.extra,
+        )
+        raise ErtAnalysisError(msg, data=data)
+
+    X_clean_scaler = StandardScaler()
+    X_clean_scaled = X_clean_scaler.fit_transform(X_clean.T)
 
     # Call fit: Learn sparse linear map only
     H = linear_boost_ic_regression(
-        U=X_full_scaled,
+        U=X_clean_scaled,
         Y=S.T,
         verbose_level=5,
     )
@@ -166,10 +196,23 @@ def analysis_EnIF(
     for param_group in updated_parameters:
         config_node = source_ensemble.experiment.parameter_configuration[param_group]
         X_local = source_ensemble.load_parameters_numpy(param_group, iens_active_index)
+
+        local_nan_mask = np.any(np.isnan(X_local), axis=1)
+        X_local_clean = X_local[~local_nan_mask]
+
+        if X_local_clean.shape[0] == 0:
+            continue
+
         X_local_scaler = StandardScaler()
-        X_scaled = X_local_scaler.fit_transform(X_local.T)
+        X_scaled = X_local_scaler.fit_transform(X_local_clean.T)
 
         graph_u_sub = config_node.load_parameter_graph()
+        if local_nan_mask.any():
+            nan_nodes = set(np.where(local_nan_mask)[0].tolist())
+            graph_u_sub.remove_nodes_from(nan_nodes)
+            graph_u_sub = nx.convert_node_labels_to_integers(
+                graph_u_sub, ordering="sorted"
+            )
 
         # This works for up to ~10^5 parameters
         Prec_u_sub = fit_precision_cholesky_approximate(
@@ -202,16 +245,18 @@ def analysis_EnIF(
         neighbor_propagation_order=15, verbose_level=1
     )
 
-    X_full = gtmap.transport(
-        X_full_scaled,
+    X_updated = gtmap.transport(
+        X_clean_scaled,
         S.T,
         observation_values,
         update_indices=update_indices,
         iterative=True,
         verbose_level=5,
         seed=random_seed,
     )
-    X_full = X_full_scaler.inverse_transform(X_full).T
+    X_updated = X_clean_scaler.inverse_transform(X_updated).T
+
+    X_full[~nan_row_mask] = X_updated
 
     # Iterate over parameters to store the updated ensemble
     log_msg = f"Storing {len(updated_parameters)} updated parameter groups"

tests/ert/unit_tests/analysis/test_es_update.py

@@ -5,6 +5,7 @@
 import numpy as np
 import polars as pl
 import pytest
+import xarray as xr
 from tabulate import tabulate
 
 from ert.analysis import (
@@ -13,6 +14,7 @@
     build_strategy_map,
     smoother_update,
 )
+from ert.analysis._enif_update import enif_update
 from ert.analysis._es_update import _create_combined_ensemble_mask
 from ert.analysis._update_commons import (
     _compute_observation_statuses,
@@ -26,6 +28,7 @@
     GenKwConfig,
     ObservationSettings,
     OutlierSettings,
+    SurfaceConfig,
 )
 from ert.storage import Ensemble, open_storage
 
@@ -1293,3 +1296,238 @@ def test_that_create_combined_ensemble_mask_handles_different_length_masks(
 ) -> None:
     result = _create_combined_ensemble_mask(ens_mask, active_realizations)
     np.testing.assert_array_equal(result, expected)
+
+
+def _make_enif_genkw(name: str) -> dict[str, Any]:
+    return GenKwConfig(
+        name=name,
+        group=name,
+        distribution={"name": "uniform", "min": 0.8, "max": 1.2},
+    ).model_dump(mode="json")
+
+
+def _populate_enif_prior(
+    prior, rng, *, ncol, nrow, ensemble_size, num_obs, genkw_names
+):
+    inactive_mask = rng.random((ncol, nrow)) < 0.55
+
+    for real in range(ensemble_size):
+        vals = rng.normal(1.0, 0.1, (ncol, nrow)).astype(np.float32)
+        vals[inactive_mask] = np.nan
+        prior.save_parameters(
+            xr.Dataset({"values": (["x", "y"], vals)}), "HEIMDAL_TOP", real
+        )
+
+    for name in genkw_names:
+        prior.save_parameters(
+            dataset=pl.concat(
+                [
+                    pl.DataFrame(
+                        {name: [float(rng.uniform(0.8, 1.2))], "realization": i}
+                    )
+                    for i in range(ensemble_size)
+                ],
+                how="vertical",
+            )
+        )
+
+    for iens in range(ensemble_size):
+        prior.save_response(
+            "gen_data",
+            pl.DataFrame(
+                {
+                    "response_key": "RESPONSE",
+                    "report_step": pl.Series(
+                        np.zeros(num_obs, dtype=int), dtype=pl.UInt16
+                    ),
+                    "index": pl.Series(np.arange(num_obs), dtype=pl.UInt16),
+                    "values": pl.Series(rng.normal(1, 0.1, num_obs), dtype=pl.Float32),
+                }
+            ),
+            iens,
+        )
+
+
+@pytest.fixture
+def enif_surface_config():
+    return SurfaceConfig(
+        name="HEIMDAL_TOP",
+        forward_init=True,
+        update=True,
+        ncol=20,
+        nrow=25,
+        xori=0.0,
+        yori=0.0,
+        xinc=25.0,
+        yinc=25.0,
+        rotation=0.0,
+        yflip=1,
+        forward_init_file="init_surf",
+        output_file="out_surf",
+        base_surface_path="base_surf",
+    )
+
+
+@pytest.fixture
+def enif_observations():
+    return [
+        {
+            "type": "general_observation",
+            "name": f"OBS_{i}",
+            "data": "RESPONSE",
+            "restart": 0,
+            "index": i,
+            "value": 1.0,
+            "error": 0.1,
+        }
+        for i in range(5)
+    ]
+
+
+@pytest.mark.filterwarnings("ignore::scipy.linalg.LinAlgWarning")
+@pytest.mark.filterwarnings("ignore::RuntimeWarning")
+def test_that_enif_with_nan_surface_cells_does_not_contaminate_genkw(
+    tmp_path, enif_surface_config, enif_observations
+):
+    rng = np.random.default_rng(42)
+    ensemble_size = 20
+
+    with open_storage(tmp_path, mode="w") as storage:
+        experiment = storage.create_experiment(
+            name="enif_surface_nan",
+            experiment_config={
+                "parameter_configuration": [
+                    _make_enif_genkw("VOLUME_MULT"),
+                    _make_enif_genkw("IN_MULT"),
+                    enif_surface_config.model_dump(mode="json"),
+                ],
+                "response_configuration": [
+                    GenDataConfig(keys=["RESPONSE"]).model_dump(mode="json")
+                ],
+                "observations": enif_observations,
+            },
+        )
+        prior = storage.create_ensemble(
+            experiment, ensemble_size=ensemble_size, iteration=0, name="prior"
+        )
+        _populate_enif_prior(
+            prior,
+            rng,
+            ncol=enif_surface_config.ncol,
+            nrow=enif_surface_config.nrow,
+            ensemble_size=ensemble_size,
+            num_obs=5,
+            genkw_names=["VOLUME_MULT", "IN_MULT"],
+        )
+
+        posterior = storage.create_ensemble(
+            experiment,
+            ensemble_size=ensemble_size,
+            iteration=1,
+            name="posterior_enif",
+            prior_ensemble=prior,
+        )
+        enif_update(
+            prior,
+            posterior,
+            observations=experiment.observation_keys,
+            parameters=list(experiment.parameter_configuration.keys()),
+            random_seed=42,
+        )
+
+        iens = np.arange(ensemble_size)
+        post_vol = posterior.load_parameters_numpy("VOLUME_MULT", iens)
+        post_in = posterior.load_parameters_numpy("IN_MULT", iens)
+        prior_surf = prior.load_parameters_numpy("HEIMDAL_TOP", iens)
+        post_surf = posterior.load_parameters_numpy("HEIMDAL_TOP", iens)
+
+        assert np.all(np.isfinite(post_vol)), (
+            f"VOLUME_MULT should be finite, "
+            f"got {np.isnan(post_vol).sum()} NaN out of {post_vol.size}"
+        )
+        assert np.all(np.isfinite(post_in)), (
+            f"IN_MULT should be finite, "
+            f"got {np.isnan(post_in).sum()} NaN out of {post_in.size}"
+        )
+
+        nan_mask = np.isnan(prior_surf)
+        assert np.array_equal(np.isnan(post_surf), nan_mask), (
+            "Inactive surface cells (NaN) should be preserved after update"
+        )
+        active = ~nan_mask
+        assert np.all(np.isfinite(post_surf[active])), (
+            "Active surface cells should remain finite after update"
+        )
+        assert not np.array_equal(post_surf[active], prior_surf[active]), (
+            "Active surface cells should be updated, not identical to prior"
+        )
+
+
+@pytest.mark.filterwarnings("ignore::scipy.linalg.LinAlgWarning")
+@pytest.mark.filterwarnings("ignore::RuntimeWarning")
+def test_that_enif_nan_filtering_preserves_genkw_regardless_of_parameter_order(
+    tmp_path, enif_surface_config, enif_observations
+):
+    rng = np.random.default_rng(42)
+    ensemble_size = 20
+
+    with open_storage(tmp_path, mode="w") as storage:
+        experiment = storage.create_experiment(
+            name="enif_reversed_order",
+            experiment_config={
+                "parameter_configuration": [
+                    enif_surface_config.model_dump(mode="json"),
+                    _make_enif_genkw("IN_MULT"),
+                    _make_enif_genkw("VOLUME_MULT"),
+                ],
+                "response_configuration": [
+                    GenDataConfig(keys=["RESPONSE"]).model_dump(mode="json")
+                ],
+                "observations": enif_observations,
+            },
+        )
+        prior = storage.create_ensemble(
+            experiment, ensemble_size=ensemble_size, iteration=0, name="prior"
+        )
+        _populate_enif_prior(
+            prior,
+            rng,
+            ncol=enif_surface_config.ncol,
+            nrow=enif_surface_config.nrow,
+            ensemble_size=ensemble_size,
+            num_obs=5,
+            genkw_names=["IN_MULT", "VOLUME_MULT"],
+        )
+
+        posterior = storage.create_ensemble(
+            experiment,
+            ensemble_size=ensemble_size,
+            iteration=1,
+            name="posterior_enif",
+            prior_ensemble=prior,
+        )
+        enif_update(
+            prior,
+            posterior,
+            observations=experiment.observation_keys,
+            parameters=list(experiment.parameter_configuration.keys()),
+            random_seed=42,
+        )
+
+        iens = np.arange(ensemble_size)
+
+        for name in ["IN_MULT", "VOLUME_MULT"]:
+            post_vals = posterior.load_parameters_numpy(name, iens)
+            assert np.all(np.isfinite(post_vals)), (
+                f"{name} should be finite regardless of parameter order"
+            )
+
+        prior_surf = prior.load_parameters_numpy("HEIMDAL_TOP", iens)
+        post_surf = posterior.load_parameters_numpy("HEIMDAL_TOP", iens)
+        active = ~np.isnan(prior_surf)
+        assert np.all(np.isfinite(post_surf[active])), (
+            "Active surface cells should remain finite after update"
+        )
+        assert not np.array_equal(post_surf[active], prior_surf[active]), (
+            "Active surface cells should be updated, not identical to prior"
+        )