Source code for pycif.plugins.modes.response_functions.base_function.base_function

import datetime
import os
import shutil
from logging import debug
from typing import Any, Union

import numpy as np
import pandas as pd

from .....utils.iterators import iter_attributes
from .....utils.path import link
from .....utils.check.errclass import CifFileNotFoundError, CifTypeError, CifValueError

# Aliases for type hinting
Mode = Any
DatetimeLike = Union[datetime.datetime, pd.Timestamp, np.datetime64]




def to_datetime(dt: DatetimeLike) -> datetime.datetime:
    """A function to convert various datetime-like objects to standard Python datetime objects.
    
    Parameters:
        dt (DatetimeLike): The input datetime-like object to be converted.
        
    Returns:
        datetime.datetime: The standard Python datetime object after conversion.
        
    Raises:
        ValueError: If the input datetime-like object cannot be converted to a standard datetime object.
        TypeError: If the input type is unexpected and cannot be converted to a datetime object.
    """
    if isinstance(dt, datetime.datetime):
        return dt

    elif isinstance(dt, pd.Timestamp):
        # Converting from pandas Timestamp to python built-in datetime
        if dt is pd.NaT:
            raise CifValueError("cannot convert 'NaT' to datetime")
        return dt.to_pydatetime()

    elif isinstance(dt, np.datetime64):
        # Converting from numpy datetime64 to python built-in datetime
        if np.isnan(dt):
            raise CifValueError("cannot convert 'NaT' to datetime")
        return dt.astype('datetime64[s]').tolist()

    else:
        raise CifTypeError("unexpected type, can not convert to datetime")





class BaseFunction:
    def __init__(
        self,
        mode: Mode,
        index: int,
        component: str,
        parameter: str,
        date_start: DatetimeLike,
        date_end: DatetimeLike,
    ) -> None:
        """A class representing a base function

        Args:
            mode (Mode): the mode plugin
            index (int): controlvect index
            component (str): component name
            parameter (str): parameter name
            date_start (datetime): start of the base function time window
            date_end (datetime): end of the base function time window
        """
        if index > 999999:
            raise CifValueError("too many response functions")

        if date_start == date_end:
            raise CifValueError(f"response function {index} has identical start "
                             f"and end dates: {date_start}")

        # References to some of the mode plugin's attributes
        self.basedir = mode.basedir
        self.datei = mode.datei
        self.controlvect = mode.controlvect
        self.datavect = mode.datavect
        self.inicond = mode.inicond

        # References to some of the mode plugin's input arguments
        self.run_mode = mode.run_mode
        self.inicond_component = mode.inicond_component
        self.reload_results = mode.reload_results
        self.ignore_tracers = mode.ignore_tracers

        self.index = index
        self.component = component
        self.parameter = parameter
        self.date_start = to_datetime(date_start)
        self.date_end = to_datetime(date_end)

        self.name = f"base_function_{index:06d}"
        self.controlvect_filename = "controlvect.pickle"

    def __repr__(self) -> str:
        s = f"<BaseFunction index={self.index}, component={self.component}, " \
            f"parameter={self.parameter}>"
        return s

    def __str__(self) -> str:
        s = f"base function {self.index:06d}: " \
            f"{self.component}/{self.parameter}, " \
            f"{self.date_start.isoformat()} -> {self.date_end.isoformat()}"
        return s

    @property
    def rundir(self) -> str:
        return os.path.join(self.basedir, self.name)

    @property
    def outdir(self) -> str:
        return os.path.join(self.rundir, f"obsoperator/{self.run_mode}_0000/obsvect")

    @property
    def obsdir(self) -> str:
        return os.path.join(self.basedir, f"base_function_obsvect_{self.index:06d}")

    @property
    def controlvect_path(self) -> str:
        return os.path.join(self.rundir, self.controlvect_filename)

    @property
    def yaml_config_path(self) -> str:
        return os.path.join(self.rundir, f"config_{self.name}.yaml")



    def is_ignored(self) -> bool:
        """True if the base fonction (component, parameter) couple is in the
        'ignore_tracer' input argument
        """
        return (self.component, self.parameter) in self.ignore_tracers




    def has_run_succesfully(self) -> bool:
        """True is the base function has run successfully (output directory is
        present)"""
        return self.reload_results and os.path.isdir(self.outdir)




    def update_inicond_date(self) -> None:
        """Updates the initial conditions tracers date if those are in the 
        control vector and the response function start date is different from
        the simulation window initial date
        """
        if self.datei == self.date_start:
            return

        for tracer_name, tracer in iter_attributes(self.inicond.parameters):
            if not tracer.iscontrol:
                continue

            if tracer.ndates != 1:
                raise CifValueError(
                    f"initial conditions component '{self.inicond_component}' "
                    "has multiple dates, response functions excepted only one"
                )

            debug(
                f"updating initial conditions tracer {self.inicond_component}/"
                f"{tracer_name} date from {tracer.dates[0].isoformat()} to "
                f"{self.date_start.isoformat()}"
            )
            tracer.dates = np.array([self.date_start])




    def dump_controlvect(self) -> None:
        """Updates the controlvector and dumps it
        """
        self.update_inicond_date()

        if self.run_mode == "fwd":
            # Setting control vector x
            self.controlvect.x = np.zeros(self.controlvect.dim)
            self.controlvect.x[self.index] = self.controlvect.xb[self.index]

        elif self.run_mode == "tl":
            # Setting control vector dx
            self.controlvect.dx = np.zeros(self.controlvect.dim)
            self.controlvect.dx[self.index] = 1.0

        os.makedirs(self.rundir, exist_ok=True)
        self.controlvect.dump(self.controlvect_path)




    def fetch_obsvect(self) -> None:
        """Fetch the response function outputs (observation vector) and
        move it to its 'obsvect' directory
        """
        if not os.path.isdir(self.obsdir):
            if not os.path.isdir(self.outdir):
                raise CifFileNotFoundError(
                    f"Response function {self.index:06d} did not produce the "
                    f"requiered outputs '{self.outdir}'"
                )

            link(self.outdir, self.obsdir)




    def flush(self) -> None:
        """Delete some temporary files to free some disk space
        """
        # Removing the control vector dump file
        if os.path.isfile(self.controlvect_path):
            debug(f"deleting control vector dump '{self.controlvect_path}'")
            os.remove(self.controlvect_path)

        # Removing the 'chain' directory
        chain_dir = os.path.join(
            self.rundir, f"obsoperator/{self.run_mode}_0000/chain")

        # Some checks to avoid deleting something else
        if os.path.isdir(chain_dir) and os.path.isabs(chain_dir):
            debug(f"deleting directory '{chain_dir}'")
            shutil.rmtree(chain_dir)