def finalize_data(
self,
dataset_model: DatasetModel,
dataset: xr.Dataset,
is_full_model: bool = False,
as_global: bool = False,
):
if not is_full_model:
global_dimension = dataset_model.get_global_dimension()
model_dimension = dataset_model.get_model_dimension()
dataset.coords["coherent_artifact_order"] = np.arange(
1, self.order + 1)
response_dimensions = (model_dimension, "coherent_artifact_order")
if dataset_model.is_index_dependent() is True:
response_dimensions = (global_dimension, *response_dimensions)
dataset["coherent_artifact_response"] = (
response_dimensions,
dataset.matrix.sel(clp_label=self.compartments()).values,
)
dataset["coherent_artifact_associated_spectra"] = (
(global_dimension, "coherent_artifact_order"),
dataset.clp.sel(clp_label=self.compartments()).values,
)
retrieve_irf(dataset_model, dataset,
dataset_model.get_global_dimension())
def retrieve_irf(dataset_model: DatasetModel, dataset: xr.Dataset,
global_dimension: str):
if not isinstance(dataset_model.irf, IrfMultiGaussian) or "irf" in dataset:
return
irf = dataset_model.irf
model_dimension = dataset_model.get_model_dimension()
dataset["irf"] = (
(model_dimension),
irf.calculate(
index=0,
global_axis=dataset.coords[global_dimension].values,
model_axis=dataset.coords[model_dimension].values,
).data,
)
center = irf.center if isinstance(irf.center, list) else [irf.center]
width = irf.width if isinstance(irf.width, list) else [irf.width]
dataset["irf_center"] = ("irf_nr",
center) if len(center) > 1 else center[0]
dataset["irf_width"] = ("irf_nr", width) if len(width) > 1 else width[0]
if irf.shift is not None:
dataset["irf_shift"] = (global_dimension,
[center[0] - p.value for p in irf.shift])
if isinstance(irf, IrfSpectralMultiGaussian) and irf.dispersion_center:
dataset["irf_center_location"] = (
("irf_nr", global_dimension),
irf.calculate_dispersion(dataset.coords["spectral"].values),
)
# center_dispersion_1 for backwards compatibility (0.3-0.4.1)
dataset["center_dispersion_1"] = dataset["irf_center_location"].sel(
irf_nr=0)
def retrieve_species_associated_data(
dataset_model: DatasetModel,
dataset: xr.Dataset,
species: list[str],
species_dimension: str,
global_dimension: str,
name: str,
is_full_model: bool,
as_global: bool,
):
model_dimension = dataset_model.get_model_dimension()
if as_global:
model_dimension, global_dimension = global_dimension, model_dimension
dataset.coords[species_dimension] = species
matrix = dataset.global_matrix if as_global else dataset.matrix
clp_dim = "global_clp_label" if as_global else "clp_label"
if len(dataset.matrix.shape) == 3:
# index dependent
dataset["species_concentration"] = (
(
global_dimension,
model_dimension,
species_dimension,
),
matrix.sel({
clp_dim: species
}).values,
)
else:
# index independent
dataset["species_concentration"] = (
(
model_dimension,
species_dimension,
),
matrix.sel({
clp_dim: species
}).values,
)
if not is_full_model:
dataset[f"species_associated_{name}"] = (
(
global_dimension,
species_dimension,
),
dataset.clp.sel(clp_label=species).data,
)
def finalize_data(
self,
dataset_model: DatasetModel,
dataset: xr.Dataset,
is_full_model: bool = False,
as_global: bool = False,
):
species_dimension = "spectral_species" if as_global else "species"
if species_dimension in dataset.coords:
return
species = []
megacomplexes = (dataset_model.global_megacomplex
if as_global else dataset_model.megacomplex)
for m in megacomplexes:
if isinstance(m, SpectralMegacomplex):
species += [
compartment for compartment in m.shape
if compartment not in species
]
dataset.coords[species_dimension] = species
matrix = dataset.global_matrix if as_global else dataset.matrix
clp_dim = "global_clp_label" if as_global else "clp_label"
dataset["species_spectra"] = (
(
dataset_model.get_model_dimension()
if not as_global else dataset_model.get_global_dimension(),
species_dimension,
),
matrix.sel({
clp_dim: species
}).values,
)
if not is_full_model:
dataset["species_associated_concentrations"] = (
(
dataset_model.get_global_dimension(),
species_dimension,
),
dataset.clp.sel(clp_label=species).data,
)
def simulate_clp(
dataset_model: DatasetModel,
parameters: ParameterGroup,
clp: xr.DataArray,
):
if "clp_label" not in clp.coords:
raise ValueError("Missing coordinate 'clp_label' in clp.")
global_dimension = next(dim for dim in clp.coords if dim != "clp_label")
global_axis = clp.coords[global_dimension]
matrices = ([
calculate_matrix(
dataset_model,
{global_dimension: index},
) for index, _ in enumerate(global_axis)
] if dataset_model.is_index_dependent() else calculate_matrix(
dataset_model, {}))
model_dimension = dataset_model.get_model_dimension()
model_axis = dataset_model.get_coordinates()[model_dimension]
result = xr.DataArray(
data=0.0,
coords=[
(model_dimension, model_axis.data),
(global_dimension, global_axis.data),
],
)
result = result.to_dataset(name="data")
for i in range(global_axis.size):
index_matrix = matrices[i] if dataset_model.is_index_dependent(
) else matrices
result.data[:, i] = np.dot(
index_matrix.matrix,
clp.isel({
global_dimension: i
}).sel({"clp_label": index_matrix.clp_labels}),
)
return result
def finalize_data(
self,
dataset_model: DatasetModel,
dataset: xr.Dataset,
is_full_model: bool = False,
as_global: bool = False,
):
if is_full_model:
return
megacomplexes = (
dataset_model.global_megacomplex if is_full_model else dataset_model.megacomplex
)
unique = len([m for m in megacomplexes if isinstance(m, DampedOscillationMegacomplex)]) < 2
prefix = "damped_oscillation" if unique else f"{self.label}_damped_oscillation"
dataset.coords[f"{prefix}"] = self.labels
dataset.coords[f"{prefix}_frequency"] = (prefix, self.frequencies)
dataset.coords[f"{prefix}_rate"] = (prefix, self.rates)
dim1 = dataset_model.get_global_axis().size
dim2 = len(self.labels)
doas = np.zeros((dim1, dim2), dtype=np.float64)
phase = np.zeros((dim1, dim2), dtype=np.float64)
for i, label in enumerate(self.labels):
sin = dataset.clp.sel(clp_label=f"{label}_sin")
cos = dataset.clp.sel(clp_label=f"{label}_cos")
doas[:, i] = np.sqrt(sin * sin + cos * cos)
phase[:, i] = np.unwrap(np.arctan2(sin, cos))
dataset[f"{prefix}_associated_spectra"] = (
(dataset_model.get_global_dimension(), prefix),
doas,
)
dataset[f"{prefix}_phase"] = (
(dataset_model.get_global_dimension(), prefix),
phase,
)
if self.index_dependent(dataset_model):
dataset[f"{prefix}_sin"] = (
(
dataset_model.get_global_dimension(),
dataset_model.get_model_dimension(),
prefix,
),
dataset.matrix.sel(clp_label=[f"{label}_sin" for label in self.labels]).values,
)
dataset[f"{prefix}_cos"] = (
(
dataset_model.get_global_dimension(),
dataset_model.get_model_dimension(),
prefix,
),
dataset.matrix.sel(clp_label=[f"{label}_cos" for label in self.labels]).values,
)
else:
dataset[f"{prefix}_sin"] = (
(dataset_model.get_model_dimension(), prefix),
dataset.matrix.sel(clp_label=[f"{label}_sin" for label in self.labels]).values,
)
dataset[f"{prefix}_cos"] = (
(dataset_model.get_model_dimension(), prefix),
dataset.matrix.sel(clp_label=[f"{label}_cos" for label in self.labels]).values,
)