import datetime
import os
from typing import Any, Literal, TypeVar
import xarray as xr
import pandas as pd
import numpy as np
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
) -> xr.Dataset:
"""Read the LMDZ-old kinetic (pressure/temperature) field from the run directory.
Opens ``kinetic.nc``, renames dimensions to CIF conventions, and trims
to the month's day count.
Args:
self: LMDZ-old model plugin instance.
ddi: period start date.
runsubdir: path to the period run directory.
Returns:
xr.Dataset with ``pmid`` and ``temp`` variables.
"""
# 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
) -> xr.DataArray:
"""Read prescribed-concentration fields from the LMDZ-old run directory.
Reads ``prescr.nc`` for all prescribed species and aligns time to the
month containing *ddi*.
Args:
self: LMDZ-old model plugin instance.
ddi: period start date.
runsubdir: path to the period run directory.
Returns:
xr.DataArray of prescribed concentrations.
"""
# 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
