dvc_session/src/feature_engineering.py

import os

import pandas as pd
import plac
from sklearn.decomposition import PCA

from src.utils import dump_yaml, read_data, read_yaml, save_results


@plac.annotations(
    data_path=("Path to source data", "option", "i", str),
    feature_path=("Path to save featurized data", "option", "f", str),
    out_path=("Path to save pca model", "option", "o", str),
)
def main(
    data_path: str = "data/split/",
    feature_path: str = "data/features/",
    out_path: str = "models/pca/",
) -> None:
    X_train, X_test, y_train, y_test = read_data(data_path)

    params = read_yaml("params.yaml", "pca")
    pca = PCA(**params).fit(X_train)

    train_feature = pd.DataFrame(pca.transform(X_train))
    test_feature = pd.DataFrame(pca.transform(X_test))
    train_feature["class"] = y_train
    test_feature["class"] = y_test

    if not os.path.isdir(feature_path):
        os.mkdir(feature_path)
    train_feature.to_csv(f"{feature_path}train.csv", index=False)
    test_feature.to_csv(f"{feature_path}test.csv", index=False)
    save_results(out_path, pca, None)

    print("Finished Feature Engineering:\nStats:")
    print(f"\tExplained Variance: {pca.explained_variance_}")
    print(f"\tExplained Variance Ratio: {pca.explained_variance_ratio_}")

    dump_yaml(
        dict(
            explained_variance=pca.explained_variance_.tolist(),
            explained_variance_ratio=pca.explained_variance_ratio_.tolist(),
        ),
        "models/pca/metrics.yaml",
    )


if __name__ == "__main__":
    plac.call(main)