Use blobs in experiment storage for rho matrices

This includes rhomatrix event, produced by distance based localization,
which yields saving the rho matrix on the experiment storage level.
Rho matrix has a name and observation keys attributes. They represent
the full description of the rho matrix.

Internally, the blobs themselves are responsible for storing the data as
the code is used on both; experiment and ensemble level.
Since the Rho matrix was not in the storage before, no storage migration
in needed.

Author: xjules

src/ert/analysis/_es_update.py

@@ -46,6 +46,7 @@
 
     from ert.config import ParameterConfig
     from ert.storage import Ensemble
+    from ert.storage.local_experiment import LocalExperiment
 
 logger = logging.getLogger(__name__)
 
@@ -215,6 +216,9 @@ def perform_ensemble_update(
             ypos=filtered_data["north"].to_numpy()[has_location],
             main_range=filtered_data["radius"].to_numpy()[has_location],
             location_mask=has_location,
+            observation_keys=filtered_data["observation_key"]
+            .to_numpy()[has_location]
+            .tolist(),
         )
 
     obs_context = ObservationContext(
@@ -298,6 +302,7 @@ def build_strategy_map(
     correlation_threshold: Callable[[int], float],
     *,
     progress_callback: Callable[[AnalysisEvent], None] | None = None,
+    experiment: LocalExperiment | None = None,
 ) -> dict[str, UpdateStrategy]:
     """Build a mapping from parameter group names to update strategies.
 
@@ -318,6 +323,8 @@ def build_strategy_map(
         threshold.
     progress_callback : Callable[[AnalysisEvent], None] | None
         Callback for reporting progress.
+    experiment : LocalExperiment | None
+        Optional experiment for loading cached rho matrices.
 
     Returns
     -------
@@ -330,11 +337,11 @@ def build_strategy_map(
     strategy_map: dict[str, UpdateStrategy] = {}
 
     field_distance_strategy = DistanceLocalizationUpdate(
-        enkf_truncation, Field, progress_callback
+        enkf_truncation, Field, progress_callback, experiment
     )
 
     surface_distance_strategy = DistanceLocalizationUpdate(
-        enkf_truncation, SurfaceConfig, progress_callback
+        enkf_truncation, SurfaceConfig, progress_callback, experiment
     )
 
     global_strategy = GlobalESUpdate(

src/ert/analysis/_update_strategies/_distance.py

@@ -2,16 +2,22 @@
 
 from __future__ import annotations
 
+import io
 import time
 from collections.abc import Callable
 from datetime import timedelta
 from typing import TYPE_CHECKING
 
 import humanize
 import numpy as np
+import scipy as sp
 from iterative_ensemble_smoother import LocalizedESMDA
 
-from ert.analysis.event import AnalysisEvent, AnalysisStatusEvent
+from ert.analysis.event import (
+    AnalysisEvent,
+    AnalysisRhoMatrixEvent,
+    AnalysisStatusEvent,
+)
 from ert.config import Field, SurfaceConfig
 from ert.field_utils import (
     AxisOrientation,
@@ -27,6 +33,7 @@
     import numpy.typing as npt
 
     from ert.config import ParameterConfig
+    from ert.storage.local_experiment import LocalExperiment
 
     from ._protocol import ObservationContext
 
@@ -39,10 +46,12 @@ def __init__(
         enkf_truncation: float,
         param_type: type[Field | SurfaceConfig],
         progress_callback: Callable[[AnalysisEvent], None],
+        experiment: LocalExperiment | None = None,
     ) -> None:
         self._enkf_truncation = enkf_truncation
         self._param_type = param_type
         self._progress_callback = progress_callback
+        self._experiment = experiment
         self._obs_loc: ObservationLocations | None = None
         self._smoother: LocalizedESMDA | None = None
         self._ensemble_size: int = 0
@@ -125,6 +134,19 @@ def update(
 
         return result
 
+    def _load_rho_from_storage(
+        self, param_name: str
+    ) -> npt.NDArray[np.floating] | None:
+        """Try to load a cached rho matrix from experiment blob storage."""
+        if self._experiment is None:
+            return None
+        obs_keys = (
+            self._obs_loc.observation_keys
+            if self._obs_loc is not None and self._obs_loc.observation_keys
+            else None
+        )
+        return self._experiment.load_rho_matrix(param_name, obs_keys)
+
     def _full_localization_matrix(
         self,
         num_params: int,
@@ -170,32 +192,50 @@ def _update_field(
         if ertbox.axis_orientation is None:
             raise ValueError("Field grid axis orientation must be defined")
 
-        xpos, ypos = transform_positions_to_local_field_coordinates(
-            ertbox.origin,
-            ertbox.rotation_angle,
-            self._obs_loc.xpos,
-            self._obs_loc.ypos,
-        )
+        cached_rho = self._load_rho_from_storage(param_config.name)
+        if cached_rho is not None:
+            rho_2d = cached_rho
+        else:
+            xpos, ypos = transform_positions_to_local_field_coordinates(
+                ertbox.origin,
+                ertbox.rotation_angle,
+                self._obs_loc.xpos,
+                self._obs_loc.ypos,
+            )
 
-        ellipse_rotation = transform_local_ellipse_angle_to_local_coords(
-            ertbox.rotation_angle,
-            np.zeros_like(self._obs_loc.main_range),
-        )
+            ellipse_rotation = transform_local_ellipse_angle_to_local_coords(
+                ertbox.rotation_angle,
+                np.zeros_like(self._obs_loc.main_range),
+            )
 
-        rho_matrix = calc_rho_for_2d_grid_layer(
-            nx=ertbox.nx,
-            ny=ertbox.ny,
-            xinc=ertbox.xinc,
-            yinc=ertbox.yinc,
-            obs_xpos=xpos,
-            obs_ypos=ypos,
-            obs_main_range=self._obs_loc.main_range,
-            obs_perp_range=self._obs_loc.main_range,
-            obs_anisotropy_angle=ellipse_rotation,
-            axis_orientation=ertbox.axis_orientation,
-        )
+            rho_matrix = calc_rho_for_2d_grid_layer(
+                nx=ertbox.nx,
+                ny=ertbox.ny,
+                xinc=ertbox.xinc,
+                yinc=ertbox.yinc,
+                obs_xpos=xpos,
+                obs_ypos=ypos,
+                obs_main_range=self._obs_loc.main_range,
+                obs_perp_range=self._obs_loc.main_range,
+                obs_anisotropy_angle=ellipse_rotation,
+                axis_orientation=ertbox.axis_orientation,
+            )
 
-        rho_2d = rho_matrix.reshape(ertbox.nx * ertbox.ny, -1)
+            rho_2d = rho_matrix.reshape(ertbox.nx * ertbox.ny, -1)
+
+            if self._obs_loc.observation_keys:
+                rho_sparse = sp.sparse.csc_matrix(rho_2d)
+                buf = io.BytesIO()
+                sp.sparse.save_npz(buf, rho_sparse)
+                self._progress_callback(
+                    AnalysisRhoMatrixEvent(
+                        param_name=param_config.name,
+                        observation_keys=self._obs_loc.observation_keys,
+                        shape=rho_2d.shape,
+                        data_type=str(rho_2d.dtype),
+                        matrix_bytes=buf.getvalue(),
+                    )
+                )
 
         for param_batch_idx in batches:
             update_idx = param_batch_idx[non_zero_variance_mask[param_batch_idx]]
@@ -242,37 +282,56 @@ def _update_surface(
 
         assert self._obs_loc is not None
 
-        xpos, ypos = transform_positions_to_local_field_coordinates(
-            (param_config.xori, param_config.yori),
-            param_config.rotation,
-            self._obs_loc.xpos,
-            self._obs_loc.ypos,
-        )
+        cached_rho = self._load_rho_from_storage(param_config.name)
+        if cached_rho is not None:
+            rho_flat = cached_rho
+        else:
+            xpos, ypos = transform_positions_to_local_field_coordinates(
+                (param_config.xori, param_config.yori),
+                param_config.rotation,
+                self._obs_loc.xpos,
+                self._obs_loc.ypos,
+            )
 
-        rotation_angle = transform_local_ellipse_angle_to_local_coords(
-            param_config.rotation,
-            np.zeros_like(self._obs_loc.main_range, dtype=np.float64),
-        )
+            rotation_angle = transform_local_ellipse_angle_to_local_coords(
+                param_config.rotation,
+                np.zeros_like(self._obs_loc.main_range, dtype=np.float64),
+            )
 
-        if param_config.yflip != 1:
-            raise ValueError(
-                f"Expected SurfaceConfig.yflip == 1, got {param_config.yflip}"
+            if param_config.yflip != 1:
+                raise ValueError(
+                    f"Expected SurfaceConfig.yflip == 1, got {param_config.yflip}"
+                )
+
+            rho_matrix = calc_rho_for_2d_grid_layer(
+                nx=param_config.ncol,
+                ny=param_config.nrow,
+                xinc=param_config.xinc,
+                yinc=param_config.yinc,
+                obs_xpos=xpos,
+                obs_ypos=ypos,
+                obs_main_range=self._obs_loc.main_range,
+                obs_perp_range=self._obs_loc.main_range,
+                obs_anisotropy_angle=rotation_angle,
+                axis_orientation=AxisOrientation.LEFT_HANDED,
             )
 
-        rho_matrix = calc_rho_for_2d_grid_layer(
-            nx=param_config.ncol,
-            ny=param_config.nrow,
-            xinc=param_config.xinc,
-            yinc=param_config.yinc,
-            obs_xpos=xpos,
-            obs_ypos=ypos,
-            obs_main_range=self._obs_loc.main_range,
-            obs_perp_range=self._obs_loc.main_range,
-            obs_anisotropy_angle=rotation_angle,
-            axis_orientation=AxisOrientation.LEFT_HANDED,
-        )
+            rho_flat = rho_matrix.reshape(-1, rho_matrix.shape[-1])
+
+            if self._obs_loc.observation_keys:
+                rho_sparse = sp.sparse.csc_matrix(rho_flat)
+                buf = io.BytesIO()
+                sp.sparse.save_npz(buf, rho_sparse)
+                self._progress_callback(
+                    AnalysisRhoMatrixEvent(
+                        param_name=param_config.name,
+                        observation_keys=self._obs_loc.observation_keys,
+                        shape=rho_flat.shape,
+                        data_type=str(rho_flat.dtype),
+                        matrix_bytes=buf.getvalue(),
+                    )
+                )
 
-        rho_flat = rho_matrix.reshape(-1, rho_matrix.shape[-1])
         for param_batch_idx in batches:
             update_idx = param_batch_idx[non_zero_variance_mask[param_batch_idx]]
             if update_idx.size == 0:

src/ert/analysis/_update_strategies/_protocol.py

@@ -4,7 +4,7 @@
 
 import time
 from collections.abc import Callable, Sequence
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Protocol, Self
 
 import numpy as np
@@ -101,6 +101,9 @@ class ObservationLocations:
     location_mask: npt.NDArray[np.bool_]
     """Boolean mask indicating which observations have valid location data."""
 
+    observation_keys: list[str] = field(default_factory=list)
+    """Observation keys for the located observations (same order as xpos/ypos)."""
+
 
 @dataclass(frozen=True)
 class ObservationContext:

src/ert/analysis/event.py

@@ -102,3 +102,12 @@ class AnalysisScalingEvent(AnalysisEvent):
     scaling_bytes: bytes = Field(exclude=True)
     num_observations: int
     num_groups: int
+
+
+class AnalysisRhoMatrixEvent(AnalysisEvent):
+    event_type: Literal["AnalysisRhoMatrixEvent"] = "AnalysisRhoMatrixEvent"
+    param_name: str
+    observation_keys: list[str]
+    shape: tuple[int, int]
+    data_type: str
+    matrix_bytes: bytes = Field(exclude=True)

src/ert/run_models/update_run_model.py

@@ -10,6 +10,7 @@
     AnalysisErrorEvent,
     AnalysisEvent,
     AnalysisMatrixEvent,
+    AnalysisRhoMatrixEvent,
     AnalysisScalingEvent,
     AnalysisStatusEvent,
     AnalysisTimeEvent,
@@ -67,6 +68,7 @@ def update_ensemble_parameters(
             enkf_truncation=self.analysis_settings.enkf_truncation,
             correlation_threshold=self.analysis_settings.correlation_threshold,
             progress_callback=progress_callback,
+            experiment=prior.experiment,
         )
 
         smoother_update(
@@ -195,6 +197,8 @@ def send_smoother_event(
                 ensemble.save_blob(event)
             case AnalysisScalingEvent():
                 ensemble.save_blob(event)
+            case AnalysisRhoMatrixEvent():
+                ensemble.experiment.save_blob(event)
             case AnalysisCompleteEvent():
                 ensemble.save_blob(event)
                 self.send_event(