import numpy as np
from .conesRegion import conesRegion
[docs]
class conesMGRegion(conesRegion):
"""
Extends conesRegion with a fine-simulation state for Multi-Grid EnKF.
Adds `stateFine` and `filterStateFine` to handle the projected state
coming from the fine simulation member.
"""
def __init__(self):
super().__init__()
self.stateFine = np.empty((0, 0), dtype="float")
self.paramFine = None
self.stateNvarFine = int()
[docs]
def setStateFine(self, state_fine):
"""
Set the fine state directly.
:param state_fine: Fine projected state array
:type state_fine: numpy ndarray
"""
self.stateFine = state_fine
return
[docs]
def setParamFine(self, param_fine):
"""
Set the fine params directly.
:param param_fine:
:type param_fine: numpy ndarray
"""
self.paramFine = param_fine
return
[docs]
def filterStateFine(self, state_proj, globalStateCells):
"""
From the global projected-fine state matrix, extract the rows
belonging to this region's cells.
:param state_proj: list of 1-D arrays (one per state variable), each of length
n_coarse_cells (total cells across all coarse procs).
:type state_proj: list of numpy ndarray
:param globalStateCells: nested list [proc][cell_id] of global cell ids for coarse mesh
:type globalStateCells: list of list of int
"""
regionCells = [x for xs in self.globalCells for x in xs]
stateCells = [x for xs in globalStateCells for x in xs]
idx = [i for i, cell in enumerate(stateCells) if cell in regionCells]
n_vars = len(state_proj)
if len(idx) == 0 or n_vars == 0:
self.stateFine = np.empty((0, 1), dtype="float")
self.stateNvarFine = n_vars
return
# Build (n_vars * n_cells_region, 1) array – same layout as coarse state
per_var = [state_proj[iv][idx] for iv in range(n_vars)]
self.stateFine = np.vstack(per_var) # shape: (n_vars * n_cells_region,)
self.stateNvarFine = n_vars
return
[docs]
def filterSamplingFine(self, sampling_proj, sampling_order):
"""
From the global fine-projected sampling matrix, extract and reorder
the values belonging to this region's observations
:param sampling_proj: The raw fine-projected sampling column gathered
:type sampling_proj: numpy ndarray
:param sampling_order: A list of observation id ordered by gather
:type sampling_order: list
:returns: None
"""
n_obs = len(self.obsList)
n_var = sum(1 for s in self.obsList[0].var_list if not s.endswith("_std"))
filtered_sampling = np.zeros((n_obs * n_var, 1))
global_order_map = {obs_id: i for i, obs_id in enumerate(sampling_order)}
for i_local, obs in enumerate(self.obsList):
i_global = global_order_map[obs.id]
filtered_sampling[i_local*n_var: (i_local+1)*n_var, 0] = sampling_proj[i_global*n_var: (i_global+1)*n_var]
self.samplingFine = filtered_sampling
return