Source code for pycif.plugins.models.lmdz_ico.ini_periods

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




def is_month_start(datetime: pd.Timestamp) -> bool:
    return (
        datetime.is_month_start
        and datetime.hour == 0
        and datetime.minute == 0
        and datetime.second == 0
    )





def ini_periods(self, **kwargs) -> None:
    # 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]:
        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