Reduced models¶

BASEDIR = '/path/where/you/saved/data/'

# append code repository to pythonpath
import sys
sys.path.append(BASEDIR + 'code')

# import from python ecosystem
import numpy as np
import pandas as pd
import os
import matplotlib.pyplot as plt
from functools import partial
import dask
from collections import defaultdict
from functools import partial
import distributed

%matplotlib inline

from data_base import DataBase
from visualize.histogram import histogram
from visualize.rasterplot import rasterplot
from data_base.analyze.temporal_binning import temporal_binning_pd

import Interface as I

whiskers = ['B1', 'B2', 'B3', 'C1', 'C2', 'C3', 'D1', 'D2', 'D3', 'E2']

models = range(1,8)

model = 1
whisker = 'C2'

# each compressed numpy archive contains a single numpy array
sa = np.load(BASEDIR + "model_{}__binned_synapse_activations_spatiotemporal_EI__{}.npz".format(model, whisker))
sa = sa['arr_0']

saexc = sa[0, :, :, :]
sainh = sa[1, :, :, :]

# get mean count of active synapses onto the whole neuron in each timebin
saexc_mean = saexc.sum(axis = 0).mean(axis = 0)
sainh_mean = sainh.sum(axis = 0).mean(axis = 0)

model = 1
whisker = 'C2'

# load spike times for a principal whisker stimulus
st = pd.read_csv(BASEDIR + "model_{}__spike_times_MCM__{}.csv".format(model, whisker))

# we will store all the parameters of the reduced models in a dictionary
# here we show how to make a reduced model from one exemplary multi-compartmental model
# the procedure was repeated for all 7 MCMs
rm_dict = {}
rm_dict[model] = {}

import modular_reduced_model_inference
from modular_reduced_model_inference import RaisedCosineBasis

RaisedCosineBasis_spatial = modular_reduced_model_inference.RaisedCosineBasis(a = 3, c = 5, phis = np.arange(3,12, 1.5), width = 30, reversed_ = False)
RaisedCosineBasis_temporal = modular_reduced_model_inference.RaisedCosineBasis(a = 2, c = 1, phis = np.arange(1,11, 0.5), width = 80, reversed_ = True)

# we also provide precomputed spatiotemporal filters for all models used in our paper
# which are visualized in the next section

model = 1
whiskers = ['C2'] # it is also possible to infer reduced models from only a subset of the data. We use this here as an example,
# with only the principal (C2) whisker. To run inference on the full dataset, uncomment the line below:
# whiskers = ['B1', 'B2', 'B3', 'C1', 'C2', 'C3', 'D1', 'D2', 'D3', 'E2']
# be aware that a large amount of RAM is necessary to handle the full dataset

mdb_list = []
for whisker in whiskers:
    mdb_list.append({'spike_times': BASEDIR + "model_{}__spike_times_MCM__{}.csv".format(model, whisker),
           'binned_synapse_activation':BASEDIR + "model_{}__binned_synapse_activations_spatiotemporal_EI__{}.npz".format(model, whisker)})

# this cell takes a few minutes to run as it loads and uncompresses all the MCM simulation data
# uncompression takes a long time so you may want to save the uncompressed array in case you want to use it again
name = 'test'
rm = modular_reduced_model_inference.Rm(name,
    mdb = mdb_list,
    tmin = 260, #timepoint_of_max_response
    tmax = 261, #timepoint_of_max_response+1
    width = 80)

# make simulation data available to the reduced model
D = modular_reduced_model_inference.DataExtractor_spatiotemporalSynapseActivation
rm.add_data_extractor('spatiotemporalSa', D())
rm.add_data_extractor('st', modular_reduced_model_inference.DataExtractor_spiketimes())
rm.add_data_extractor('y', modular_reduced_model_inference.DataExtractor_spikeInInterval())
rm.add_data_extractor('ISI', modular_reduced_model_inference.DataExtractor_ISI())

st = rm.data_extractors['st'].get() # spike times

# define a strategy which evaluates simulation trials based on spatiotemporal filters
RaisedCosineBasis_spatial = RaisedCosineBasis(a = 3, c = 5, phis = np.arange(3,12, 1.5), width = 30, reversed_ = False)
RaisedCosineBasis_temporal = RaisedCosineBasis(a = 2, c = 1, phis = np.arange(1,11, 0.5), width = 80, reversed_ = True)
len_z = len(RaisedCosineBasis_spatial.phis) # number of parameters of the spatial filter
len_t = len(RaisedCosineBasis_temporal.phis)  # number of parameters of the temporal filter
strategy = modular_reduced_model_inference.Strategy_spatiotemporalRaisedCosine('SAspatiotemporalRaisedCosine', RaisedCosineBasis_spatial, RaisedCosineBasis_temporal)

rm.add_strategy(strategy)

# add solver
solver = modular_reduced_model_inference.Solver_COBYLA('cobyla')
strategy.add_solver(solver)

for lv in range(3): # make 3 random train/test splits
    # rm.DataSplitEvaluation.add_random_split('.7_{}'.format(lv), percentage_train=.7)
    rm.DataSplitEvaluation.add_isi_dependent_random_split('ISI50_.7_{}'.format(lv),min_isi = 50, percentage_train=.7)

result_df = rm.get_results(client)
# If you get ValueError: Only one class present in y_true. ROC AUC score is not defined in that case, it means that the
# reduced model failed (this happens for example due to a lack of spikes in the interval used for inference)

RaisedCosineBasis_spatial = RaisedCosineBasis(a = 3, c = 5, phis = np.arange(3,12, 1.5), width = 30, reversed_ = False)
RaisedCosineBasis_temporal = RaisedCosineBasis(a = 2, c = 1, phis = np.arange(1,11, 0.5), width = 80, reversed_ = True)
strategy = modular_reduced_model_inference.Strategy_spatiotemporalRaisedCosine(
    'SAspatiotemporalRaisedCosine',
    RaisedCosineBasis_spatial,
    RaisedCosineBasis_temporal)
# manually do what happens in rm.add_strategy(strategy) in order to get the shape of the basis functions
# this avoids having to load in the dataset
strategy.len_z = len(RaisedCosineBasis_spatial.phis) # number of parameters of the spatial filter
strategy.len_t = len(RaisedCosineBasis_temporal.phis)  # number of parameters of the temporal filter
strategy.groups = [('EXC',), ('INH',)] # excitatory and inhibitory
strategy.convert_x = partial(strategy._convert_x_static, strategy.groups, strategy.len_z)

rm_dict[model]['spatiotemporal_filters'] = spatiotemporal_filters

# calculate all the WNI values at relevant timepoints - we will need them again to estimate the ISI-dependent penalty
@dask.delayed
def calculate_WNI_values(sa_file, kernel_dict, model, whisker, outdir):

    timebins = len(range(220, 271))

    sa = np.load(sa_file)
    sa = sa['arr_0']
    n_cells = sa.shape[2]
    WNI_df = pd.DataFrame(index = range(n_cells), columns = range(220, 271))
    SAexc = sa[0,:,:,:]
    SAinh = sa[1,:,:,:]


    SAexc = np.asarray(SAexc)
    SAinh = np.asarray(SAinh)

    s_exc = kernel_dict['s_exc']
    s_inh = kernel_dict['s_inh']
    t_exc = kernel_dict['t_exc']
    t_inh = kernel_dict['t_inh']



    SAinh_cumulative = np.ndarray((n_cells, 400))
    SAexc_cumulative = np.ndarray((n_cells, 400))

    for t, timebin in enumerate(range(220-80, 271)): # for all timebins relevant to period of interest
        ## get excitatory and inhibitory input, spatially binned for the CURRENT timebin
        SAexc_timebin = SAexc[:, :, timebin]
        SAinh_timebin = SAinh[:, :, timebin]

        ## apply spatial kernel to the current timebin
        SAexc_timebin = sum([o*s for o, s in zip(SAexc_timebin, s_exc)])
        SAinh_timebin = sum([o*s for o, s in zip(SAinh_timebin, s_inh)])



        for cell in range(n_cells):
            SAinh_cumulative[cell][timebin] = SAinh_timebin[cell]
            SAexc_cumulative[cell][timebin] = SAexc_timebin[cell]



    ## apply temporal kernel to timebins of interest
    for t, timebin in enumerate(range(220, 271)):
#         print timebin
        for cell in range(n_cells):

            SAexc_window = SAexc_cumulative[cell][timebin-79:timebin+1]
            SAinh_window = SAinh_cumulative[cell][timebin-79:timebin+1]



            SAexc_window = sum([o*s for o, s in zip(SAexc_window, t_exc)])
            SAinh_window = sum([o*s for o, s in zip(SAinh_window, t_inh)])

            ## get weighted net input for each cell
            WNI = SAexc_window + SAinh_window
            WNI_df.iloc[cell, t] = WNI
    name = 'model_{}__{}.csv'.format(model, whisker)
    WNI_df.to_csv(outdir + name)

# we provide WNI values for the exemplary spatiotemporal filters generated in this notebook under the following key:
outdir = BASEDIR + 'precalculated_WNI_values__model_1__C2'

# to calculate them yourself, make a folder to store them in:
outdir = BASEDIR + 'precalculated_WNI_values/'
os.mkdir(outdir)

# takes about 20 mins per whisker on a single core
delayeds = []
index = []

for whisker in whiskers:
    sa_file = BASEDIR + "model_{}__binned_synapse_activations_spatiotemporal_EI__{}.npz".format(model, whisker)
    d = calculate_WNI_values(sa_file, spatiotemporal_filters, model, whisker, outdir)
    delayeds.append(d)
    index.append((model, whisker))

# to prevent memory issues - each worker will load around 25GB of data
# adjust the number of cores working per node according to your memory limits
workers_per_node = 5
from itertools import groupby
workers = client.scheduler_info()['workers'].keys()
workers = sorted(workers, key = lambda x: x.split(':')[1])
w = [list(w)[:workers_per_node] for group, w in groupby(workers, lambda x: x.split(':')[1])]
w = [w for w in w for w in w]

futures = client.compute(delayeds, workers=w)

def group_bins(helper_df, min_items = 10):
    # https://codereview.stackexchange.com/questions/12753/taking-the-first-few-items-of-a-list-until-a-target-sum-is-exceeded/12759

    total_items = 0
    rows_to_group = []
    groupings = []

    for row in range(len(helper_df.index)):

        items = helper_df.iloc[row, 3]
        total_items += items # keep a running total of number of datapoints
        rows_to_group.append(row)

        if total_items >= min_items: # once we have enough datapoints, save the grouping and continue
            total_items = 0
            groupings.append(rows_to_group)
            rows_to_group = []

        elif row == range(len(helper_df.index))[-1]: # or if it's the last bin, which might get skipped otherwise
            groupings.append([row])

    # check if the last bin is too small, merge with second to last if it is
    last_binsize = 0
    for row in groupings[-1]:
        last_binsize += helper_df.iloc[row, 3]

    if last_binsize < min_items:
        new_grouping = groupings[:-2]
        last_bins = groupings[-2:]
        new_grouping.append([i for sublist in last_bins for i in sublist])
        groupings = new_grouping

    return groupings

def signchange(x,y):
    if x / abs(x) == y / abs(y):
        return False
    else:
        return True

def linear_interpolation_between_pairs(X,Y, x, top = 'inf', bottom = 'min'):
    if x > max(X):
        if top == 'inf':
            result = np.inf
        elif top == 'max':
            result = max(Y)
    elif x < min(X):
        if bottom == 'inf':
            result = -np.inf
        elif bottom == 'min':
            result = min(Y)
    elif float(x) in X: # don't need to interpolate, causes assertion error
        result = Y[X.index(x)]
    else:
        pair = [lv for lv in range(len(X)-1) if signchange(X[lv]-x, X[lv+1]-x)]
#         X[lv] <= x < X[lv + 1]]
        assert(len(pair) == 1)
        pair = pair[0]
        m = (Y[pair+1]-Y[pair]) / (X[pair+1]-X[pair])
        c = Y[pair]-X[pair]*m

        result = m*x+c
    return result

def variable_stepsize_nonlinearity(wni_values, spike_times, spike_before = None, lookup_series_stepsize = 3, min_items = 10):
    pdf2 = pd.DataFrame(dict(wni_values = wni_values, spike = spike_times))
    if spike_before is not None:
        pdf = pdf2[~spike_before]
    else:
        pdf = pdf2

    bins = range(int(np.floor(pdf['wni_values'].min())), int(np.floor(pdf['wni_values'].max())) + lookup_series_stepsize + 1, lookup_series_stepsize)
    pdf['bin'] = pd.cut(pdf['wni_values'], bins, include_lowest=True)

    helper_df = pd.DataFrame()

    edges = []
    spike_prob = []
    n_items = []
    spikes = []
    for b in set(pdf['bin']):
        b_pdf = pdf[pdf['bin'] == b]
        n_items.append(len(b_pdf))
        edges.append(int(b.split(',')[0][1:]))
        spike_prob.append(np.mean(b_pdf['spike']))
        spikes.append(list(b_pdf['spike']))

    helper_df['bin_start'] = edges
    helper_df['spike_prob'] = spike_prob
    helper_df['spikes'] = spikes
    helper_df['n_items'] = n_items
    helper_df.sort_values('bin_start', inplace = True)

    bin_groupings = group_bins(helper_df, min_items = min_items)
    new_df = pd.DataFrame(columns = ['bin_start', 'spike_prob'], index = range(len(bin_groupings)))

    bin_starts = []
    spike_probs = []
    for g in bin_groupings:
        if len(g) == 1:
            bin_starts.append(helper_df.iloc[g[0], 0])
            spike_probs.append(helper_df.iloc[g[0], 1])
        else:
            bin_starts.append(helper_df.iloc[g[0], 0])
            spikes = []
            for row in g:
                spikes.append(helper_df.iloc[row, 2])

            spikes = [i for sublist in spikes for i in sublist]
            assert len(spikes) >= min_items
            spike_probs.append(np.mean(spikes))


    points = []
    for n, b in enumerate(bin_starts):
        midpoint = b + (lookup_series_stepsize * len(bin_groupings[n]))/2.
        points.append((midpoint, spike_probs[n]))

    points = sorted(points)
    x_points = [p[0] for p in points]
    y_points = [p[1] for p in points]

    diff = max(x_points) - min(x_points)
    index = np.arange(int(min(x_points) - 0.3 * diff), int(max(x_points) + 0.3 * diff))

    LUT = pd.Series(index = index)
    for i in index:
        LUT[i] = linear_interpolation_between_pairs(x_points,y_points, i, top = 'max')

    return LUT

from model_data_base.analyze.spike_detection import spike_in_interval

outdir = BASEDIR + 'precalculated_WNI_values/'

rm_dict[model]['nonlinearity'] = nonlinearity

outdir = BASEDIR + 'precalculated_WNI_values/'

def fetch_ISI_WNI_data(model, rm_inference_timepoint):
    st = pd.concat([pd.read_csv(BASEDIR + "model_{}__spike_times_MCM__{}.csv".format(model, whisker)) for whisker in whiskers])
    # is there a spike in the estimation period?
    st['spike_in_interval'] = spike_in_interval(st, rm_inference_timepoint, rm_inference_timepoint + 1)
    # how long ago was the most recent spike?
    st[st>rm_inference_timepoint] = np.NaN
    st['ISI'] = st.max(axis=1) - rm_inference_timepoint

    # collect WNI values and add them to the dataframe
    wni = pd.DataFrame()
    for whisker in whiskers:
        # adapt according to where you saved (this loads the provided precalculated WNI values for C2 whisker)
        wni_path = 'precalculated_WNI_values__model_1__{}'.format(whisker)
#         wni_path = outdir + '{}__{}.csv'.format(model, whisker)
        wni = pd.concat([wni, pd.read_csv(wni_path)])

    st['wni_value'] = np.asarray(wni.loc[:, str(rm_inference_timepoint)])


    spike_df = st[st.spike_in_interval & (st.ISI > -80) & (st.ISI < 0)]
    no_spike_df = st[~st.spike_in_interval & (st.ISI > -80) & (st.ISI < 0)]

    spike_isi = spike_df['ISI']
    spike_wni = spike_df['wni_value']

    no_spike_isi = no_spike_df['ISI']
    no_spike_wni = no_spike_df['wni_value']

    return spike_isi, spike_wni, no_spike_isi, no_spike_wni, st

def get_isi_boundary_95(spike_isi, spike_wni):
    def signchange(x,y):
        if x / abs(x) == y / abs(y):
            return False
        else:
            return True

    def linear_interpolation_between_pairs(X,Y, x):
        if x > max(X):
            result = np.inf
        elif x < min(X):
            result = min(Y)
        elif float(x) in X: # don't need to interpolate, causes assertion error
            result = Y[X.index(x)]
        else:
            pair = [lv for lv in range(len(X)-1) if signchange(X[lv]-x, X[lv+1]-x)]
    #         X[lv] <= x < X[lv + 1]]
            assert(len(pair) == 1)
            pair = pair[0]
            m = (Y[pair+1]-Y[pair]) / (X[pair+1]-X[pair])
            c = Y[pair]-X[pair]*m

            result = m*x+c
        return result

    l = zip(spike_isi, spike_wni)
    if len(l) < 2:
        return I.np.nan
    l = sorted(l, key = lambda x: x[1]) # sort by WNI values
    to_drop = len(spike_wni) / 20 # drop lowest 5%
    l = l[to_drop:]

    l = sorted(l, key = lambda x: x[0], reverse = True) # sort by ISI again

    points = []
    points.append([l[0][0], l[0][1]])
    y_min = l[0][1]
    for x, y in l:
        if y < y_min:
            points.append([x, y])
            y_min = y

    x_points = [p[0] for p in points]
    y_points = [p[1] for p in points]

    ISI_boundary = pd.Series(index = range(-80, 0))

    for ISI in ISI_boundary.index:
        ISI_boundary[ISI] = linear_interpolation_between_pairs(x_points, y_points, ISI)

    return ISI_boundary

spike_isi, spike_wni, no_spike_isi, no_spike_wni, st = fetch_ISI_WNI_data(1, rm_inference_timepoint)

ISI_penalty = get_isi_boundary_95(spike_isi, spike_wni)

def plot_ISI_cloud(ax, spike_isi, spike_wni, no_spike_isi, no_spike_wni, s = 2):
    ax.scatter(no_spike_isi, no_spike_wni, c = 'k', s = s, rasterized = True)
    ax.scatter(spike_isi, spike_wni, c = 'r', s = s, rasterized = True)
    ax.legend(['no spike', 'spike'])
    ax.set_xlim(-85, 5)
    ax.set_ylim(-30, 120)
    ax.set_xlabel('inter spike interval (ms)')
    ax.set_ylabel('weighted net input')

# when using the ISI-dependent penalty for simulations, you should subtract the baseline value, i.e.
ISI_penalty_simulations = ISI_penalty - min(ISI_penalty)

rm_dict[model]['ISI_penalty'] = ISI_penalty

import json
import seaborn as sns

# load reduced models which we provide, or use ones that you generated above
rm_dict = {}
for model in range(1,8):
    with open(BASEDIR + "model_{}__reduced_model_parameters.txt".format(model), "r") as fp:
        rm = json.load(fp)
    rm_dict[model] = rm

class ReducedModel:
    def __init__(self, kernel_dict, LUT, ISI_penalty, spatial_bin_size = 50):
        self.kernel_dict = kernel_dict
        self.LUT = LUT
        self.ISI_penalty = ISI_penalty
        self.spatial_bin_size = spatial_bin_size

    def run(self, SAexc, SAinh, return_WNI = False):
        '''Apply the reduced model to synaptic input to get a list containing output spike times.
        SAexc, SAinh: numpy arrays containing spatiotemporally binned synaptic inputs, shape (spatial bins, trials, temporal bins)
        return_WNI: bool
            full_trace returns a dataframe containing WNI values at all timepoints along with spike times.
            False returns spike times dataframe only.'''
        s_exc = self.kernel_dict['s_exc']
        s_inh = self.kernel_dict['s_inh']
        t_exc = self.kernel_dict['t_exc']
        t_inh = self.kernel_dict['t_inh']

        assert SAexc.shape == SAinh.shape

        LUT = self.LUT

        WNI_boundary = self.ISI_penalty

        spatial_bins = SAexc.shape[0]
        n_trials = SAexc.shape[1]
        timebins = SAexc.shape[2]

        WNI_df = pd.DataFrame(index = range(n_trials), columns = range(timebins))

        SAinh_cumulative = np.zeros((n_trials, timebins))
        SAexc_cumulative = np.zeros((n_trials, timebins))

        spike_times_df = pd.DataFrame(index = range(n_trials))

        for timebin in range(timebins):
            ## get excitatory and inhibitory input, spatially binned for the CURRENT timebin
            SAexc_timebin = np.ndarray((spatial_bins, n_trials, 1))
            for dist in range(spatial_bins):
                for cell in range(n_trials):
                    SAexc_timebin[dist][cell] = SAexc[dist][cell][timebin]

            SAinh_timebin = np.ndarray((spatial_bins, n_trials, 1))
            for dist in range(spatial_bins):
                for cell in range(n_trials):
                    SAinh_timebin[dist][cell] = SAinh[dist][cell][timebin]

            ## apply spatial kernel to the current timebin
            SAexc_timebin = sum([o*s for o, s in zip(SAexc_timebin, s_exc)])
            SAinh_timebin = sum([o*s for o, s in zip(SAinh_timebin, s_inh)])



            for cell in range(n_trials):
                SAinh_cumulative[cell][timebin] = SAinh_timebin[cell]
                SAexc_cumulative[cell][timebin] = SAexc_timebin[cell]



            ## apply temporal kernel
            spikes = []
            for cell in range(n_trials):
                if timebin - 80 >= 0:
                    SAexc_window = SAexc_cumulative[cell][timebin-79:timebin+1]
                    SAinh_window = SAinh_cumulative[cell][timebin-79:timebin+1]
                else:
                    SAexc_window = SAexc_cumulative[cell][0:timebin+1]
                    SAinh_window = SAinh_cumulative[cell][0:timebin+1]


                SAexc_window = sum([o*s for o, s in zip(SAexc_window, t_exc[-len(SAexc_window):])])
                SAinh_window = sum([o*s for o, s in zip(SAinh_window, t_inh[-len(SAinh_window):])])

                ## get weighted net input for each cell
                WNI = SAexc_window + SAinh_window
                WNI_df.iloc[cell, timebin] = WNI

                # check the last 80 bins for a spike, apply WNI penalty if there was
                min_index = 0 if timebin < 80 else timebin-80
                if sum(spike_times_df.iloc[cell, min_index:timebin]) > 0: #if there was a spike in the last 80 ms
                    last_spike_interval = 80 - np.where(spike_times_df.iloc[cell, min_index:timebin] == 1)[0][-1]
                    penalty = WNI_boundary[-last_spike_interval]
                    WNI -= penalty


                ## get spike probability from WNI
                if WNI > LUT.index.max():
                    spiking_probability = LUT[LUT.index.max()]
                elif WNI < LUT.index.min():
                    spiking_probability = LUT[LUT.index.min()]
                else:
                    spiking_probability = LUT[np.round(WNI)]

                ## spike or not?
                if spiking_probability > np.random.uniform():
                    spikes.append(1)
                else:
                    spikes.append(0)

            spike_times_df[str(timebin)] = spikes

        if return_WNI:
            return spike_times_df, WNI_df
        else:
            return spike_times_df

model = 1
filters = rm_dict[model]['spatiotemporal_filters']
nonlinearity = pd.Series(rm_dict[model]['nonlinearity'])
nonlinearity.index = nonlinearity.index.astype(int)
nonlinearity = nonlinearity.sort_index()

penalty = pd.Series(rm_dict[model]['ISI_penalty'])
penalty.index = penalty.index.astype(int)
penalty = penalty.sort_index()

rm = ReducedModel(filters, nonlinearity, penalty)

sa = np.load(BASEDIR + "model_{}__binned_synapse_activations_spatiotemporal_EI__{}.npz".format(model, whisker))
sa = sa['arr_0']
SAexc = sa[0,:,:,:]
SAinh = sa[1,:,:,:]

spike_times = rm.run(SAexc[:, 4000:6000, :], SAinh[:, 4000:6000, :]) # select the same trials as we visualised for the MCM

# reformat the spike times so that instead of 1 or 0 for each timebin per cell, each cell has a list of its spike times
st_reformat = pd.DataFrame(index = range(len(spike_times.index)), columns = range(max(spike_times.sum(axis = 1))))
for cell in spike_times.index:
    nspikes = 0
    for timebin in spike_times.columns:
        if spike_times.iloc[cell, int(timebin)] == 1:
            st_reformat.iloc[cell, nspikes] = int(timebin)
            nspikes += 1