SPATZ/spatz/observer/observer.py

from typing import Any, List, Dict, AnyStr
from numpy.typing import ArrayLike

from spatz.dataset import Dataset
from spatz.logger import Logger
from spatz.transforms import Transform


class Observer:
    def __init__(self, dataset: Dataset, logger: Logger, attributes: List[str]):
        self._dataset = dataset
        self._logger = logger
        self.__attrs = attributes

    def _get_name(self) -> AnyStr:
        return 'general'

    def set_dataset(self, dataset: Dataset):
        self._dataset = dataset

    def set_logger(self, logger: Logger):
        self._logger = logger
    
    def _log(self, name: AnyStr, value: Any):
        self._logger.write(name, value, self._get_name())

    def get_start_value(self) -> ArrayLike:
        """
        Returns:
            ArrayLike: Returns the values of the observed attributes at the start of the simulation.
        """
        return self(t=self._dataset.get_start_time())

    def __call__(self, t: float | None = None) -> ArrayLike:
        data = self._dataset.fetch_values(self.__attrs, t)

        for attrib, value in zip(self.__attrs, data):
            self._log(attrib, value)

        return data