drift/transformations/pca.py

from __future__ import annotations
from transformations.base import Transformation
from typing import Optional
from copy import deepcopy
from sklearn.decomposition import PCA
import pandas as pd


class PCATransformation(Transformation):

    pca: PCA

    def __init__(self, ratio_components_to_keep: float, sliding_window_size: int):
        self.ratio_components_to_keep = ratio_components_to_keep
        self.sliding_window_size = sliding_window_size

    def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None) -> None:
        self.pca = PCA(
            n_components=min(
                int(len(X.columns) * self.ratio_components_to_keep),
                self.sliding_window_size,
            )
        )
        self.pca.fit(X, y)

    def fit_transform(
        self, X: pd.DataFrame, y: Optional[pd.Series] = None
    ) -> pd.DataFrame:
        self.fit(X, y)
        return self.transform(X)

    def transform(self, X: pd.DataFrame) -> pd.DataFrame:
        X = pd.DataFrame(self.pca.transform(X), index=X.index)
        X.columns = ["PCA_" + str(i) for i in range(1, len(X.columns) + 1)]
        return X

    def clone(self) -> PCATransformation:
        return deepcopy(self)

    def get_name(self) -> str:
        return "PCA"