lag

class Lag(OneToOneObservationRelatedTransformer):
    """Transform features by shifting your variable by any number of datapoints (timesteps e.g. weeks in a time serie)
       of your choice.

    Args:
        lag (float): Number of datapoints (timesteps e.g. weeks in a time serie) shifted. Common lags are positive and
                     no more than 2-3 weeks.
    """

    def __init__(self, lag: Union[int, float]) -> None:
        self.lag = lag

    @property
    def parameters(self) -> Dict[str, float]:
        return self.__dict__

    @property
    def is_linear(self) -> bool:
        return True

    # ------- METHODS -------

    def indexed_cumulative_sum(self, serie: np.ndarray) -> np.ndarray:
        return serie

    # ------- TRANSFORMERS -------

    @staticmethod
    def _transformer(serie: np.ndarray, lag: Union[int, float]) -> np.ndarray:  # type: ignore
        if serie.ndim == 2:
            # Apply the transformer to each row of the 2D array
            return np.column_stack([Lag._transformer_1d(serie[:, i], lag) for i in range(serie.shape[0])]).transpose()
        else:
            return Lag._transformer_1d(serie, lag)

    @staticmethod
    def _transformer_1d(serie: np.ndarray, lag: Union[int, float]) -> np.ndarray:
        serie = np.array(serie)

        if abs(lag) < FLOAT_TOLERANCE_THRESHOLD:
            return serie

        part_int, part_float = divmod(lag, 1)

        week_before = (1 - part_float) * shift(serie, int(part_int), cval=serie[0])
        week_after = part_float * shift(serie, int(part_int) + 1, cval=0)
        return week_before + week_after

    # ------- CHECKERS -------

    def check_params(self, serie: np.ndarray):
        """Check if parameters respect their application scope."""
        if self.lag >= serie.shape[0]:
            logger.warning(f"WARNING: lag is higher than data size. lag: '{self.lag}', data size: '{serie.shape[0]}'.")