from __future__ import annotations
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 ......utils.check.errclass import CifValueError
from ......utils.hdf5 import _hdf5_lock
[docs]
def make_obs(
self,
dict_datastores: dict[tuple[str, str], pd.DataFrame],
ddi: datetime.datetime,
input_type: str,
runsubdir: str | PathLike,
mode: Literal["fwd", "adj"],
) -> 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
if self.reset_obs[ddi]:
self.reset_obs[ddi] = False
self.chunk_indexes[ddi] = {
comp: {spec: None for spec in self.chemistry.active_species}
for comp in self.output_components
}
def get_tracer_names(parameter: pd.Series) -> pd.Series:
return 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]
# Initialize datastore to dump
datei, _ = self.subsimu_intervals[ddi]
ds = []
# Loop over tracers to concatenate datastores to dump
for datastore in dict_datastores.values():
tracer_name = get_tracer_names(datastore["metadata"]["parameter"])
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")
time_units = f"seconds since {datei:%Y-%m-%d %H:00:00}"
# 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")
ds_tracer = xr.Dataset(
# fmt: off
{
"time": (["index"], time, {
"standard_name": "time",
"long_name": "time",
"units": time_units,
"calendar": "proleptic_gregorian",
}),
"obs": (["index"], maindata["obs"].values.astype("float32"), {
"standard_name": "observation",
"long_name": "observation (unused, for reference only)",
"units": "ppm",
}),
"itrac": (["index"], tracer_index + 1 , {
"standard_name": "tracer_index",
"long_name": "tracer index",
}),
"ilev" : (["index"], level + 1, {
"standard_name": "model_level_index",
"long_name": "model level index",
}),
}
# fmt: on
)
if mode == "adj":
# fmt: off
ds_tracer["obs_ad"] = (["index"], maindata["adj_out"].values, {
"standard_name": "adjoint_observation",
"long_name": "adjoint observation",
"units": "ppm",
})
# fmt: on
if self.grid == "regular":
# NOTE: In CIF latitude index is "i" and longitude index is "j" (row-major)
# In DISPERSION latitude index is "j" and longitude index is "i" (column-major)
# fmt: off
ds_tracer["ilon"] = (["index"], metadata["j"].values.astype("int32") + 1, {
"standard_name": "model_longitude_index",
"long_name": "model longitude index (starts from 1)",
})
ds_tracer["ilat"] = (["index"], metadata["i"].values.astype("int32") + 1, {
"standard_name": "model_latitude_index",
"long_name": "model latitude index (starts from 1)",
})
# fmt: on
elif self.grid == "dynamico":
# fmt: off
ds_tracer["icell"] = (["index"], metadata["i"].values.astype("int32") + 1, {
"standard_name": "model_cell_index",
"long_name": "model cell index (starts from 1)",
})
# fmt: on
else:
raise CifValueError(f"Unknown grid type: '{self.grid}'")
ds.append(ds_tracer)
# Stop here if ther is no observation to include
if ds == []:
return
# Concatenate trids
ds = xr.concat(ds, dim="index")
# Concat observations to obs.nc if it already exists
obs_file = Path(runsubdir, "obs.nc")
if self.iniobs[ddi] and obs_file.exists():
with _hdf5_lock:
with xr.open_dataset(obs_file, decode_times=False) as ds_orig:
nobs_orig = ds_orig.sizes["index"]
ds = xr.concat([ds_orig, ds], dim="index")
else:
nobs_orig = 0
# Get spatial and temporal location of observations
h_vars = ["ilon", "ilat"] if self.grid == "regular" else ["icell"]
cols = ["time", "itrac", "ilev"] + h_vars
df = ds[cols].to_dataframe()
# Get unique obs location to avoid double extraction of parameters
groupby_duplicates = df.groupby(cols)
group_ids = groupby_duplicates.ngroup()
unique_index = group_ids.drop_duplicates()
# Aggregate adj_out
if mode == "adj":
obs_ad = ds["obs_ad"].to_dataframe().groupby(group_ids).sum()
obs_ad = obs_ad.values[unique_index]
# Truncate to current obs batch
group_ids = group_ids.iloc[nobs_orig:]
# Get index of observations to extract (unique locations)
indices = pd.Series(range(len(unique_index)), index=unique_index)
indices = indices.loc[group_ids].values
# Save index per tracer
itrac = df.drop_duplicates().iloc[indices]["itrac"] - 1
for i, spec in enumerate(self.chemistry.active_species):
if (itrac == i).any():
self.chunk_indexes[ddi][input_type][spec] = indices[itrac == i]
# Update NetCDF file
ds = ds.sel(index=unique_index.index)
if mode == "adj":
ds["obs_ad"][:] = obs_ad.flatten()
with _hdf5_lock:
ds.to_netcdf(obs_file)
# obs.nc already exists and should be updated if another tracer is added
self.iniobs[ddi] = True
