from __future__ import annotations
import datetime
import os
from os import PathLike
from typing import Any, Literal, TypeVar
import numpy as np
import pandas as pd
import xarray as xr
from .....utils.check.errclass import CifFileNotFoundError, CifNotImplementedError, CifValueError
# Aliases for type hinting
Model = Any
DataType = TypeVar("DataType", xr.DataArray, xr.Dataset)
[docs]
def strip_spatial_coords(data: DataType) -> DataType:
"""Drop 'lev', 'lat' and 'lon' coordinates from a Dataset or DataArray to
ensure taht futher operation between DataArrays will be index based (like
numpy operation) and not coordinate based.
Args:
data (xr.DataArray or xr.Dataset)
Returns:
xr.DataArray or xr.Dataset: without spatial coordinates
"""
return data.drop_vars(['lev', 'lat', 'lon'], errors="ignore")
[docs]
def read_inicond_file(
self: Model,
runsubdir: str,
mode: Literal["fwd", "tl"],
) -> xr.DataArray:
"""Read LMDZ initial conditions file (.nc or .bin) for a given mode (fwd or tl)
Args:
self (Model): LMDZ model plugin
runsubdir (str): sub simulation run directory
mode (str): 'fwd' (forward) or 'tl' (tangent-linear)
"""
if mode == "fwd":
file_basename = "start"
elif mode == "tl":
file_basename = "start_tl"
else:
raise CifValueError(f"unexpected mode '{mode}', should be 'fwd' or 'tl'")
# Dimensions
nspec = len(self.chemistry.acspecies.attributes)
nlev = self.domain.nlev
nlat = self.domain.nlat
nlon = self.domain.nlon
nc_file = os.path.join(runsubdir, f"{file_basename}.nc")
if os.path.isfile(nc_file):
inicond = xr.DataArray(
data=np.zeros((nspec, nlev, nlat, nlon)), dims=['spec', 'lev', 'lat', 'lon']
)
# Reading initial conditions from NetCDF file
with xr.open_dataset(nc_file) as ds:
# Looping over active species
for index, spec in enumerate(self.chemistry.acspecies.attributes):
spec_plg = getattr(self.chemistry.acspecies, spec)
var_name = f"q{spec_plg.restart_id:02d}"
inicond[index, ...] = ds[var_name].values[0, ...]
else:
# Reading initial conditions from Fortran binary file
bin_file = os.path.join(runsubdir, f"{file_basename}.bin")
if not os.path.isfile(bin_file):
raise CifFileNotFoundError(
"could not find initial condition file " f"'{nc_file}' or '{bin_file}'"
)
# Version info
lmdz_version = self.plugin.version
if lmdz_version == "std":
data = np.fromfile(bin_file, offset=4)
elif lmdz_version == "acc":
data = np.fromfile(bin_file)
else:
raise CifValueError(f"unexpected LMDZ version '{lmdz_version}'")
data = data.reshape((nlon, nlat, nlev, nspec), order="F").T
inicond = xr.DataArray(data=data, dims=['spec', 'lev', 'lat', 'lon'])
return inicond
[docs]
def chemfield_time_coord(ddi: datetime.datetime) -> xr.DataArray:
"""Build the daily time coordinate for a chemistry-field input dataset.
Returns a 1-D DataArray of daily timestamps spanning the full calendar
month that contains *ddi*.
Args:
ddi: any date within the target month.
Returns:
xr.DataArray: daily date range of length ``days_in_month``.
"""
di = pd.to_datetime(ddi)
di = di if di.is_month_start else di - pd.offsets.MonthBegin()
ntime = di.days_in_month
time = xr.DataArray(data=pd.date_range(di, periods=ntime, freq='1D'), dims='time')
return time
[docs]
def read_kinetic(
self: Model,
ddi: datetime.datetime,
runsubdir: str | PathLike,
) -> xr.Dataset:
"""Read the LMDZ kinetic (pressure/temperature) field from the run directory.
Opens ``kinetic.nc``, renames dimensions to CIF conventions, trims to
the month's day count, and strips spatial coordinate metadata.
Args:
self: LMDZ model plugin instance.
ddi: period start date (determines the target month).
runsubdir: path to the period run directory.
Returns:
xr.Dataset with ``pmid`` and ``temp`` variables on
``(time, lev, lat, lon)`` coordinates.
"""
# Version info
lmdz_version = self.plugin.version
# Coords
time = chemfield_time_coord(ddi)
(ntime,) = time.shape
kinetic_file = os.path.join(runsubdir, "kinetic.nc")
with xr.open_dataset(kinetic_file) as ds:
ds = ds[['pmid', 'temp']].load()
if lmdz_version == "acc":
# Adding looping longitude
ds = xr.concat([ds, ds.isel(lon=[0])], dim='lon')
ds = ds.rename({'time_counter': 'time', 'presnivs': 'lev'})
ds = ds.isel(time=slice(None, ntime))
ds = strip_spatial_coords(ds)
ds['time'] = time
return ds
[docs]
def read_prescr(
self: Model,
ddi: datetime.datetime,
runsubdir: str | PathLike,
) -> xr.DataArray:
"""Read prescribed-concentration fields from the LMDZ run directory.
Reads concentration data for all prescribed species from ``prescr.nc``
in *runsubdir*, aligns the time axis to the month containing *ddi*, and
strips spatial coordinate metadata.
Args:
self: LMDZ model plugin instance (carries
``chemistry.prescrconcs.attributes``).
ddi: period start date (determines the target month).
runsubdir: path to the period run directory.
Returns:
xr.DataArray of prescribed concentrations on
``(time, lev, lat, lon)`` coordinates.
"""
# Version info
lmdz_version = self.plugin.version
# Coords
time = chemfield_time_coord(ddi)
# Dimensions
nspec = len(self.chemistry.prescrconcs.attributes)
(ntime,) = time.shape
nlev = self.domain.nlev
nlat = self.domain.nlat
nlon = self.domain.nlon
prescr = xr.DataArray(
data=np.zeros((nspec, ntime, nlev, nlat, nlon)),
dims=['spec', 'time', 'lev', 'lat', 'lon'],
coords={'time': time},
)
# Looping over prescribed species
for index, spec in enumerate(self.chemistry.prescrconcs.attributes):
scale_file = os.path.join(runsubdir, f"mod_scale_{spec}.bin")
if os.path.isfile(scale_file):
raise CifNotImplementedError("scaling files are not implemented")
prescr_file = os.path.join(runsubdir, f"prescr_{spec}.nc")
with xr.open_dataset(prescr_file) as ds:
data = ds[spec].values[:ntime, ...]
if lmdz_version == "acc":
data = np.concatenate([data, data[..., [0]]], axis=3)
prescr[index, ...] = data
return prescr
