# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python (sbi)
#     language: python
#     name: sbi
# ---

# +
# %load_ext autoreload
# %autoreload 2

import numpy as np
import os
import pickle
import matplotlib.pyplot as plt
import seaborn as sns
import itertools

import torch
torch.manual_seed(0)

import tensorflow as tf
from neural_circuits.LRRNN import tf_num_params, torch_num_params

import warnings
warnings.filterwarnings('ignore')

"""RNN stable amplification."""
# -

# ## Supp Fig 2

# +
colors = sns.color_palette()
Ns = [2, 5, 10, 25, 50, 100, 250]
epi_num_params = []
snpe_num_params = []
for N in Ns:
    epi_num_params.append(tf_num_params(N))
    snpe_num_params.append(torch_num_params(N))
    
yticks = [0, 1e5, 2e5, 3e5, 4e5, 5e5]
yticklabels = ["%dk" % (y/1000) for y in yticks]
    
fig, ax = plt.subplots(1,1)
plt.plot(Ns, epi_num_params, '-o', c=colors[0], label="EPI", alpha=0.7)
plt.plot(Ns, snpe_num_params, '-o', c=colors[1], label="SNPE", alpha=0.7)
plt.legend()
ax.set_yticks(yticks)
ax.set_yticklabels(yticklabels)
plt.show()
# -

#

StarcoderdataJupyterScriptsDedupFiltered