import datetime
import os
from typing import Any, Dict, Iterable, Iterator, Sequence, Tuple, Union, Set
import numpy as np
import pandas as pd
from .....utils.iterators import iter_attributes, iter_tracers
from .....utils.path import link
from .base_function import BaseFunction
from .....utils.check.errclass import CifAttributeError, CifFileNotFoundError, CifValueError
# Aliases for type hinting
Mode = Any
ControlVect = Any
DatetimeLike = Union[datetime.datetime, pd.Timestamp, np.datetime64]
Component = str
Parameter = str
TracerId = Tuple[Component, Parameter]
[docs]
def chain_outputs(
in_out_dict: Dict[TracerId, Set[TracerId]],
list_output: Set[TracerId]
) -> Set[TracerId]:
"""For each tracer id in 'list_output', if it is an input tracer (a key in
'in_out_dict') then replace it with the corresponding set of output tracers.
This is done recursively.
Args:
in_out_dict (dict((str, str) -> set of (str, str))): dictionary with
input tracers as keys and a set of output tracers as values
list_output (set of (str, str)): set of output tracers
Returns:
set of (str, str): chained outputs
"""
chained_outputs = set()
for output in list_output:
if output in in_out_dict:
chained_outputs = chained_outputs.union(
chain_outputs(in_out_dict, in_out_dict[output]))
else:
chained_outputs.add(output)
return chained_outputs
[docs]
class BaseFunctionSamplingBatch(BaseFunction, Iterable):
def __init__(
self,
mode: Mode,
batch_index: int,
indices: Union[Sequence[int], np.ndarray],
components: Sequence[Component],
parameters: Sequence[Parameter],
date_start: DatetimeLike,
date_end: DatetimeLike
) -> None:
"""A class representing a sampling batch of base functions.
BaseFunctionSamplingBatch objects are iterable and produce a list
of all non-ignored base function within th batch.
Args:
mode (Mode): the mode plugin
batch_index (int): sampling batch index
indices (array or list of int): controlvect indices
components (str): component names
parameters (list of str): parameter names
date_start (datetime): start of the base functions time window
date_end (datetime): end of the base functions time window
"""
if batch_index > 999999:
raise CifValueError("too many response function sampling batch")
super().__init__(mode, -1, "_Batch_", "_Batch_", date_start, date_end)
self.index = batch_index
self.indices = indices
self.components = components
self.parameters = parameters
self.name = f"base_function_batch_{batch_index:06d}"
# Base functions in the sampling batch
self._base_function_list = [
BaseFunction(mode, i, comp, param, date_start, date_end)
for i, comp, param in zip(indices, components, parameters)
]
# Number of base functions that will actually be run
self.n_samples = len(list(self))
# List of tracers in the batch sample
self.tracer_list = list(zip(self.components, self.parameters))
# List of tracers taht should not be propagated by the sampling batch
self.dont_propagate = [
[comp, param] for (comp, param), _ in iter_tracers(self.datavect)
if (comp, param) not in self.tracer_list
]
# Mapping between input and output tracers
self.in_out_dict = get_transforms_in_out_mapping(self.controlvect)
# Filtering out the input tracers that are not in the batch
self.in_out_dict = {
trid_in: list_trid_out
for trid_in, list_trid_out in self.in_out_dict.items()
if trid_in in self.tracer_list
}
# Output parameters from controlvect transforms
out_parameters = set(
param_out
for list_trid_out in self.in_out_dict.values()
for _, param_out in list_trid_out
)
# Output parameters of the response functions in the sampling batch
out_parameters = out_parameters.union(self.parameters)
# Observation tracers taht are independant from this base function
# sampling batch
self.independant_obs = [
[comp, param]
for (comp, param), tracer in iter_tracers(self.datavect)
if tracer.isobs and param not in out_parameters
]
def __iter__(self) -> Iterator[BaseFunction]:
return (base_function for base_function in self._base_function_list
if not base_function.is_ignored())
def __repr__(self) -> str:
return f"<BaseFunctionSamplingBatch indices={self.indices}>"
def __str__(self) -> str:
s = f"base function batch sampling {self.index:06d}:"
for base_function in self.iter_all():
if base_function.is_ignored():
s += f"\n - {base_function} (ignored)"
else:
s += f"\n - {base_function}"
return s
[docs]
def iter_all(self) -> Iterator[BaseFunction]:
"""All base functions within the batch (ignored or not)"""
return iter(self._base_function_list)
@property
def obsdir(self) -> str:
raise CifAttributeError(
"BaseFunctionSamplingBatch has no 'obsdir' attribute")
[docs]
def is_ignored(self) -> bool:
# Base functions batches are always run
return False
[docs]
def dump_controlvect(self) -> None:
self.update_inicond_date()
samples_var = np.zeros((self.n_samples, self.controlvect.dim))
if self.run_mode == "fwd":
# Setting control vector x
self.controlvect.x = np.zeros(self.controlvect.dim)
for sample_ind, base_function in enumerate(self):
x_ind = base_function.index
self.controlvect.x[x_ind] = self.controlvect.x[x_ind]
samples_var[sample_ind, x_ind] = self.controlvect.x[x_ind]
self.controlvect.x_ens = samples_var
elif self.run_mode == "tl":
# Setting control vector dx
self.controlvect.dx = np.zeros(self.controlvect.dim)
for sample_ind, base_function in enumerate(self):
x_ind = base_function.index
self.controlvect.dx[x_ind] = 1.0
samples_var[sample_ind, x_ind] = 1.0
self.controlvect.dx_ens = samples_var
self.controlvect.x_ens = np.repeat(
self.controlvect.x[np.newaxis, :], self.n_samples, axis=0)
os.makedirs(self.rundir, exist_ok=True)
self.controlvect.dump(self.controlvect_path, ensemble=True)
[docs]
def fetch_obsvect(self) -> None:
# Getting run directory
if not os.path.isdir(self.outdir):
raise CifFileNotFoundError(
f"Sampling batch {self.index:06d} did not produce the "
f"requiered outputs '{self.outdir}'"
)
# List of lists to list of tuples
independant_obs = [tuple(trid) for trid in self.independant_obs]
# Looping over all tracers in the observation vector
for (component, parameter), tracer in iter_tracers(self.datavect):
if not tracer.isobs:
continue
# Looping over response functions in sampling batch
for sample_ind, base_function in enumerate(self):
# Skipping ignored tracers response function
if base_function.is_ignored():
continue
dst_path = os.path.join(
base_function.obsdir, component, parameter, "monitor.nc")
if os.path.isfile(dst_path):
continue
sample_name = f"{parameter}__sample#{sample_ind:03d}"
src_path = os.path.join(
self.outdir, component, sample_name, "monitor.nc")
if not os.path.isfile(src_path):
if (component, parameter) in independant_obs:
src_path = os.path.join(
self.outdir, component, parameter, "monitor.nc")
else:
raise CifFileNotFoundError(
f"Sampling batch {self.index:06d} did not produce "
f"the requiered outputs '{src_path}'"
)
os.makedirs(os.path.dirname(dst_path), exist_ok=True)
link(src_path, dst_path)
