End-to-End Training Pipeline

import os
import gc


import numpy as np
import pandas as pd

import dask.dataframe as dd
from joblib import parallel_backend
from dask.diagnostics import ProgressBar

from sklearn.preprocessing import RobustScaler
from sklearn.model_selection import train_test_split
from sklearn.feature_selection import chi2, f_classif
from sklearn.metrics import confusion_matrix, classification_report

from rgf.sklearn import RGFClassifier
from sklearn.tree import DecisionTreeClassifier
from sklearn.ensemble import RandomForestClassifier

from matplotlib import pyplot as plt
import seaborn as sns

from src import utils

sns.set()

# dask progress bar
pbar = ProgressBar()
pbar.register()

# # dask parallel backend client
# client = Client()

# define thresholds as timedelta
BAD_THRESHOLD_NDAYS = np.timedelta64(14, "D")
WARNING_THRESHOLD_NDAYS = np.timedelta64(42, "D")

# inferred int32 types cause a type mismatch (int vs float) error when dask sees a null value
# null values cannot be interpreted as ints
custom_dtypes = {
    "date": "object",
    "serial_number": "object",
    "model": "object",
    "capacity_bytes": "float32",
    "failure": "float32",
    "smart_1_normalized": "float32",
    "smart_1_raw": "float32",
    "smart_2_normalized": "float32",
    "smart_2_raw": "float32",
    "smart_3_normalized": "float32",
    "smart_3_raw": "float32",
    "smart_4_normalized": "float32",
    "smart_4_raw": "float32",
    "smart_5_normalized": "float32",
    "smart_5_raw": "float32",
    "smart_7_normalized": "float32",
    "smart_7_raw": "float32",
    "smart_8_normalized": "float32",
    "smart_8_raw": "float32",
    "smart_9_normalized": "float32",
    "smart_9_raw": "float32",
    "smart_10_normalized": "float32",
    "smart_10_raw": "float32",
    "smart_11_normalized": "float32",
    "smart_11_raw": "float32",
    "smart_12_normalized": "float32",
    "smart_12_raw": "float32",
    "smart_13_normalized": "float32",
    "smart_13_raw": "float32",
    "smart_15_normalized": "float32",
    "smart_15_raw": "float32",
    "smart_16_normalized": "float32",
    "smart_16_raw": "float32",
    "smart_17_normalized": "float32",
    "smart_17_raw": "float32",
    "smart_22_normalized": "float32",
    "smart_22_raw": "float32",
    "smart_23_normalized": "float32",
    "smart_23_raw": "float32",
    "smart_24_normalized": "float32",
    "smart_24_raw": "float32",
    "smart_168_normalized": "float32",
    "smart_168_raw": "float32",
    "smart_170_normalized": "float32",
    "smart_170_raw": "float32",
    "smart_173_normalized": "float32",
    "smart_173_raw": "float32",
    "smart_174_normalized": "float32",
    "smart_174_raw": "float32",
    "smart_177_normalized": "float32",
    "smart_177_raw": "float32",
    "smart_179_normalized": "float32",
    "smart_179_raw": "float32",
    "smart_181_normalized": "float32",
    "smart_181_raw": "float32",
    "smart_182_normalized": "float32",
    "smart_182_raw": "float32",
    "smart_183_normalized": "float32",
    "smart_183_raw": "float32",
    "smart_184_normalized": "float32",
    "smart_184_raw": "float32",
    "smart_187_normalized": "float32",
    "smart_187_raw": "float32",
    "smart_188_normalized": "float32",
    "smart_188_raw": "float32",
    "smart_189_normalized": "float32",
    "smart_189_raw": "float32",
    "smart_190_normalized": "float32",
    "smart_190_raw": "float32",
    "smart_191_normalized": "float32",
    "smart_191_raw": "float32",
    "smart_192_normalized": "float32",
    "smart_192_raw": "float32",
    "smart_193_normalized": "float32",
    "smart_193_raw": "float32",
    "smart_194_normalized": "float32",
    "smart_194_raw": "float32",
    "smart_195_normalized": "float32",
    "smart_195_raw": "float32",
    "smart_196_normalized": "float32",
    "smart_196_raw": "float32",
    "smart_197_normalized": "float32",
    "smart_197_raw": "float32",
    "smart_198_normalized": "float32",
    "smart_198_raw": "float32",
    "smart_199_normalized": "float32",
    "smart_199_raw": "float32",
    "smart_200_normalized": "float32",
    "smart_200_raw": "float32",
    "smart_201_normalized": "float32",
    "smart_201_raw": "float32",
    "smart_218_normalized": "float32",
    "smart_218_raw": "float32",
    "smart_220_normalized": "float32",
    "smart_220_raw": "float32",
    "smart_222_normalized": "float32",
    "smart_222_raw": "float32",
    "smart_223_normalized": "float32",
    "smart_223_raw": "float32",
    "smart_224_normalized": "float32",
    "smart_224_raw": "float32",
    "smart_225_normalized": "float32",
    "smart_225_raw": "float32",
    "smart_226_normalized": "float32",
    "smart_226_raw": "float32",
    "smart_231_normalized": "float32",
    "smart_231_raw": "float32",
    "smart_232_normalized": "float32",
    "smart_232_raw": "float32",
    "smart_233_normalized": "float32",
    "smart_233_raw": "float32",
    "smart_235_normalized": "float32",
    "smart_235_raw": "float32",
    "smart_240_normalized": "float32",
    "smart_240_raw": "float32",
    "smart_241_normalized": "float32",
    "smart_241_raw": "float32",
    "smart_242_normalized": "float32",
    "smart_242_raw": "float32",
    "smart_250_normalized": "float32",
    "smart_250_raw": "float32",
    "smart_251_normalized": "float32",
    "smart_251_raw": "float32",
    "smart_252_normalized": "float32",
    "smart_252_raw": "float32",
    "smart_254_normalized": "float32",
    "smart_254_raw": "float32",
    "smart_255_normalized": "float32",
    "smart_255_raw": "float32",
}

# read all the data into one dataframe
# NOTE: assumes aws credentials are stored as env vars or in ~/.aws/credentials
DATA_ROOT_DIR = "/home/kachauha/Downloads/"
df4 = dd.read_parquet(
    os.path.join(DATA_ROOT_DIR, "data_Q4_2018_parquet"), engine="pyarrow", index=False
)
df3 = dd.read_parquet(
    os.path.join(DATA_ROOT_DIR, "data_Q3_2018_parquet"), engine="pyarrow", index=False
)
df = dd.concat(dfs=[df3, df4], interleave_partitions=True)

# deal with only seagate in this notebook
seagate_df = utils.optimal_repartition_df(df[df["model"].str.startswith("S")])

# convert from str to datetime
seagate_df["date"] = seagate_df["date"].astype("datetime64")

# get failed serial numbers
failed_sers = seagate_df[seagate_df["failure"] == 1]["serial_number"].compute()
failed_sers.head()

# start off with critical columns as indicated from 3 sources
CRITICAL_STATS = [
    1,
    5,
    7,
    10,
    184,
    187,
    188,
    189,
    190,
    193,
    194,
    196,
    197,
    198,
    201,
    240,
    241,
    242,
]

# NOTE - THESE LISTS ARE SUBJECT TO CHANGE
crit_cols_raw = ["smart_{}_raw".format(i) for i in CRITICAL_STATS]
crit_cols_normalized = ["smart_{}_normalized".format(i) for i in CRITICAL_STATS]

print("Seagate")
seagate_nan_ct = utils.get_nan_count_percent(
    seagate_df[crit_cols_raw + crit_cols_normalized]
)
seagate_nan_ct.compute()

seagate_corr_cols = ["failure", "capacity_bytes"] + crit_cols_raw + crit_cols_normalized
seagate_corr = seagate_df[seagate_corr_cols].corr().compute()
seagate_corr["failure"].abs().sort_values(ascending=False)

# re-selection of critical columns based on number of nans and correlations
CRITICAL_STATS = [1, 5, 7, 10, 187, 188, 190, 193, 194, 197, 198, 240, 241, 242]
crit_cols_raw = ["smart_{}_raw".format(i) for i in CRITICAL_STATS]
crit_cols_normalized = ["smart_{}_normalized".format(i) for i in CRITICAL_STATS]

seagate_corr_cols = ["failure", "capacity_bytes"] + crit_cols_raw + crit_cols_normalized
seagate_corr = seagate_df[seagate_corr_cols].corr().compute()
seagate_corr["failure"].abs().sort_values(ascending=False)

fig, ax = plt.subplots(figsize=(10, 10))
sns.heatmap(
    seagate_corr,
    ax=ax,
    vmin=-1,
    vmax=1,
    square=True,
    linewidths=0.5,
    cmap="RdBu",
    cbar_kws={"shrink": 0.5},
)

seagate_corr["smart_5_normalized"]["smart_198_normalized"]

# make sure we remove nans - sklearn doesnt deal well with those
test_df = seagate_df[
    [
        "failure",
        "smart_5_normalized",
        "smart_7_normalized",
        "smart_10_normalized",
        "smart_197_normalized",
        "smart_198_normalized",
    ]
].dropna()
test_df.isna().any().any().compute()

# split
Xtest = test_df.drop("failure", axis=1)
ytest = test_df["failure"]

with parallel_backend("dask"):
    # inspect deeper into highly correlated features
    print(chi2(Xtest, ytest))

# want big values
with parallel_backend("dask"):
    # inspect deeper into highly correlated features
    print(f_classif(Xtest, ytest))

# for now, keep the highly correlated ones but remove 194. remove 240, 242 as well (too low corr)
CRITICAL_STATS = [1, 5, 7, 10, 187, 188, 190, 193, 197, 198, 241]
crit_cols_raw = ["smart_{}_raw".format(i) for i in CRITICAL_STATS]
crit_cols_normalized = ["smart_{}_normalized".format(i) for i in CRITICAL_STATS]
seagate_df = seagate_df[
    ["date", "serial_number", "model", "capacity_bytes", "failure"]
    + crit_cols_raw
    + crit_cols_normalized
]

# what is the nan situation like
utils.get_nan_count_percent(seagate_df).compute()

# failed drives that have nan values
failed_df = seagate_df[seagate_df["serial_number"].isin(failed_sers)]
len_failed_df = len(failed_df)
utils.get_nan_count_percent(failed_df, divisor=len_failed_df).compute()

# # =============================== FOR DASK =============================== #
# # create meta of the resulting failed_df otherwise dask complains
# rul_meta = seagate_df._meta
# rul_meta = rul_meta.assign(rul_days=rul_meta['date'].max()-rul_meta['date'])
# # ======================================================================== #

# # get remaining useful life as diff(today, maxday)
# # reset index coz result is multiindexed. drop=True coz serial_number already exists as a col
# seagate_df = seagate_df.groupby('serial_number').apply(
#     utils.append_rul_days_column,
#     meta=rul_meta
# ).reset_index(drop=True)

# # how many drives had nan values close to fail day
# utils.get_nan_count_percent(failed_df[failed_df['rul_days']<=BAD_THRESHOLD_NDAYS], divisor=len_failed_df).compute()

# # how many drives had nan values exactly on fail day
# utils.get_nan_count_percent(failed_df[failed_df['failure']==1], divisor=len_failed_df).compute()

# drop the nans that occur in working hard drives, keep the ones in failing ones for now
seagate_df = seagate_df[
    (seagate_df["serial_number"].isin(failed_sers))
    | (
        (~seagate_df["smart_1_raw"].isna())
        & (~seagate_df["smart_5_raw"].isna())
        & (~seagate_df["smart_187_raw"].isna())
        & (~seagate_df["smart_193_raw"].isna())
        & (~seagate_df["smart_241_raw"].isna())
    )
]

# shouldnt have any nan values from working drives now
utils.get_nan_count_percent(seagate_df).compute()

# wont lose too many points if we just drop nans for failed ones
seagate_df = seagate_df[
    (~seagate_df["smart_1_raw"].isna())
    & (~seagate_df["smart_5_raw"].isna())
    & (~seagate_df["smart_187_raw"].isna())
    & (~seagate_df["smart_193_raw"].isna())
    & (~seagate_df["smart_241_raw"].isna())
]

# ensure no more nans
utils.get_nan_count_percent(seagate_df).compute()

# =============================== FOR DASK =============================== #
# create meta of the resulting failed_df otherwise dask complains
rul_meta = seagate_df._meta
rul_meta = rul_meta.assign(rul_days=rul_meta["date"].max() - rul_meta["date"])
# ======================================================================== #

# get remaining useful life as diff(today, maxday)
# reset index coz result is multiindexed. drop=True coz serial_number already exists as a col
seagate_df = (
    seagate_df.groupby("serial_number")
    .apply(utils.append_rul_days_column, meta=rul_meta)
    .reset_index(drop=True)
)

# remove working drive data that is recorded after [quarter end minus 6 weeks]
# because we dont know (as of quarter end) if those drives survived more than 6 weeks or not
seagate_df = seagate_df[
    (seagate_df["serial_number"].isin(failed_sers))
    | (seagate_df["rul_days"] >= WARNING_THRESHOLD_NDAYS)
]
# print(dd.compute(seagate_df.shape))

# NOTE: assignment must be done in th
# df.head()is order otherwise it wont be correct. FIXME
# assign all as good initially
seagate_df["status"] = 0

# overwrite those which have rul less than 6 weeks as warning
seagate_df["status"] = seagate_df["status"].mask(
    seagate_df["rul_days"] < WARNING_THRESHOLD_NDAYS, 1
)

# overwrite those which have rul less than 2 weeks as bad
seagate_df["status"] = seagate_df["status"].mask(
    seagate_df["rul_days"] < BAD_THRESHOLD_NDAYS, 2
)

seagate_df = utils.optimal_repartition_df(seagate_df)

num_working_serials = 7500
working_sers = seagate_df[~seagate_df["serial_number"].isin(failed_sers)][
    "serial_number"
].unique()
working_repr_sers = working_sers.sample(
    frac=(num_working_serials / len(working_sers))
).compute()

# preprocess only downsampled data
seagate_df = seagate_df[
    (seagate_df["serial_number"].isin(failed_sers))
    | seagate_df["serial_number"].isin(working_repr_sers)
].compute()

def rolling_featurize(
    df,
    window=6,
    drop_cols=("date", "failure", "capacity_bytes", "rul"),
    group_cols=("serial_number"),
    cap=True,
):
    """
    Extracts 6-day rolling features (6-day mean, std, coefficient of variation) from raw data
    in a pandas dataframe
    """
    # save the status labels
    # FIXME: if this is not a df, then earlier versions of pandas (0.19) complains
    statuses = df[["status"]]

    # group by serials, drop cols which are not to be aggregated
    if drop_cols is not None:
        grouped_df = df.drop(drop_cols, axis=1).groupby(group_cols)
    else:
        grouped_df = df.groupby(group_cols)

    # feature columns
    featcols = grouped_df.first().columns

    # get mean value in last 6 days
    means = grouped_df.rolling(window)[featcols].mean()

    # get std in last 6 days
    stds = grouped_df.rolling(window)[featcols].std()

    # coefficient of variation
    cvs = stds.divide(means, fill_value=0)

    # rename before mergeing
    means = means.rename(columns={col: "mean_" + col for col in means.columns})
    stds = stds.rename(columns={col: "std_" + col for col in stds.columns})
    cvs = cvs.rename(columns={col: "cv_" + col for col in cvs.columns})

    # combine features into one df
    res = means.merge(stds, left_index=True, right_index=True)
    res = res.merge(cvs, left_index=True, right_index=True)

    # drop rows where all columns are nans
    res = res.dropna(how="all")

    # fill nans created by cv calculation
    res = res.fillna(0)

    # capacity of hard drive
    if cap:
        capacities = (
            df[["serial_number", "capacity_bytes"]].groupby("serial_number").max()
        )
        res = res.merge(capacities, left_index=True, right_index=True)

    # bring serial number back as a col instead of index, preserve the corresponding indices
    res = res.reset_index(level=[0])

    # add status labels back.
    res = res.merge(statuses, left_index=True, right_index=True)

    return res

# MUST make sure indices are unique before processing
seagate_df = seagate_df.reset_index(drop=True)
seagate_feats_df = rolling_featurize(
    seagate_df,
    window=6,
    drop_cols=["date", "model", "capacity_bytes", "failure", "rul_days", "status"],
    group_cols=["serial_number"],
)

# infinities get added to df - remove these
seagate_feats_df = seagate_feats_df.replace([np.inf, -np.inf], np.nan).dropna()

seagate_feats_df.head()

# ========================================================================= #
#                             UNTESTED FUNCTION                             #
# ========================================================================= #
def test_expanding(
    df,
    drop_cols=("date", "failure", "capacity_bytes", "rul"),
    group_cols=("serial_number"),
    cap=True,
    num_days=False,
):
    # group by serials, drop cols which are not to be aggregated
    if drop_cols is not None:
        grouped_df = df.drop(drop_cols, axis=1).groupby(group_cols)
    else:
        grouped_df = df.groupby(group_cols)

    # rolling window mean. drop the group_cols since they dont vary in a group
    # .drop(group_cols, axis=1)
    means = grouped_df.expanding()[
        [col for col in grouped_df.first().columns if col not in group_cols]
    ].mean()
    means = means.rename(columns={col: "mean_" + col for col in means.columns})

    # rolling window srd. drop the group_cols since they dont vary in a group
    stds = grouped_df.expanding()[
        [col for col in grouped_df.first().columns if col not in group_cols]
    ].std(ddof=0)
    stds = stds.rename(columns={col: "std_" + col for col in stds.columns})

    # combine features into one df
    feats = means.merge(stds, left_index=True, right_index=True)

    # capacity of hard drive
    if cap:
        capacities = (
            df[["serial_number", "capacity_bytes"]].groupby("serial_number").max()
        )
        feats = feats.merge(capacities, left_index=True, right_index=True)

    # number of days of observed data available
    if num_days:
        days_per_drive = grouped_df.size().to_frame("num_days")
        feats = feats.merge(days_per_drive, left_index=True, right_index=True)

    return feats

# split into X,Y
X_arr = seagate_feats_df.drop(["serial_number", "status"], axis=1)
Y_arr = seagate_feats_df[["serial_number", "status"]]

del seagate_feats_df
gc.collect()

# split by serial number
# failed serials left after reduction
failed_sers_red = pd.Series(Y_arr["serial_number"].unique())
failed_sers_red = failed_sers_red[failed_sers_red.isin(failed_sers)]

# working serials left after reduction
working_sers_red = pd.Series(Y_arr["serial_number"].unique())
working_sers_red = working_sers_red[~working_sers_red.isin(failed_sers)]

# split working and failed
working_train, working_test = train_test_split(
    working_sers_red, test_size=0.2, random_state=42
)
failed_train, failed_test = train_test_split(
    failed_sers_red, test_size=0.2, random_state=42
)

# use serial numbers to generate train/test set
# CHECKED OK - train/test ratio 0.8, fail/work and overall both
X_train_work = X_arr[Y_arr["serial_number"].isin(working_train)]
X_train_fail = X_arr[Y_arr["serial_number"].isin(failed_train)]
X_train = pd.concat([X_train_work, X_train_fail])
del X_train_work, X_train_fail

Y_train_work = Y_arr[Y_arr["serial_number"].isin(working_train)]["status"]
Y_train_fail = Y_arr[Y_arr["serial_number"].isin(failed_train)]["status"]
Y_train = pd.concat([Y_train_work, Y_train_fail])
del Y_train_work, Y_train_fail

X_test_work = X_arr[Y_arr["serial_number"].isin(working_test)]
X_test_fail = X_arr[Y_arr["serial_number"].isin(failed_test)]
X_test = pd.concat([X_test_work, X_test_fail])
del X_test_work, X_test_fail

Y_test_work = Y_arr[Y_arr["serial_number"].isin(working_test)]["status"]
Y_test_fail = Y_arr[Y_arr["serial_number"].isin(failed_test)]["status"]
Y_test = pd.concat([Y_test_work, Y_test_fail])
del Y_test_work, Y_test_fail

# trecolumnsas regular dataset and scale
# with parallel_backend('dask'):
scaler = RobustScaler()
X_train_scaled = scaler.fit_transform(X_train)

# with joblib.parallel_backend('dask'):
dt_clf = DecisionTreeClassifier(class_weight="balanced")
dt_clf.fit(X_train_scaled, Y_train)

# get preds
dt_preds = dt_clf.predict(scaler.transform(X_test))
dt_confmat = confusion_matrix(Y_test, dt_preds)

print(np.around(dt_confmat, decimals=2))
print(classification_report(Y_test, dt_preds, target_names=["good", "warning", "bad"]))

rf_clf = RandomForestClassifier(n_estimators=128, class_weight="balanced", n_jobs=-1)
rf_clf.fit(X_train_scaled, Y_train)

# get preds
rf_preds = rf_clf.predict(scaler.transform(X_test))
rf_confmat = confusion_matrix(Y_test, rf_preds)

print(np.around(rf_confmat, decimals=2))
print(classification_report(Y_test, rf_preds, target_names=["good", "warning", "bad"]))

rgf_clf = RGFClassifier(max_leaf=2500, l2=1e-10, learning_rate=0.01)
rgf_clf.fit(X_train_scaled, Y_train)

# get preds
rgf_preds = rgf_clf.predict(scaler.transform(X_test))
rgf_confmat = confusion_matrix(Y_test, rgf_preds)

print(np.around(rgf_confmat / rgf_confmat.sum(axis=1, keepdims=True), decimals=2))
print(classification_report(Y_test, rgf_preds, target_names=["good", "warning", "bad"]))