from __future__ import annotations
from datetime import datetime
from logging import debug
import numpy as np
import pandas as pd
from ....utils.dates import date_range
[docs]
def is_month_start(datetime: pd.Timestamp) -> bool:
"""Return ``True`` if *datetime* falls exactly at the start of a calendar month (midnight)."""
return (
datetime.is_month_start
and datetime.hour == 0
and datetime.minute == 0
and datetime.second == 0
)
[docs]
def ini_periods(self, **kwargs) -> None:
"""Compute temporal discretisation for the LMDZ-ico model.
Splits the full window into sub-periods capped at one calendar month
each, then builds per-period date arrays for time steps, inputs, fluxes,
chemistry fields, and meteo mass fluxes.
Sets on *self*:
* ``subsimu_dates`` — period boundary datetimes.
* ``subsimu_intervals`` — dict mapping period start to (start, end) tuple.
* ``tstep_dates`` — per-period full-month time-step arrays.
* ``input_dates`` — per-period daily initial/end-concentration dates.
* ``flux_input_dates`` — per-period flux input dates.
* ``chem_input_dates`` — per-period chemistry field dates.
* ``meteo_input_dates`` — per-period mass-flux dates.
* ``iniobs``, ``reset_obs`` — per-period obs bookkeeping flags.
* ``chunk_indexes`` — per-period per-component per-species index cache.
* ``runsimu`` — per-period flag controlling whether to run or use cached output.
* ``chain`` — per-period flag for adjoint restart chaining.
Args:
self: LMDZ-ico plugin instance with ``datei``, ``datef``,
``period_freq``, ``dt``, ``flux_freq``, ``chem_freq``,
``mass_fluxes_freq``, and ``chemistry.active_species`` set.
**kwargs: unused.
"""
# Keep old behaviour for month sub-simulations (maybe not needed)
if self.period_freq in ("MS", "1MS"):
ref_datei = pd.Timestamp(year=self.datei.year, month=self.datei.month, day=1)
else:
ref_datei = self.datei
dates = pd.date_range(ref_datei, self.datef, freq=self.period_freq)
subsimu_dates = []
# Simulation sub-periods, cut to a maximum length of 1 month
for di, df in zip(dates[:-1], dates[1:]):
subsimu_dates.append(di)
# Both datetimes are in the same month
if di.year == df.year and di.month == df.month:
continue
next_month = (di + pd.offsets.MonthBegin()).month
# df is the first day of the month after di
if df.month == next_month and is_month_start(df):
continue
# There is at least one month between di and df
in_between = pd.date_range(
di,
df,
freq="MS",
normalize=True,
inclusive="neither",
)
subsimu_dates.extend(in_between.to_list())
# List of sub-simulation windows
subsimu_dates.append(dates[-1])
self.subsimu_dates = pd.to_datetime(subsimu_dates).to_pydatetime() # type: ignore # pylint: disable=no-member
if self.subsimu_dates[0] < self.datei:
self.subsimu_dates[0] = self.datei
if self.subsimu_dates[-1] < self.datef:
self.subsimu_dates = np.append(self.subsimu_dates, self.datef)
self.subsimu_intervals = {
ddi: (ddi, ddf)
for ddi, ddf in zip(self.subsimu_dates[:-1], self.subsimu_dates[1:])
}
debug_msg = "LMDZ sub-simulation windows:\n"
for dk, (di, df) in self.subsimu_intervals.items():
debug_msg += f"- {dk}: {di} - {df}\n"
debug(debug_msg[:-1])
def get_input_dates(
freq: str | pd.Timedelta,
full_month: bool = False,
) -> dict[datetime, np.ndarray]:
"""Build per-period date arrays at *freq* resolution.
Args:
freq: date frequency (pandas-compatible string or Timedelta).
full_month: if ``True``, span the full calendar month of each
period start rather than just the period window.
Returns:
dict mapping period-start date to its date array.
"""
dates = {}
for ddi, ddf in self.subsimu_intervals.values():
if full_month:
ddi_month = pd.Timestamp(year=ddi.year, month=ddi.month, day=1)
ddf_month = ddi_month + pd.DateOffset(months=1)
dates[ddi] = date_range(ddi_month, ddf_month, period=freq) # type: ignore
else:
dates[ddi] = date_range(ddi, ddf, period=freq) # type: ignore
return dates
# List of time steps
# WARNING: This part is critical for observations, modify with caution
self.tstep_dates = get_input_dates(self.dt, full_month=True)
# TODO: make the date inputs flexible (daily so far)
self.input_dates = get_input_dates("1D", full_month=False)
# Surface fluxes input dates
self.flux_input_dates = get_input_dates(self.flux_freq, full_month=False)
# Chemical fields input dates (full months)
self.chem_input_dates = get_input_dates(self.chem_freq, full_month=True)
# Mass fluxes input dates (full months at 3-hourly resolution)
# Need to set input dates for meteo plugin here
self.meteo_input_dates = get_input_dates(self.mass_fluxes_freq, full_month=True)
# Initializes dictionary to keep in memory whether observations were
# already dumped for a given period
self.iniobs = {ddi: False for ddi in self.subsimu_dates}
self.reset_obs = {ddi: True for ddi in self.subsimu_dates}
# Initializes dictionary to keep in memory observations index information
# Only "concs" component is used there
self.chunk_indexes = {
ddi: {
comp: {spec: None for spec in self.chemistry.active_species}
for comp in self.output_components
}
for ddi in self.subsimu_dates
}
# Run model or approximation with fake_end
self.runsimu = {ddi: True for ddi in self.subsimu_dates}
# Keep in memory whether a given period has a successor period
# The info is used by the adjoint to fetch or not the restart file
self.chain = {ddi: True for ddi in self.subsimu_dates[:-1]}
self.chain[self.subsimu_dates[-2]] = False
