SPATZ/spatz/logger.py

import numpy as np
import pandas as pd

from typing import Any, Tuple
from numpy.typing import ArrayLike

from abc import abstractmethod


class Advanceable:
    def __init__(self) -> None:
        self.reset()

    def step(self, dt: float):
        """Advances the simulation data in time.

        Args:
            dt (float): The step in time to make.
        """
        self.__t += dt
        self._on_step(dt)

    def reset(self):
        """
        Reset the Avanceable object to its initial state.
        """
        self.__t = 0
        self._on_reset()

    @abstractmethod
    def _on_step(self, dt: float):
        pass

    @abstractmethod
    def _on_reset(self):
        pass

    def get_time(self) -> float:
        """
        Returns:
            float: Returns the current time of the Advanceable.
        """
        return self.__t


class Logger(Advanceable):
    def __init__(self) -> None:
        super().__init__()
        self.__idx = -1

    def _on_step(self, _: float):
        self.__df = pd.concat([self.__df, pd.DataFrame()], ignore_index=True)
        self.__idx += 1
        self.__df.loc[self.__idx, 'time'] = self.get_time()

    def _on_reset(self):
        self.__df = pd.DataFrame.from_dict({'time': [self.get_time()]}).astype(np.float64)

    def write(self, attrib: str, value: Any, domain: str = 'all'):
        """Writes a value to the logger.

        Args:
            attrib (str): The name of the value to log.
            value (Any): The value to log.
            domain (str, optional): The domain the value belongs to. Defaults to 'any'.
        """
        name = domain + '/' + attrib

        if name not in self.__df.columns:
            self.__df[name] = pd.Series([pd.NA] * len(self.__df))

        self.__df.loc[self.__idx, name] = value

    def get_dataframe(self) -> pd.DataFrame:
        return self.__df