<jupyter_start><jupyter_text>4-neuron model<jupyter_code>neuron_inds = {'E':0, 'P':1, 'S':2, 'V':3}
V_INIT = tf.constant(-65.*np.ones((1,4,1)), dtype=np.float32)


k = 0.3
n = 2.
v_rest = -70.

dt = 0.005

N = 5
T = 100

def f_r(v):
    return k*(tf.nn.relu(v-v_rest)**n)

# Dim is [M,N,|r|,T]
def SSSN_sim_traj(h):
    h = h[:,None,:,None]

    W = load_W()
    sigma_eps = .2*np.array([1., 0.5, 0.5, 0.5])
    tau = np.array([0.02, 0.01, 0.01, 0.01])
    tau_noise = np.array([0.05, 0.05, 0.05, 0.05])

    W = W[None,:,:,:]
    sigma_eps = sigma_eps[None,None,:,None]
    tau = tau[None,None,:,None]
    tau_noise = tau_noise[None,None,:,None]
   
    _v_shape = tf.ones((h.shape[0], N, 4, 1), dtype=tf.float32)
    v_init = _v_shape*V_INIT
    eps_init = 0.*_v_shape
    y_init = tf.concat((v_init, eps_init), axis=2)
    
    def f(y):
        v = y[:,:,:4,:]
        eps = y[:,:,4:,:]
        B = tf.random.normal(eps.shape, 0., np.sqrt(dt))
       
        dv = (-v + v_rest + h + eps + tf.matmul(W, f_r(v))) / tau
        deps = (-eps + (np.sqrt(2.*tau_noise)*sigma_eps*B/dt)) / tau_noise
        
        return tf.concat((dv, deps), axis=2)
        
    v_t = euler_sim_stoch_traj(f, y_init, dt, T)
    
    return v_t

def SSSN_sim(h):
    h = h[:,None,:,None]

    W = load_W()
    sigma_eps = .2*np.array([1., 0.5, 0.5, 0.5])
    tau = np.array([0.02, 0.01, 0.01, 0.01])
    tau_noise = np.array([0.05, 0.05, 0.05, 0.05])

    W = W[None,:,:,:]
    sigma_eps = sigma_eps[None,None,:,None]
    tau = tau[None,None,:,None]
    tau_noise = tau_noise[None,None,:,None]
   
    _v_shape = tf.ones((h.shape[0], N, 4, 1), dtype=tf.float32)
    v_init = _v_shape*V_INIT
    eps_init = 0.*_v_shape
    y_init = tf.concat((v_init, eps_init), axis=2)
    
    def f(y):
        v = y[:,:,:4,:]
        eps = y[:,:,4:,:]
        B = tf.random.normal(eps.shape, 0., np.sqrt(dt))
       
        dv = (-v + v_rest + h + eps + tf.matmul(W, f_r(v))) / tau
        deps = (-eps + (np.sqrt(2.*tau_noise)*sigma_eps*B/dt)) / tau_noise
        
        return tf.concat((dv, deps), axis=2)
        
    v_ss = euler_sim_stoch(f, y_init, dt, T)
    return v_ss

#def SSSN_stats(h):
#    v_t = SSSN_sim_traj(h)[:,:,:4,:]
#    r_t = v_t
#    avg_r = tf.reduce_mean(r_t[:,:,:,(T//2):], axis=3)
#   std_r = tf.math.reduce_std(r_t[:,:,:,(T//2):], axis=3)
#    
#    mean_avg_r = tf.reduce_mean(avg_r, axis=1)
#    mean_std_r = tf.reduce_mean(std_r, axis=1)
#   std_std_r = tf.math.reduce_std(std_r, axis=1)
#    
#    return tf.concat((mean_avg_r, mean_std_r, std_std_r), axis=1)
lb_h = 0.*np.ones((4,))
ub_h = 25.*np.array([1., 1., 1., 1.])

lb_dh = -10.
ub_dh = 10.

h = Parameter("h", 4, lb=lb_h, ub=ub_h)
dh = Parameter("dh", 1, lb=lb_dh, ub=ub_dh)
parameters = [h, dh]

diff_prod_mean = -.25
diff_sum_mean = 0.
def SV_flip(h, dh):
    dh_pattern = tf.constant(np.array([[1., 1., 0., 0.]], dtype=np.float32))
    dh = dh*dh_pattern
   
    ss1 = tf.reduce_mean(f_r(SSSN_sim(h)[:,:,:4]), axis=1)
    ss2 = tf.reduce_mean(f_r(SSSN_sim(h+dh)[:,:,:4]), axis=1)
    
    diff1 = (ss1[:,2]-ss1[:,3]) / tf.norm(ss1, axis=1, keepdims=False)
    diff2 = (ss2[:,2]-ss2[:,3]) / tf.norm(ss2, axis=1, keepdims=False)
    diff_prod = diff1*diff2
    diff_sum = diff1+diff2
    T_x = tf.stack((diff_prod, 
                    diff_sum, 
                    (diff_prod - diff_prod_mean) ** 2, 
                    (diff_sum-diff_sum_mean) ** 2), axis=1)

    return T_x
# Define model
name = "V1Circuit"
model = Model(name, parameters)
model.set_eps(SV_flip)
D = sum([param.D for param in model.parameters])
lb = np.concatenate([param.lb for param in model.parameters])
ub = np.concatenate([param.ub for param in model.parameters])

M = 1000
num_keep = 100

mean = [diff_prod_mean, 0.]
std = [0.05, 0.2]
def accept_inds(T_x, mean ,std):
    return np.logical_and(np.logical_and(mean[0]-2*std[0] < T_x[:,0], T_x[:,0] < mean[0]+2*std[0]),
                          np.logical_and(mean[1]-2*std[1] < T_x[:,1], T_x[:,1] < mean[1]+2*std[1]))

num_found = 0
z = None
T_x = None
x = []
while (num_found < num_keep):
    _z = np.zeros((M, D))
    for j in range(D):
        _z[:,j] = np.random.uniform(lb[j], ub[j], (M,))
    _T_x = model.eps(_z).numpy()

    inds = accept_inds(_T_x, mean, std)
    _z = _z[inds, :]
    _T_x = _T_x[inds, :]
    num_found += _z.shape[0]

    if (z is None):
        z = _z
        T_x = _T_x
    else:
        z = np.concatenate((z, _z), axis=0)
        T_x = np.concatenate((T_x, _T_x), axis=0)
    print('%d/%d\r' % (num_found, num_keep), end='')
    x.append(num_found)

num_found = z.shape[0]

# reduce to num_keep
z = z[:num_keep]
T_x = T_x[:num_keep]
num_plot = 500
fontsize = 14
z_labels = [r'$h_E$', r'$h_P$', r'$h_S$', r'$h_V$', r'$dh$']
fig = pairplot(z[:num_plot], np.arange(5), z_labels,
               fontsize=22,
               lb=lb, ub=ub,
               c = T_x[:,0])
plt.savefig(figdir + 'SV_flip_ABC_z.png')
               #c_label='prod')
def plot_responses(z):
    h = z[:4]
    dh = z[4]
    gamma = np.concatenate((np.array([0., 1.]), np.linspace(-.1, 3., 101)), axis=0)
    dh_pattern = tf.constant(np.array([[1., 1., 0., 0.]], dtype=np.float32))
    _h = h + dh*gamma[:,None]*dh_pattern
    v_t = SSSN_sim_traj(_h).numpy()
    #T_x = SSN_stats(_h)
    
    neurons = ['E', 'P', 'S', 'V']
    """plt.figure()
    for i in range(4):
        plt.errorbar(gamma[2:], T_x[2:,i], T_x[2:,i+4], c=colors[i])
    plt.legend(neurons)
    plt.ylabel('means')
    plt.show()
    
    plt.figure()
    for i in range(4):
        plt.errorbar(gamma[2:], T_x[2:,i+4], T_x[2:,i+8], c=colors[i])
    plt.ylabel('noise')
    plt.legend(neurons)
    plt.show()"""


    
    t = dt*np.arange(T+1)
    
    for gamma_ind in [0, 1]:
        plt.figure()
        for neuron in neurons:
            ind = neuron_inds[neuron]
            plt.plot(t, v_t[gamma_ind,0,ind,:], c=colors[ind], lw=3)
        plt.plot([t[0], t[-1]], [-70, -70], 'k--', lw=3)
        h_gamma = _h[gamma_ind]
        plt.title('h = [%.2f %.2f %.2f %.2f]mV' % tuple([h_gamma[j] for j in range(4)]))
        plt.xlabel('time')
        plt.ylabel(r'$v_\alpha(t)$')
        plt.legend(neurons, title=r'$\alpha$')
        plt.show()
    
        
    fig, axs = plt.subplots(1,2,figsize=(12,5))
    for neuron in neurons:
        ind = neuron_inds[neuron]
        axs[0].plot(gamma[2:], v_t[2:,0,ind,-1], c=colors[ind], lw=3)
        axs[1].plot(gamma[2:], f_r(v_t[2:,0,ind,-1]), c=colors[ind], lw=3)
    axs[0].set_ylabel('v')
    axs[1].set_ylabel('r')
    plt.show()
    return None

for i in range(4):
    plot_responses(z[i])<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.21215319662243667, 'fraction_numerical': 0.04221954161640531, 'mean_word_length': 2.620633187772926, 'pattern_counts': {'":': 0, '<': 10, '<?xml version=': 0, '>': 5, '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}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '124745', 'license': 'Apache-2.0', 'n_tokens_mistral': 3222, 'n_tokens_neox': 3047, 'n_words': 534, 'path': 'notebooks/SSSN/.ipynb_checkpoints/SSSN_explore1_old-checkpoint.ipynb', 'repo_name': 'cunningham-lab/epi', 'size': '6632'}