03.ppseq_crossvalidation.py

# %% import and definition
import itertools as itt
import os
import pickle as pkl

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns
import torch
import xarray as xr
from tqdm.auto import tqdm

from routine.plotting import ppseq_plot_scatter, ppseq_plot_temp
from routine.ppseq import PPSeqE, unique_temp

SPK_PATH = "./intermediate/ppseq/spks_ds.nc"
INT_PATH = "./intermediate/ppseq_crossvalid"
FIG_PATH = "./figs/ppseq_crossvalid"
PARAM_NSPLT = 5
PARAM_NITER = 50
PARAM_TEMP_DUR = [5, 10, 25, 50, 100]
PARAM_NTEMP = [1, 5, 10]
PARAM_UNQ_TOL = 1e-1

os.makedirs(INT_PATH, exist_ok=True)
os.makedirs(FIG_PATH, exist_ok=True)

# %% ppseq
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
ds_spks = xr.load_dataset(SPK_PATH).rename(unit_id="cell")
spk = ds_spks["spks_ds"].dropna("frame", how="all")
spk = spk.where(spk > 2, other=0)
sess = np.unique(spk.coords["session"])
for ss in sess:
    print("processing session: {}".format(ss))
    spk_ss = spk.sel(frame=spk["session"] == ss)
    splt_idxs = np.array_split(np.array(spk_ss.coords["frame"]), PARAM_NSPLT)
    for temp_dur, ntemp, isplt in tqdm(
        list(itt.product(PARAM_TEMP_DUR, PARAM_NTEMP, range(PARAM_NSPLT)))
    ):
        train_idx = np.concatenate(splt_idxs[:isplt] + splt_idxs[isplt + 1 :])
        test_idx = splt_idxs[isplt]
        train_data = spk_ss.sel(frame=train_idx)
        test_data = spk_ss.sel(frame=test_idx)
        model = PPSeqE(
            num_templates=ntemp,
            num_neurons=int(spk_ss.sizes["cell"]),
            template_duration=temp_dur,
            alpha_a0=1.5,
            beta_a0=0.2,
            alpha_b0=1,
            beta_b0=0.1,
            alpha_t0=1.2,
            beta_t0=0.1,
            device=device,
        )
        torch.manual_seed(0)
        lps_fit, amplitudes_fit = model.fit(
            torch.from_numpy(np.array(train_data)), num_iter=PARAM_NITER
        )

        lps_train, amplitudes_train = model.predict(
            torch.from_numpy(np.array(train_data)), num_iter=PARAM_NITER
        )
        lps_test, amplitudes_test = model.predict(
            torch.from_numpy(np.array(test_data)), num_iter=PARAM_NITER
        )
        with open(
            os.path.join(
                INT_PATH, "model-{}-{}-{}-{}.pkl".format(ss, ntemp, temp_dur, isplt)
            ),
            "wb",
        ) as pklf:
            pkl.dump(
                {
                    "ss": ss,
                    "ntemp": ntemp,
                    "temp_dur": temp_dur,
                    "isplt": isplt,
                    "model": model,
                    "X_train": train_data,
                    "X_test": test_data,
                    "lps_fit": lps_fit,
                    "lps_train": lps_train,
                    "lps_test": lps_test,
                    "amp_fit": amplitudes_fit.cpu(),
                    "amp_train": amplitudes_train.cpu(),
                    "amp_test": amplitudes_test.cpu(),
                },
                pklf,
            )

# %% plot results
fig_path_indv = os.path.join(FIG_PATH, "indv")
fig_path_agg = os.path.join(FIG_PATH, "agg")
os.makedirs(fig_path_indv, exist_ok=True)
os.makedirs(fig_path_agg, exist_ok=True)
pkl_files = list(filter(lambda fn: fn.endswith(".pkl"), os.listdir(INT_PATH)))
res_df = []
param_df = []
for pklf in tqdm(pkl_files):
    with open(os.path.join(INT_PATH, pklf), "rb") as pklf:
        ds = pkl.load(pklf)
    model = ds["model"]
    nunq = unique_temp(model, ds["amp_train"], tol=PARAM_UNQ_TOL, modify_model=False)
    res_df.append(
        pd.DataFrame(
            {
                "ss": ds["ss"],
                "ntemp": ds["ntemp"],
                "temp_dur": ds["temp_dur"],
                "isplt": ds["isplt"],
                "nunq": nunq,
                "split": ["train", "test"],
                "lps": [
                    np.array(ds["lps_train"]).max().item(),
                    np.array(ds["lps_test"]).max().item(),
                ],
            }
        )
    )
    param_df.append(
        pd.DataFrame(
            {
                "ss": ds["ss"],
                "ntemp": ds["ntemp"],
                "temp_dur": ds["temp_dur"],
                "isplt": ds["isplt"],
                "temp_widths": np.median(model.template_widths, axis=1),
            }
        )
    )
    fig_temp = ppseq_plot_temp(ds["X_train"], model, ds["amp_train"])
    fig_temp.write_html(
        os.path.join(
            fig_path_indv,
            "temp-{}-{}-{}-{}.html".format(
                ds["ss"], ds["ntemp"], ds["temp_dur"], ds["isplt"]
            ),
        )
    )
    fig_scatter = ppseq_plot_scatter(ds["X_train"], model, ds["amp_train"])
    fig_scatter.write_html(
        os.path.join(
            fig_path_indv,
            "scatter-{}-{}-{}-{}-train.html".format(
                ds["ss"], ds["ntemp"], ds["temp_dur"], ds["isplt"]
            ),
        )
    )
    fig_scatter = ppseq_plot_scatter(ds["X_test"], model, ds["amp_test"])
    fig_scatter.write_html(
        os.path.join(
            fig_path_indv,
            "scatter-{}-{}-{}-{}-test.html".format(
                ds["ss"], ds["ntemp"], ds["temp_dur"], ds["isplt"]
            ),
        )
    )
res_df = pd.concat(res_df, ignore_index=True)
param_df = pd.concat(param_df, ignore_index=True)

# %% plot agg results
g = sns.FacetGrid(
    res_df,
    row="split",
    col="ss",
    sharey=False,
    sharex=True,
    margin_titles=True,
    legend_out=True,
    aspect=2,
)
g.map_dataframe(sns.boxplot, x="temp_dur", y="lps", dodge=True, hue="ntemp", fill=False)
g.map_dataframe(sns.swarmplot, x="temp_dur", y="lps", dodge=True, hue="ntemp")
g.add_legend()
g.figure.savefig(os.path.join(fig_path_agg, "summary.svg"), bbox_inches="tight")

# %% plot agg params
fig, ax = plt.subplots()
ax = sns.boxplot(
    param_df, x="temp_dur", y="temp_widths", hue="ntemp", dodge=True, fill=False, ax=ax
)
fig.savefig(os.path.join(fig_path_agg, "temp_widths.svg"), bbox_inches="tight")