Distance localization: let parameters without a grid be updated or left out

src/pipt/localization/distance_localization.py

@@ -286,6 +286,12 @@ def __init__(
                 - **dict**: Pre-built ``{(data_type, time, param): LocalizationEntry}``
                   dict (rarely used; prefer the other formats).
 
+            **unlocalized** : list of str, *optional*
+                State parameters that every datum updates with weight one, i.e.
+                that are not localized -- typically scalars, which have no
+                position. A parameter that is neither here nor in an entry gets
+                weight zero: it is not updated. Default: ``[]``.
+
 
         data : pd.DataFrame, optional
             Observed data whose **index** contains the assimilation time
@@ -452,6 +458,15 @@ def __init__(
         self.prior_info    = prior_info if prior_info is not None else {}
         self.ensemble_size = ensemble_size
 
+        # -- parameters updated by every datum with weight one
+        self.unlocalized = [str(p).lower() for p in info.get("unlocalized", [])]
+        known = {p.lower() for p in self.parameters or []}
+        if self.parameters is not None and set(self.unlocalized) - known:
+            raise ValueError(
+                f"'unlocalized' names {sorted(set(self.unlocalized) - known)}, "
+                f"which are not among the state's parameters {sorted(known)}."
+            )
+
         # -- select and instantiate the kernel (optional; entry rows may supply it instead)
         taperfunc = info.get("taper_func")
         if taperfunc is not None and taperfunc not in self._kernel_map:
@@ -520,14 +535,14 @@ def __call__(
 
                 for param in curr_param:
                     key = (data_name.lower(), time, param.lower())
-                    if key in self._entries and self._entries[key].taper is not None:
-                        mask = self._resolve_mask(key)
-                        for i in range(n_obs):
-                            obs_blocks[i].append(mask)
+                    if param.lower() in self.unlocalized:
+                        masks = self._constant_mask(param, n_obs, 1.0)
+                    elif key in self._entries and self._entries[key].taper is not None:
+                        masks = [self._resolve_mask(key)] * n_obs
                     else:
-                        zero_masks = self._zero_mask(param, n_obs)
-                        for i in range(n_obs):
-                            obs_blocks[i].append(zero_masks[i])
+                        masks = self._constant_mask(param, n_obs, 0.0)
+                    for i in range(n_obs):
+                        obs_blocks[i].append(masks[i])
 
                 for blocks in obs_blocks:
                     sparse_blocks = [sparse.csc_matrix(b.reshape(1, -1)) for b in blocks]
@@ -771,9 +786,9 @@ def _resolve_mask(self, key: Tuple[str, float, str]) -> np.ndarray:
         """Return the repositioned spatial mask for an entry key, over the active cells.
 
         The kernel is placed on the full grid, but the state holds only the active
-        cells, and :meth:`_zero_mask` already reduces to them. Returning the full grid
+        cells, and :meth:`_constant_mask` already reduces to them. Returning the full grid
         here made the two disagree: a localized parameter contributed one row per grid
-        cell and an unlocalized one a row per active cell, so the operator came out
+        cell and one without an entry a row per active cell, so the operator came out
         with the wrong number of rows altogether.
         """
         entry = self._entries[key]
@@ -841,15 +856,17 @@ def _place_kernel(self, kernel: np.ndarray, position: List[int]) -> np.ndarray:
         result[z_pos, gx0:gx1, gy0:gy1] = kernel[kx0:kx1, ky0:ky1]
         return result
 
-    def _zero_mask(self, param: str, n_obs: int) -> List[np.ndarray]:
-        """Return zero-valued masks for a parameter with no localization entry."""
-        p       = self.prior_info[param]
-        n_cells = p["nx"] * p["ny"] * p["nz"]
+    def _constant_mask(self, param: str, n_obs: int, value: float) -> List[np.ndarray]:
+        """Masks of ``value`` over all of ``param``: zero for a parameter with no entry, one for an unlocalized one.
 
-        if n_obs > 1:
-            mat = np.zeros((n_obs, n_cells))
-            return [mat[i, self.actnum] if self.actnum is not None else mat[i]
-                    for i in range(n_obs)]
-
-        vec = np.zeros(n_cells)
-        return [vec[self.actnum] if self.actnum is not None else vec]
+        A field's size is its grid's, reduced to the active cells as :meth:`_resolve_mask`
+        does. A parameter without a grid is as many independent scalars as its mean has
+        entries (ref ``StateLayout.from_prior_info``), and has no cells to be inactive.
+        """
+        p = self.prior_info[param]
+        if "nx" in p:
+            vec = np.full(p["nx"] * p["ny"] * p["nz"], value)
+            vec = vec[self.actnum] if self.actnum is not None else vec
+        else:
+            vec = np.full(np.size(p["mean"]), value)
+        return [vec] * n_obs

tests/assimilation/test_distance_loc.py

@@ -476,6 +476,33 @@ def test_unconfigured_param_gives_zero_weights(self):
         assert np.any(dense[:n_cells] > 0),    "perm weights should have non-zero entries"
         np.testing.assert_array_equal(dense[n_cells:], 0.0)
 
+    def test_a_parameter_without_a_grid_and_an_entry_gives_zero_weights(self):
+        """Its zero mask was sized from a grid it does not have: KeyError 'nx'."""
+        prior_info = {"sor": {"mean": [0.2]}, "mult": {"mean": [1.0, 2.0, 3.0]}}
+        loc = DistanceLocalization(
+            _make_info(), data=_make_data(), parameters=["perm", "sor", "mult"], prior_info=prior_info
+        )
+        dense = loc().toarray()
+        n_cells = NZ * NX * NY
+        assert dense.shape == (n_cells + 1 + 3, 1)
+        np.testing.assert_array_equal(dense[n_cells:], 0.0)
+
+    def test_an_unlocalized_parameter_gets_weight_one_from_every_datum(self):
+        data = pd.DataFrame({"pressure": [5], "rate": [3]}, index=[1.0])
+        info = {**_make_info(), "unlocalized": ["SOR", "other"]}
+        grid = {"nx": NX, "ny": NY, "nz": NZ}
+        prior_info = {"perm": grid, "sor": {"mean": [0.2]}, "other": grid}
+        loc = DistanceLocalization(info, data=data, parameters=["perm", "sor", "other"], prior_info=prior_info)
+        dense = loc().toarray()
+        n_cells = NZ * NX * NY
+        assert dense.shape == (2 * n_cells + 1, 2)
+        np.testing.assert_array_equal(dense[n_cells:], 1.0)
+        np.testing.assert_array_equal(dense[:n_cells, 1], 0.0)  # "rate" has no entry for perm
+
+    def test_unlocalized_must_name_parameters_of_the_state(self):
+        with pytest.raises(ValueError, match="not among the state's parameters"):
+            DistanceLocalization({**_make_info(), "unlocalized": ["sro"]}, data=_make_data(), parameters=["perm", "sor"])
+
     def test_two_configured_params_correct_output_shape(self):
         """With two configured parameters the output must span both cell-spaces."""
         data = _make_data()