<jupyter_start><jupyter_text>Training networks of PRC neurons to memorize random inputThis notebook contains code for Fig. 7 of [<NAME> _et al_. (2021)](https://doi.org/10.1101/2021.03.25.437091).To reproduce these findings, first run `memorize.py`.<jupyter_code>import os
from copy import deepcopy

import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.gridspec as gs
from matplotlib.patheffects import Normal, Stroke
import numpy as np
import seaborn as sns
from ezephys import pltools

import memorize<jupyter_output><empty_output><jupyter_text>Illustrate training of one network<jupyter_code>FIGURE_WIDTH = 7

def savefig(fname, **pltargs):
    plt.savefig(fname + '.png', dpi=300, bbox_inches='tight', **pltargs)
    plt.savefig(fname + '.svg', dpi=300, bbox_inches='tight', **pltargs)
memorize.torch.manual_seed(42)

DEMO_NET = 'BAP'

demo_net_before_train = memorize.get_networks()[DEMO_NET]
demo_net_after_train = deepcopy(demo_net_before_train)
demo_optimizer = memorize.get_optimizers({DEMO_NET: demo_net_after_train})[DEMO_NET]
demo_x, demo_y = memorize.generate_data()

demo_optimizer.optimize(demo_x, demo_y, memorize.EPOCHS, progress_bar='notebook')
def tensor_to_events(tensor):
    return [np.nonzero(tensor[i, :]).flatten() for i in range(tensor.shape[0])]

def _visual_spike(voltage, spikes, spike_height=50):
    assert voltage.shape == spikes.shape
    voltage = deepcopy(voltage)
    voltage[spikes > 0.0] = spike_height
    return voltage

def plot_voltage_with_spikes(recorder, example_nb: int, unit_nb: int, ax=None, **pltargs):
    if ax is None:
        ax = plt.gca()
        
    ax.plot(
        _visual_spike(
            recorder.recorded['somatic_subunit.linear'][example_nb, :, unit_nb].detach().numpy(),
            recorder.recorded['output'][example_nb, :, unit_nb].detach().numpy()
        ),
        **pltargs
    )

y_palette = {0: 'xkcd:violet', 1: 'xkcd:forest green'}
spec_outer = gs.GridSpec(1, 3)
spec_hidden = gs.GridSpecFromSubplotSpec(demo_net_before_train.nb_units_by_layer[1], 1, spec_outer[:, 1])
spec_output = gs.GridSpecFromSubplotSpec(demo_net_before_train.nb_units_by_layer[2], 1, spec_outer[:, 2])

plt.figure(figsize=(FIGURE_WIDTH, 2), dpi=120)

plt.subplot(spec_outer[:, 0])
plt.title('Input layer')
plt.eventplot(tensor_to_events(demo_x[0, ...].T), linelengths=4, color='black')
plt.ylabel('Input units')
pltools.hide_border('trb', trim=True)

_, demo_activity = demo_net_before_train.run_snn(demo_x)

for i in range(demo_net_before_train.nb_units_by_layer[1]):
    plt.subplot(spec_hidden[i, :])
    plot_voltage_with_spikes(demo_activity['l1'], 0, i, color='black', clip_on=False)
    plt.ylim(-23, 52)
    
    if i == 0:
        plt.title('Hidden layer')
    
    pltools.hide_border()
    
for i in range(demo_net_before_train.nb_units_by_layer[2]):
    plt.subplot(spec_output[i, :])
    plt.plot(
        demo_activity['l2'].recorded['output'][0, :, i].detach().numpy(),
        color=y_palette[i],
        clip_on=False
    )
    plt.ylim(-12, 12)
    
    if i == 0:
        plt.title('Output layer')
    
    pltools.hide_border()

savefig(os.path.join('..', 'data', 'single_example_BAP'))

del spec_outer, spec_hidden, spec_output, i
def predicted_label(recording, example_nb):
    return recording.recorded['output'][example_nb, ...].max(axis=0)[0].max(axis=0)[1]

def annotate_correctness(recording, labels, example_nb, ax=None):
    if ax is None:
        ax = plt.gca()
        
    if predicted_label(demo_activity['l2'], example_nb) == demo_y[i]:
        mark, color = '✓', 'green'
    else:
        mark, color = '✗', 'red'
        
    text = ax.text(1, 1, mark, color=color, transform=ax.transAxes, ha='center', va='top')
    text.set_path_effects([Stroke(linewidth=2, foreground='white'), Normal()])
EXAMPLE_GRID_SHAPE = (3, 3)

spec_outer = gs.GridSpec(1, 4, hspace=0.2)
spec_input = gs.GridSpecFromSubplotSpec(*EXAMPLE_GRID_SHAPE, spec_outer[:, 0])
spec_before = gs.GridSpecFromSubplotSpec(*EXAMPLE_GRID_SHAPE, spec_outer[:, 1])
spec_loss = gs.GridSpecFromSubplotSpec(2, 2, spec_outer[:, 2], width_ratios=[0.4, 1])
spec_after = gs.GridSpecFromSubplotSpec(*EXAMPLE_GRID_SHAPE, spec_outer[:, 3])

plt.figure(figsize=(FIGURE_WIDTH, 2), dpi=120)

# Show input

for i in range(np.prod(EXAMPLE_GRID_SHAPE)):
    if i == 0:
        ax0 = plt.subplot(spec_input[i])
    else:
        plt.subplot(spec_input[i], sharey=ax0)
    plt.eventplot(tensor_to_events(demo_x[i, ...].T), linelengths=4, color=y_palette[int(demo_y[i])])
    pltools.hide_border()

    if i == EXAMPLE_GRID_SHAPE[0] // 2:
        plt.title('Input')

del ax0

# Show output before training

_, demo_activity = demo_net_before_train.run_snn(demo_x)

for i in range(np.prod(EXAMPLE_GRID_SHAPE)):
    if i == 0:
        ax0 = plt.subplot(spec_before[i])
    else:
        plt.subplot(spec_before[i], sharey=ax0)
    
    for j in range(demo_net_before_train.nb_units_by_layer[2]):
        plt.plot(
            demo_activity['l2'].recorded['output'][i, :, j].detach().numpy(),
            color=y_palette[j],
            clip_on=False
        )
    pltools.hide_border()
    
    annotate_correctness(demo_activity['l2'], demo_y, i)
    
    if i == EXAMPLE_GRID_SHAPE[0] // 2:
        plt.title('Before training')
    
# Show loss and accuracy

plt.subplot(spec_loss[0, 1])
plt.plot(demo_optimizer.loss_history, 'k-')
plt.xticks([0, 300, 600], ['', '', ''])
plt.ylim(0, plt.ylim()[1])
plt.ylabel('Loss')
pltools.hide_border('tr')

plt.subplot(spec_loss[1, 1])
plt.plot(demo_optimizer.accuracy_history, 'k-')
plt.xticks([0, 300, 600])
plt.ylim(0.4, 1)
plt.ylabel('Accuracy')
plt.xlabel('Epoch')
pltools.hide_border('tr')
    
# Show output after training

_, demo_activity = demo_net_after_train.run_snn(demo_x)


for i in range(np.prod(EXAMPLE_GRID_SHAPE)):
    if i == 0:
        ax0 = plt.subplot(spec_after[i])
    else:
        plt.subplot(spec_after[i], sharey=ax0)
    
    for j in range(demo_net_after_train.nb_units_by_layer[2]):
        plt.plot(
            demo_activity['l2'].recorded['output'][i, :, j].detach().numpy(),
            color=y_palette[j],
            clip_on=False
        )
    pltools.hide_border()
    
    if i == EXAMPLE_GRID_SHAPE[0] // 2:
        plt.title('After training')
        
    annotate_correctness(demo_activity['l2'], demo_y, i)
    
savefig(os.path.join('..', 'data', 'multi_example_BAP'))<jupyter_output><empty_output><jupyter_text>Performance comparisonRun `memorize.py` to train five different PRC models several times with different seeds. The loss and accuracy during training are saved to CSV files that we'll inspect here.<jupyter_code>DATA_FILE_PREFIX = os.path.join('..', 'data', 'memorization_training_results_')

performance = []

for i in range(10):
    single_seed_performance = pd.read_csv(DATA_FILE_PREFIX + str(i) + '.csv')
    single_seed_performance['seed'] = i

    performance.append(single_seed_performance)

performance = pd.concat(performance).reset_index(drop=True)
performance.head()
def performance_bandplot(model_name, metric_loss=True, ax=None, **pltargs):
    if ax is None:
        ax = plt.gca()

    if metric_loss:
        metric = 'loss'
    else:
        metric = 'accuracy'

    alpha = min(pltargs.pop('alpha', 1), 0.8)

    this_model = performance.query('model_name == @model_name').sort_values(
        'epoch'
    )
    if 'label' not in pltargs:
        pltargs['label'] = model_name

    this_std = (
        this_model.groupby(['model_name', 'epoch']).std().sort_values('epoch')
    )
    this_mean = (
        this_model.groupby(['model_name', 'epoch']).mean().sort_values('epoch')
    )

    label = pltargs.pop('label', None)
    maincolor = pltargs.pop('color', None)
    ax.fill_between(
        this_mean.index.get_level_values('epoch'),
        this_mean[metric] - this_std[metric],
        this_mean[metric] + this_std[metric],
        alpha=alpha,
        facecolor=maincolor,
        edgecolor='none',
        **pltargs
    )
    ax.plot(
        this_mean.index.get_level_values('epoch'),
        this_mean[metric],
        label=label,
        color=maincolor,
        **pltargs
    )
palette = {
    'One compartment': 'gray',
    'No BAP': 'xkcd:ocean',
    'BAP': 'xkcd:cherry',
    'Parallel subunits, no BAP': 'xkcd:iris',
    'Parallel subunits + BAP (full PRC model)': 'xkcd:blood orange'
}
spec_outer = gs.GridSpec(1, 3, width_ratios=(0.5, 1, 1))

plt.figure(figsize=(FIGURE_WIDTH, 2), dpi=120)

plt.subplot(spec_outer[:, 1])
plt.axhline(0.5, color='k', ls='--', dashes=(10, 5), lw=0.7, zorder=-1)
for model_name, colour in palette.items():
    this_mean = performance_bandplot(
        model_name, metric_loss=False, color=colour, alpha=0.3
    )
plt.legend(loc='upper right', bbox_to_anchor=(-0.2, 1))
plt.xlabel('Epoch')
plt.ylabel('Accuracy')
plt.ylim(0.36, 1.05)
pltools.hide_border('tr')

plt.subplot(spec_outer[:, 2])
plt.axhline(0.5, color='k', ls='--', dashes=(10, 5), lw=0.7, zorder=-1)
box_data = performance.query('epoch in [0, 1999]').copy()
box_data.loc[:, 'Epoch'] = box_data.epoch.astype(str)
box_data.replace('0', 'Before training', inplace=True)
box_data.replace('1999', 'After training', inplace=True)
sns.boxplot(
    x='model_name',
    y='accuracy',
    hue='Epoch',
    data=box_data,
)
plt.legend(title='', loc='upper left', bbox_to_anchor=(1, 1))
plt.xticks(rotation=45, ha='right')
plt.ylabel('')
plt.ylim(0.36, 1.05)
pltools.hide_border('tr')
plt.gca().set_yticklabels([])

savefig(os.path.join('..', 'data', 'performance_comparison'))
plt.show()<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.15515620102392644, 'fraction_numerical': 0.021836798662626684, 'mean_word_length': 4.091489361702128, 'pattern_counts': {'":': 0, '<': 14, '<?xml version=': 0, '>': 15, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 1, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '3', 'dir': 'jupyter-structured-clean-dedup', 'id': '666265', 'license': 'CC-BY-4.0', 'n_tokens_mistral': 3778, 'n_tokens_neox': 3590, 'n_words': 634, 'path': 'training/memorize-seed-results.ipynb', 'repo_name': 'nauralcodinglab/linear-nonlinear-dendrites', 'size': '9577'}