Step 1: Exploring Data

import os

import numpy as np
import pandas as pd

import matplotlib
import seaborn as sns
from matplotlib import pyplot as plt

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

from joblib import parallel_backend
from dask.distributed import Client

from umap import UMAP
from sklearn.decomposition import PCA
from sklearn.preprocessing import StandardScaler

from src import utils

sns.set()

pd.set_option("display.max_rows", 75)
pd.set_option("display.max_columns", 500)

pbar = ProgressBar()
pbar.register()

# 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 (local) data into one dataframe
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
)
df2 = dd.read_parquet(
    os.path.join(DATA_ROOT_DIR, "data_Q2_2018_parquet"), engine="pyarrow", index=False
)
df1 = dd.read_parquet(
    os.path.join(DATA_ROOT_DIR, "data_Q1_2018_parquet"), engine="pyarrow", index=False
)
df = utils.optimal_repartition_df(
    dd.concat(dfs=[df1, df2, df3, df4], interleave_partitions=True)
)

# split into working vs failed drives serial number
failed_sers = df[df["failure"] == 1][["serial_number", "model"]].compute()
working_sers = (
    df[~df["serial_number"].isin(failed_sers["serial_number"])][
        ["serial_number", "model"]
    ]
    .drop_duplicates()
    .compute()
)

# sanity check - make sure we havent labelled a serial number as BOTH working and failed
assert ~working_sers["serial_number"].isin(failed_sers["serial_number"]).any()

# how many of each class
print(f"{len(failed_sers)} failed drives, {len(working_sers)} working drives")

# split data into vendors, print how many data points available for each vendor
seagate_df = df[df["model"].str.startswith("S")]
print(dd.compute(seagate_df.shape))
hgst_df = df[df["model"].str.startswith("HG")]
print(dd.compute(hgst_df.shape))
toshiba_df = df[df["model"].str.startswith("T")]
print(dd.compute(toshiba_df.shape))
wdc_df = df[df["model"].str.startswith("W")]
print(dd.compute(wdc_df.shape))
hitachi_df = df[df["model"].str.startswith("Hi")]
print(dd.compute(hitachi_df.shape))

# how many drives by model, and what is the fail percentage
model_stats = pd.merge(
    failed_sers["model"].value_counts().to_frame("failed_count").reset_index(),
    working_sers["model"].value_counts().to_frame("working_count").reset_index(),
    how="outer",
).fillna(0)

# rename the index column as the model column
model_stats.rename(columns={"index": "model"}, inplace=True)

# total count of model, raw value and as a percent of total drives
model_stats["total_count"] = model_stats["failed_count"] + model_stats["working_count"]
model_stats["total_percent"] = np.around(
    100 * model_stats["total_count"] / model_stats["total_count"].sum(), decimals=2
)

# percentage of instances that have failed, per model
model_stats["fail_percent"] = np.around(
    100 * model_stats["failed_count"] / model_stats["total_count"], decimals=2
)

# add manufacturer/vendor column
model_stats["vendor"] = model_stats["model"].apply(utils.get_vendor)

model_stats.sort_values(by=["total_count"], ascending=False)

# how many failed drives in each vendor
for mdf in (seagate_df, hgst_df, toshiba_df, wdc_df, hitachi_df):
    num_failed_sers = (mdf["failure"] == 1).sum().compute()
    num_working_sers = mdf["serial_number"].nunique().compute() - num_failed_sers
    print(
        "{} failed drives\n{} working drives".format(num_failed_sers, num_working_sers)
    )
    print(
        "================================================================================"
    )

# get vendor-wise stats for failed vs working
vendor_stats = model_stats.groupby("vendor").sum()
vendor_stats["fail_percent"] = np.around(
    100
    * vendor_stats["failed_count"]
    / (vendor_stats["failed_count"] + vendor_stats["working_count"]),
    decimals=2,
)
vendor_stats

# within drives of a manufacturer, what is the distribution of models like
for mdf in (seagate_df, hgst_df, toshiba_df, wdc_df, hitachi_df):
    valcts = mdf["model"].value_counts().compute()
    sns.barplot(x=valcts.index, y=(valcts.values / valcts.sum()))
    plt.xticks(rotation=90)
    plt.show()

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]

# number of nans
print("All data")
nan_count = utils.get_nan_count_percent(
    df[crit_cols_raw + crit_cols_normalized]
).compute()
nan_count

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

print("WDC")
wdc_nan_ct = utils.get_nan_count_percent(
    wdc_df[crit_cols_raw + crit_cols_normalized]
).compute()
wdc_nan_ct

print("HGST")
hgst_nan_ct = utils.get_nan_count_percent(
    hgst_df[crit_cols_raw + crit_cols_normalized]
).compute()
hgst_nan_ct

print("Hitachi")
hitachi_nan_ct = utils.get_nan_count_percent(
    hitachi_df[crit_cols_raw + crit_cols_normalized]
).compute()
hitachi_nan_ct

print("Toshiba")
toshiba_nan_ct = utils.get_nan_count_percent(
    toshiba_df[crit_cols_raw + crit_cols_normalized]
).compute()
toshiba_nan_ct

# let's see if there are unusual amounts of nan's within each model too
# if there is, then that means there was probably something wrong with data collection
top_models_seagate = ["ST12000NM0007", "ST4000DM000", "ST8000NM0055"]
for model in top_models_seagate:
    # get nan counts and percents, but print only those which are not 0 or 1
    nanct = utils.get_nan_count_percent(seagate_df[seagate_df["model"] == model])
    print(nanct[(nanct["percent"] != 0) & (nanct["percent"] != 1)].compute())

# general description
# NOTE: the columns with all values NAN must be removed otherwise a value error is thrown
summary = df[nan_count.index[nan_count["percent"] != 1]].describe().compute()
summary

# number of rows to use for plotting - sampling is done iff number of rows is too much to keep in memory
num_rows_to_sample = 5000000

# plot histograms for overall data
# it would probably be more insightful to do this per vendor, but start with overall
# dont make the hist if all values are nan
hist_cols = nan_count.index[
    nan_count["percent"] < 0.25
]  # + ['failure', 'capacity_bytes']
len_df = len(df)
for col in hist_cols:
    # get all the values for this column
    if num_rows_to_sample < len_df:
        data = df[col].sample(frac=num_rows_to_sample / len_df).compute()
    else:
        data = df[col].compute()

    # plot only the non null values
    print(
        data.isna().sum(),
        "out of",
        data.shape[0],
        "are NaN values. These are not shown on the graph below",
    )
    sns.distplot(data[~data.isna()])
    plt.show()

# correlation with failure
# corr_cols = ['failure', 'capacity_bytes'] + crit_cols_raw + crit_cols_normalized
corr_cols = ["failure", "capacity_bytes"] + list(
    nan_count.index[nan_count["percent"] != 1]
)
downsampled_sers = (
    failed_sers["serial_number"].sample(n=500).values.tolist()
    + working_sers["serial_number"].sample(n=75000).values.tolist()
)
corr_mat = df[df["serial_number"].isin(downsampled_sers)][corr_cols].corr().compute()
fig, ax = plt.subplots(figsize=(10, 10))
sns.heatmap(
    corr_mat,
    ax=ax,
    vmin=-1,
    vmax=1,
    square=True,
    linewidths=0.5,
    cmap="RdBu",
    cbar_kws={"shrink": 0.5},
)

# only those attributes with less than NAN_PERCENT_THRESHOLD% values as NaN's
# will be selected for computing the correaltion matirx
NAN_PERCENT_THRESHOLD = 0.5

# correlation matrix for seagaet drives
seagate_corr_cols = ["failure", "capacity_bytes"] + list(
    seagate_nan_ct.index[seagate_nan_ct["percent"] < NAN_PERCENT_THRESHOLD]
)
seagate_corr = (
    seagate_df[seagate_df["serial_number"].isin(downsampled_sers)][seagate_corr_cols]
    .corr()
    .compute()
)
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},
)

# correlation matrix for wdc drives
wdc_corr_cols = ["failure", "capacity_bytes"] + list(
    wdc_nan_ct.index[wdc_nan_ct["percent"] < NAN_PERCENT_THRESHOLD]
)
wdc_corr = wdc_df[wdc_corr_cols].corr().compute()
fig, ax = plt.subplots(figsize=(10, 10))
sns.heatmap(
    wdc_corr,
    ax=ax,
    vmin=-1,
    vmax=1,
    square=True,
    linewidths=0.5,
    cmap="RdBu",
    cbar_kws={"shrink": 0.5},
)

# correlation matrix for hgst drives
hgst_corr_cols = ["failure", "capacity_bytes"] + list(
    hgst_nan_ct.index[hgst_nan_ct["percent"] < NAN_PERCENT_THRESHOLD]
)
hgst_corr = hgst_df[hgst_corr_cols].corr().compute()
fig, ax = plt.subplots(figsize=(10, 10))
sns.heatmap(
    hgst_corr,
    ax=ax,
    vmin=-1,
    vmax=1,
    square=True,
    linewidths=0.5,
    cmap="RdBu",
    cbar_kws={"shrink": 0.5},
)

# correlation matrix for hitachi drives
hitachi_corr_cols = ["failure", "capacity_bytes"] + list(
    hitachi_nan_ct.index[hitachi_nan_ct["percent"] < NAN_PERCENT_THRESHOLD]
)
hitachi_corr = hitachi_df[hitachi_corr_cols].corr().compute()
fig, ax = plt.subplots(figsize=(10, 10))
sns.heatmap(
    hitachi_corr,
    ax=ax,
    vmin=-1,
    vmax=1,
    square=True,
    linewidths=0.5,
    cmap="RdBu",
    cbar_kws={"shrink": 0.5},
)

# correlation matrix for toshiba drives
toshiba_corr_cols = ["failure", "capacity_bytes"] + list(
    toshiba_nan_ct.index[toshiba_nan_ct["percent"] < NAN_PERCENT_THRESHOLD]
)
toshiba_corr = toshiba_df[toshiba_corr_cols].corr().compute()
fig, ax = plt.subplots(figsize=(10, 10))
sns.heatmap(
    toshiba_corr,
    ax=ax,
    vmin=-1,
    vmax=1,
    square=True,
    linewidths=0.5,
    cmap="RdBu",
    cbar_kws={"shrink": 0.5},
)

# TODO: Might be better to call compute to get the combined data of all serials in subset
# as opposed to calling it for earch serial number in for loop

# NOTE: running this cell will take a VERY long time (~1hr on intel i7 w/ 16GB ram)
# adjust NUM_DRIVES_TO_SAMPLE to select a small subset to use for plotting
NUM_DRIVES_TO_SAMPLE = 50

# plots for smart stat 5
# cols_to_plot = ['smart_5_raw', 'smart_7_raw']
cols_to_plot = crit_cols_raw + crit_cols_normalized
cols_to_plot.remove("smart_201_raw")  # too many nans
cols_to_plot.remove("smart_201_normalized")  # too many nans

# one figure per smart stat
figs = [plt.figure(i, figsize=(10, 10)) for i in range(len(cols_to_plot))]
axes = [f.add_subplot(111) for f in figs]
for colname, ax in zip(cols_to_plot, axes):
    ax.set_title(
        "{} vs Time for {} Failed Drives".format(colname, NUM_DRIVES_TO_SAMPLE)
    )
    ax.set_xlabel("Number of Days")
    ax.set_ylabel(colname)

# keep track of what hard drives were used to generate the data. NOTE: only the first 16 chars of ser will be saved
failed_ser_subset = np.empty(shape=(NUM_DRIVES_TO_SAMPLE), dtype="<S16")

# make the figures
for i, ser in enumerate(
    failed_sers["serial_number"].sample(NUM_DRIVES_TO_SAMPLE, random_state=42)
):
    # log serial numbers which are being used
    failed_ser_subset[i] = ser
    print("{} / {}. Drive serial number {}".format(i + 1, NUM_DRIVES_TO_SAMPLE, ser))

    # get teh data to make the figures
    drive_data = df[df["serial_number"] == ser][cols_to_plot].compute()

    # dummy x axis data
    xvals = [i for i in range(drive_data.shape[0])]

    # make the plot
    for ax, c in zip(axes, cols_to_plot):
        ax.plot(xvals, drive_data[c])

# save the figures
for f, c in zip(figs, cols_to_plot):
    f.savefig("../../../reports/figures/{}_failed.png".format(c))

# save the serial numbres used in figures
np.save("failed_graphs_serials", failed_ser_subset)

# sample to use for plotting
NUM_DRIVES_TO_SAMPLE = 50

# one figure per smart stat
for colname, ax in zip(cols_to_plot, axes):
    ax.cla()  # clear data from previous plots
    ax.set_title(
        "{} vs Time for {} Working Drives".format(colname, NUM_DRIVES_TO_SAMPLE)
    )
    ax.set_xlabel("Number of Days")
    ax.set_ylabel(colname)

# keep track of what hard drives were used to generate the data. NOTE: only the first 16 chars of ser will be saved
working_ser_subset = np.empty(shape=(NUM_DRIVES_TO_SAMPLE), dtype="<S16")

# make the figures
for i, ser in enumerate(
    working_sers["serial_number"].sample(NUM_DRIVES_TO_SAMPLE, random_state=42)
):
    # log serial numbers which are being used
    working_ser_subset[i] = ser
    print("{} / {}. Drive serial number {}".format(i + 1, NUM_DRIVES_TO_SAMPLE, ser))

    # get teh data to make the figures
    drive_data = df[df["serial_number"] == ser][cols_to_plot].compute()

    # dummy x axis data
    xvals = [i for i in range(drive_data.shape[0])]

    # make the plot
    for ax, c in zip(axes, cols_to_plot):
        ax.plot(xvals, drive_data[c])

# save the figures
for f, c in zip(figs, cols_to_plot):
    f.savefig("../../../reports/figures/{}_working.png".format(c))

# save the serial numbres used in figures
np.save("working_graphs_serials", working_ser_subset)

plt.close("all")

# get the groupby object. there should be one gourp per serial nunmbers
# get only those columns which have less than 50 percent nan values
grouped_data_cols = ["serial_number", "model", "capacity_bytes"] + list(
    nan_count.index[nan_count["percent"] < 0.5]
)
groups = df[grouped_data_cols].groupby("serial_number")

# get simple stats to represent time series of each hard drive
group_means = groups.mean().compute().add_prefix("mean_")
group_stds = groups.std().compute().add_prefix("std_")
# group_days = groups.size().compute().to_frame('days')
# group_days.index.name = None    # to match the other agg results, so that it can be concatenated easily

# put the stats together into one df
group_stats = pd.concat([group_means, group_stds], axis=1)

# make serial number a column instead of index. will be easier for calc later
group_stats = group_stats.reset_index()

# need to add failed label
group_stats["failure"] = group_stats["serial_number"].isin(failed_sers["serial_number"])

group_stats.head()

# FIXME: find a smarter way to deal with nans. drop columns not rows
# for now, just drop it. this drops ~25k observations and leaves us with ~83k
clean_group_stats = group_stats.dropna(how="any")

# make sure we still have enough failure drive data
print(clean_group_stats["failure"].sum(), "failed drives data retained")

clean_group_stats[clean_group_stats["failure"]].head()

# scale the data and find the top principal components
scaler = StandardScaler()
pca = PCA(n_components=3, random_state=42, whiten=True).fit_transform(
    scaler.fit_transform(clean_group_stats.drop(["serial_number", "failure"], axis=1))
)

# plot pca
colors = ["blue", "red"]
fig = plt.figure(figsize=(15, 15))
ax = fig.add_subplot(111, projection="3d")
ax.scatter(
    pca[:, 0],
    pca[:, 1],
    pca[:, 2],
    c=clean_group_stats["failure"],
    cmap=matplotlib.colors.ListedColormap(colors),
)

# plot umap
client = Client()
with parallel_backend("dask"):
    umap = UMAP(n_components=3, random_state=42).fit_transform(
        scaler.fit_transform(
            clean_group_stats.drop(["serial_number", "failure"], axis=1)
        )
    )

# plot umap
colors = ["blue", "red"]
fig = plt.figure(figsize=(15, 15))
ax = fig.add_subplot(111, projection="3d")
ax.scatter(
    umap[:, 0],
    umap[:, 1],
    umap[:, 2],
    c=clean_group_stats["failure"],
    cmap=matplotlib.colors.ListedColormap(colors),
)

# tsne embeddings may be more meaningful than pca, but it is extremely slow
# tsne = TSNE(n_components=3, random_state=42)\
#         .fit_transform(scaler.fit_transform(clean_group_stats.drop(['serial_number', 'failure'], axis=1)))

# columns to extract from dataset for analyzing
analysis_cols = ["date", "serial_number", "smart_9_raw", "smart_9_normalized"]

# columns for which to make histogram
hist_cols = ["smart_9_raw", "smart_9_normalized"]

# since failure=1 is marked on the last day that a drive worked, the entries on that day woudl be the most recent ones
# NOTE: this is a property of backblaze dataset, and may not generalize
failed_last_smart9_df = df[df["failure"] == 1][analysis_cols].compute()

# do the same for working drives
working_grouped = df[~df["serial_number"].isin(failed_sers["serial_number"])][
    hist_cols + ["serial_number"]
].groupby("serial_number")
working_last_smart9_df = working_grouped.max()
dd.compute(working_last_smart9_df.shape)

# size is not too bad -- we can bring it into memory
working_last_smart9_df = working_last_smart9_df.compute()

# working and failed plots together
for col in hist_cols:
    # get figure and axes for plotting current column
    fig, ax = plt.subplots()

    # plot failed and working onto the axes one by one
    data = working_last_smart9_df[col]
    sns.distplot(data[~data.isna()], ax=ax, label="working")
    data = failed_last_smart9_df[col]
    sns.distplot(data[~data.isna()], ax=ax, label="failed")
    plt.legend()
    plt.show()