<commit_before>"""
Functions tailored to multi-trial neural data.
"""

import matplotlib.pyplot as plt
import numpy as np
from ipywidgets import interact
from tensortools.kruskal import _validate_kruskal
from tensorly import kruskal_to_tensor   
from jetpack import breathe, bars

DEFAULT_COLOR_CYCLE = [(0.4470588235294118, 0.6196078431372549, 0.807843137254902),
                       (1.0, 0.6196078431372549, 0.2901960784313726),
                       (0.403921568627451, 0.7490196078431373, 0.3607843137254902),
                       (0.9294117647058824, 0.4, 0.36470588235294116),
                       (0.6784313725490196, 0.5450980392156862, 0.788235294117647),
                       (0.6588235294117647, 0.47058823529411764, 0.43137254901960786),
                       (0.9294117647058824, 0.592156862745098, 0.792156862745098),
                       (0.6352941176470588, 0.6352941176470588, 0.6352941176470588),
                       (0.803921568627451, 0.8, 0.36470588235294116),
                       (0.42745098039215684, 0.8, 0.8549019607843137)]

def interact_reconstruction(data, model, avg_trial=False, condition=None, color_cycle=DEFAULT_COLOR_CYCLE, **kwargs):
    """Returns a function for interacive plotting of model reconstruction error
    """

    ndim, rank = _validate_kruskal(model)

    if ndim != 3 or data.ndim != 3:
        raise ValueError('Data and model must be a 3rd-order tensor.')

    N, T, K = data.shape
    tax = np.linspace(0,2,T)

    if avg_trial:
        Xest = kruskal_to_tensor(model)
        if condition is None:
            condition = np.ones(K, np.int)
        Cavg = [np.mean(model[2][condition == c], axis=0) for c in np.unique(condition)]
        Xavg = [np.mean(data[:, :, condition == c], axis=2) for c in np.unique(condition)]
        Xavg_est = [np.mean(Xest[:, :, condition == c], axis=2) for c in np.unique(condition)]

    # plotting function with trial averaging
    def _yes_trial_avg(neuron):
        fig, axes = plt.subplots(2, 2, gridspec_kw=dict(width_ratios=[3,1]), **kwargs)

        [axes[0,0].plot(tax, x[neuron], '-', color=c, lw=3, alpha=0.3) for x, c in zip(Xavg, color_cycle)]
        [axes[0,0].plot(tax, x[neuron], '--', color=c, lw=3) for x, c in zip(Xavg_est, color_cycle)]

        [axes[1,0].plot(tax, model[1]*cavg, '-', lw=3, color='c', alpha=0.5) for cavg, c in zip(Cavg, color_cycle)]
        axes[1,0].set_ylim(axes[0,0].get_ylim())
        # axes[1,0].set_title('latent factors', y=0.8)
        # axes[1,0].set_xlabel('time (s)')
        
        axes[0,1].bar(range(1,rank+1), model[0][neuron], align='center', color=(0.5,0.5,0.5))
        axes[0,1].axhline(0, color='k', lw=1)
        axes[0,1].set_xlim([0,rank+1])
        axes[0,1].set_xticks(range(1,rank+1))
        axes[0,1].set_title('neuron loadings')

        bars(np.array(Cavg), ax=axes[1,1], colors=color_cycle)
        axes[1,1].set_title('trial loadings')
        axes[1,1].axhline(0, color='k', lw=1)

        plt.tight_layout()
        [breathe(ax=ax) for ax in axes.ravel()]

    # plotting function without trial averaging
    def _no_trial_avg(neuron, trial):
        fig, axes = plt.subplots(2, 2, gridspec_kw=dict(width_ratios=[3,1]), **kwargs)
<commit_msg>minor<commit_after>"""
Functions tailored to multi-trial neural data.
"""

import matplotlib.pyplot as plt
import numpy as np
from ipywidgets import interact
from tensortools.kruskal import _validate_kruskal
from tensorly import kruskal_to_tensor   
from jetpack import breathe, bars

DEFAULT_COLOR_CYCLE = [(0.4470588235294118, 0.6196078431372549, 0.807843137254902),
                       (1.0, 0.6196078431372549, 0.2901960784313726),
                       (0.403921568627451, 0.7490196078431373, 0.3607843137254902),
                       (0.9294117647058824, 0.4, 0.36470588235294116),
                       (0.6784313725490196, 0.5450980392156862, 0.788235294117647),
                       (0.6588235294117647, 0.47058823529411764, 0.43137254901960786),
                       (0.9294117647058824, 0.592156862745098, 0.792156862745098),
                       (0.6352941176470588, 0.6352941176470588, 0.6352941176470588),
                       (0.803921568627451, 0.8, 0.36470588235294116),
                       (0.42745098039215684, 0.8, 0.8549019607843137)]

def interact_reconstruction(data, model, avg_trial=False, condition=None, color_cycle=DEFAULT_COLOR_CYCLE, **kwargs):
    """Returns a function for interacive plotting of model reconstruction error
    """

    ndim, rank = _validate_kruskal(model)

    if ndim != 3 or data.ndim != 3:
        raise ValueError('Data and model must be a 3rd-order tensor.')

    N, T, K = data.shape
    tax = np.linspace(0,2,T)

    if avg_trial:
        Xest = kruskal_to_tensor(model)
        if condition is None:
            condition = np.ones(K, np.int)
        Cavg = [np.mean(model[2][condition == c], axis=0) for c in np.unique(condition)]
        Xavg = [np.mean(data[:, :, condition == c], axis=2) for c in np.unique(condition)]
        Xavg_est = [np.mean(Xest[:, :, condition == c], axis=2) for c in np.unique(condition)]


    # plotting function with trial averaging
    def _yes_trial_avg(neuron):
        fig, axes = plt.subplots(2, 2, gridspec_kw=dict(width_ratios=[3,1]), **kwargs)

        [axes[0,0].plot(tax, x[neuron], '-', color=c, lw=3, alpha=0.3) for x, c in zip(Xavg, color_cycle)]
        [axes[0,0].plot(tax, x[neuron], '--', color=c, lw=3) for x, c in zip(Xavg_est, color_cycle)]
        axes[0,0].set_ylim((0, max([np.max(x) for x in Xavg])))

        [axes[1,0].plot(tax, model[0][neuron]*model[1]*cavg, '-', lw=3, color=c, alpha=0.5) for cavg, c in zip(Cavg, color_cycle)]
        axes[1,0].set_title('latent factors', y=1)
        axes[1,0].set_xlabel('time (s)')
        
        axes[0,1].bar(range(1,rank+1), model[0][neuron], align='center', color=(0.5,0.5,0.5))
        axes[0,1].axhline(0, color='k', lw=1)
        axes[0,1].set_xlim([0,rank+1])
        axes[0,1].set_xticks(range(1,rank+1))
        axes[0,1].set_title('neuron loadings')

        bars(np.array(Cavg), ax=axes[1,1], colors=color_cycle)
        axes[1,1].set_title('trial loadings')
        axes[1,1].axhline(0, color='k', lw=1)

        plt.tight_layout()
        [breathe(ax=ax) for ax in axes.ravel()]

    # plotting function without trial averaging
    def _no_trial_avg(neuron, trial):
        fig, axes = plt.subplots(2, 2, gridspec_kw=dict(width_ratios=[3,1]), **kwargs)

{'fraction_non_alphanumeric': 0.1474269819193324, 'fraction_numerical': 0.16148972338123938, 'mean_word_length': 2.965686274509804, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '728755', 'n_tokens_mistral': 3047, 'n_tokens_neox': 2395, 'n_words': 467}