from __future__ import annotations
from datetime import datetime
from os import PathLike
from pathlib import Path
from typing import Literal
import numpy as np
import pandas as pd
import xarray as xr
from pandas import DataFrame, MultiIndex, Series
from xarray import Dataset
from ......utils.check.errclass import (
CifFileNotFoundError,
CifRuntimeError,
CifValueError,
)
from ......utils.hdf5 import _hdf5_lock
_attributes = {
"time": {
"long_name": "time",
"calendar": "proleptic_gregorian",
},
"itrac": {
"standard_name": "tracer_index",
"long_name": "tracer index",
},
"ilev": {
"standard_name": "model_level_index",
"long_name": "model level index",
},
"ilon": {
"standard_name": "model_longitude_index",
"long_name": "model longitude index",
},
"ilat": {
"standard_name": "model_latitude_index",
"long_name": "model latitude index",
},
"icell": {
"standard_name": "model_cell_index",
"long_name": "model cell index",
},
"obs": {
"standard_name": "observation",
"long_name": "observation (unused, for reference only)",
"units": "ppm",
},
"obs_ad": {
"standard_name": "adjoint_observation",
"long_name": "adjoint observation",
"units": "ppm",
},
}
[docs]
def to_dataframe(
self,
dict_datastores: dict[tuple[str, str], DataFrame],
input_type: str,
mode: Literal["fwd", "adj"],
) -> DataFrame | None:
df_list = []
for datastore in dict_datastores.values():
parameter = datastore[("metadata", "parameter")]
tracer_name = parameter.str.replace("__sample#", "_").str.lower()
# Include only part of the datastore with datastore in active species
active_species = [spec.lower() for spec in self.chemistry.active_species]
mask = tracer_name.isin(active_species)
maindata = datastore["maindata"].loc[mask, :]
metadata = datastore["metadata"].loc[mask, :]
if np.any(metadata["dtstep"].values != 1):
raise CifValueError("'dstep' values different from 1 is not supported")
# Time coordinate
# WARNING: This part is critical, modify with caution
time = (metadata["tstep"] * self.dt.total_seconds()).astype("int32")
# Tracer index
tracer_name = tracer_name.values
tracer_index = np.zeros(tracer_name.shape, dtype=np.int32)
for i, spec in enumerate(active_species):
tracer_index[tracer_name == spec] = i
# Level index
# Assumes that stations with no level are in first level
# TODO: make it general
level = np.asarray(metadata["level"].values)
level[~pd.notnull(level)] = 0
level = level.astype("int32")
df = pd.DataFrame(
{
"time": time,
"itrac": tracer_index + 1,
"ilev": level + 1,
},
)
if self.grid == "regular":
df["ilon"] = metadata["j"].values.astype("int32") + 1
df["ilat"] = metadata["i"].values.astype("int32") + 1
elif self.grid == "dynamico":
df["icell"] = metadata["i"].values.astype("int32") + 1
else:
raise CifValueError(f"Unknown grid type: '{self.grid}'")
if input_type == "concs":
if "obs" in maindata:
df["obs"] = maindata["obs"].values.astype("float32")
if mode == "adj":
df["obs_ad"] = maindata["adj_out"].values
df_list.append(df)
if not df_list:
return
df = pd.concat(df_list, ignore_index=True)
return df
[docs]
def get_indices(df: DataFrame, grid: Literal["regular", "dynamico"]) -> MultiIndex:
h_vars = ["ilon", "ilat"] if grid == "regular" else ["icell"]
cols = ["itrac", "time", "ilev"] + h_vars
indices = pd.MultiIndex.from_frame(df[cols])
return indices
[docs]
def agg_duplicates(
df: DataFrame, grid: Literal["regular", "dynamico"]
) -> tuple[np.ndarray, DataFrame]:
indices = get_indices(df, grid) # type: ignore
if not indices.has_duplicates:
return df.index.to_numpy(), df
inverse, _ = pd.factorize(indices, sort=False)
# Remove duplicates (keep first)
first = ~pd.Index(inverse).duplicated(keep="first")
df_unique = df.iloc[first]
df_unique.reset_index(drop=True, inplace=True)
# Sum adjoint values
if "obs_ad" in df.columns:
df_unique.loc[:, "obs_ad"] = df["obs_ad"].groupby(inverse).sum().values
return inverse, df_unique
[docs]
def write_obs(df: DataFrame, path: str | PathLike[str], datei: datetime) -> None:
df = df.sort_values("time")
ds = xr.Dataset.from_dataframe(df)
for varname, attrs in _attributes.items():
if varname in ds:
ds[varname].attrs = attrs
for varname in ("itrac", "ilev", "ilat", "ilon", "icell"):
if varname in ds:
ds[varname].attrs["comment"] = "starts from 1"
# WARNING: This part is critical, modify with caution
ds["time"].attrs["units"] = f"seconds since {datei:%Y-%m-%d %H}:00:00"
ds.to_netcdf(path)
[docs]
def make_obs(
self,
dict_datastores: dict[tuple[str, str], DataFrame],
ddi: datetime,
input_type: str,
runsubdir: str | PathLike,
mode: Literal["fwd", "adj"],
do_simu: bool = True,
) -> None:
"""Write observation input file for one sub-period"""
# If empty datastore, do nothing
if all(isinstance(datastore, dict) for datastore in dict_datastores.values()):
return
if all(datastore.size == 0 for datastore in dict_datastores.values()):
return
active_species = [spec.lower() for spec in self.chemistry.active_species]
obs_file = Path(runsubdir, "obs.nc")
if self.init_obs[ddi] and not obs_file.exists():
raise CifFileNotFoundError(f"'{obs_file.name}' does not exist")
if self.reset_obs[ddi]:
self.reset_obs[ddi] = False
# Reset indexes cache
self.chunk_indexes[ddi] = {
comp: {spec: None for spec in self.chemistry.active_species}
for comp in self.output_components
}
# Remove obs.nc if it already exists (most likely from a previous run)
obs_file.unlink(missing_ok=True)
# Prepare observations dataset
df = to_dataframe(self, dict_datastores, input_type, mode)
if df is None:
return
if input_type == "concs":
# Get already existing observations
if self.init_obs[ddi]:
with _hdf5_lock:
with xr.open_dataset(obs_file, decode_times=False) as ds_ref:
# Restore original order (previously sorted by time)
df_ref = ds_ref.to_dataframe().sort_index()
nobs_ref = len(df_ref)
# Concatenate datasets
df_all = pd.concat([df_ref, df], ignore_index=True)
else:
df_all = df
nobs_ref = 0
# Aggregate duplicates
indices, df_unique = agg_duplicates(df_all, self.grid)
indices = indices[nobs_ref:]
# Write new file
write_obs(df_unique, obs_file, self.subsimu_intervals[ddi][0])
if not self.skip_obs_chunking:
# For each tracer save the corresponding indices
for _, tracer in dict_datastores:
tracer = tracer.replace("__sample#", "_")
itrac = active_species.index(tracer.lower()) + 1
self.chunk_indexes[ddi][input_type][tracer] = indices[df.itrac == itrac]
# Observation file has been initialized
self.init_obs[ddi] = True
elif self.init_obs[ddi] and not self.skip_obs_chunking:
with _hdf5_lock:
with xr.open_dataset(obs_file, decode_times=False) as ds_ref:
# Restore original order (previously sorted by time)
df_ref = ds_ref.to_dataframe().sort_index()
indices_ref = get_indices(df_ref, self.grid) # Must have no duplicate
indices = get_indices(df, self.grid) # Can have duplicates
indexer = indices_ref.get_indexer(indices)
if np.any(indexer == -1):
raise CifValueError(
"Observations indices not found in 'obs.nc' for "
+ f"{input_type} ({ddi:%Y-%m-%d %H:%M:%S})"
)
# For each tracer save get the indices in the observation file corresponding to
# the observation locations
for _, tracer in dict_datastores:
tracer = tracer.replace("__sample#", "_")
itrac = active_species.index(tracer.lower()) + 1
self.chunk_indexes[ddi][input_type][tracer] = indexer[df.itrac == itrac]
else:
raise CifRuntimeError(
f"attempting to set observations locations for '{input_type}' any 'concs'"
)
