diff --git a/src/pipt/localization/distance_localization.py b/src/pipt/localization/distance_localization.py index 04c1f43..90bd05a 100644 --- a/src/pipt/localization/distance_localization.py +++ b/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 diff --git a/tests/assimilation/test_distance_loc.py b/tests/assimilation/test_distance_loc.py index 7a11b66..54b6b82 100644 --- a/tests/assimilation/test_distance_loc.py +++ b/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()