<jupyter_start><jupyter_text>Build train and test sets<jupyter_code># generate training data
P_train = 100
x_lim = 4

x1 = np.linspace(-x_lim, x_lim, P_train)
x2 = np.linspace(-x_lim, x_lim, P_train)
X1_train, X2_train = np.meshgrid(x1, x2)
    
def func(x1, x2):
    y = 0.3*(x1**2 + x2**2)
    return y

Y_targ_train = func(X1_train, X2_train)
Y_ravelled = Y_targ_train.ravel()
X_train = np.vstack((X1_train.ravel(), X2_train.ravel())).T
Y_train = Y_ravelled[:, None]

T = 4 # simulation time
dt = 3e-03 # time step
t_span = np.arange(0, T, dt)
num_bins = t_span.size
buffer_bins = int(1/dt)
buffer_zeros = int(buffer_bins/2)
K = 2
x_sample = np.zeros((K, num_bins))

data_index = 1
print(X_train)
print(X_train[data_index, :][:, None])
x_sample[:, buffer_zeros:] = X_train[data_index, :][:, None]
print(x_sample)
print(x_sample.shape)
# create test data
P_test = 500
x_lim = 4
x1 = np.linspace(-x_lim, x_lim, P_test)
x2 = np.linspace(-x_lim, x_lim, P_test)
X1_test, X2_test = np.meshgrid(x1, x2)

Y_targ_test = func(X1_test, X2_test)
Y_ravelled_test = Y_targ_test.ravel()
X_test = np.vstack((X1_test.ravel(), X2_test.ravel())).T
Y_test = Y_ravelled_test[:, None]

# Plot a 3D surface
fig = plt.figure(figsize=(8, 9))
ax1 = fig.add_subplot(111, projection='3d')
ax1.set_xlabel('$x_1$')
ax1.set_ylabel('$x_2$')
ax1.set_zlabel('$y$')
# Plot a 3D surface
target_surface = ax1.plot_surface(X1_test, X2_test, Y_targ_test, alpha=0.5)
fig.suptitle('Target surface', fontsize=18)
plt.show()<jupyter_output><empty_output><jupyter_text>Network parameters<jupyter_code># setting up dimensions and initial parameters
M = 1
K = 2
N = 50
leak = 2

# initialize my gamma matrix
random_state = np.random.RandomState(seed=3)
D_weights_init = random_state.rand(M, N)
D_weights_init = D_weights_init / np.linalg.norm(D_weights_init, axis=0)
G_weights_init = D_weights_init.copy().T
F_weights_init = random_state.randn(K, N).T
omega_init = -G_weights_init @ D_weights_init
thresholds_init = 2*random_state.rand(N) - 1

F_weights_2_init = random_state.randn(N, N).T
G_weights_2_init = G_weights_init
omega_2_init = omega_init
thresholds_2_init = 2*random_state.rand(N) - 1<jupyter_output><empty_output><jupyter_text>Before learningWe run the network with the inital parameters and compute an average readout for each input sample. Thus, we can get the SNN surface before learning.Also, since the output dimension of the network is one, we can use a maxout function - refer to equation (10) of the paper - to get the network output (without discretization error). We also check which neuron contributes to the code for each input sample (active constraints) by looking at the argument that maximizes the maxout function.<jupyter_code>T = 4 # simulation time
dt = 3e-03 # time step
t_span = np.arange(0, T, dt)
num_bins = t_span.size
buffer_bins = int(1/dt)
buffer_zeros = int(buffer_bins/2)
x_sample = np.zeros((K, num_bins))

# initialize network parameters
D_weights = D_weights_init.copy()
G_weights = G_weights_init.copy()
F_weights = F_weights_init.copy()
omega = omega_init.copy()
thresholds = thresholds_init.copy()

G_weights_2 = G_weights_2_init.copy()
F_weights_2 = F_weights_2_init.copy()
omega_2 = omega_2_init.copy()
thresholds_2 = thresholds_2_init.copy()


y_readout = []

for data_index in range(X_train.shape[0]):
    x_sample[:, buffer_zeros:] = X_train[data_index, :][:, None]

    rates = snn_cvx.run_snn_trial(
        x_sample,
        F_weights,
        omega,
        thresholds,
        dt,
        leak,
    )
    
    y_readout += [np.copy(D_weights[0, :] @ rates)]
    
average_readouts = np.array(y_readout)[:, buffer_zeros + 500:].mean(axis=1)
# Plot a 3D surface
fig = plt.figure(figsize=(8, 9))
ax1 = fig.add_subplot(111, projection='3d')
ax1.set_xlabel('$x_1$')
ax1.set_ylabel('$x_2$')
ax1.set_zlabel('$y$')
# Plot a 3D surface
readout_surf_snn = ax1.plot_surface(X1_train, X2_train, average_readouts.reshape(Y_targ_train.shape), alpha=0.5)
fig.suptitle('SNN surface before learning', fontsize=18)
plt.show()
# plot contours and active-inactive neurons
active_neurons_init = np.zeros(X_test.shape[0]) * np.nan
y_predict_init = np.zeros(X_test.shape[0])
for i, x in enumerate(X_test):
    y_out_init, n_act = snn_cvx.run_maxout(x, F_weights_init, G_weights_init, thresholds_init)
    y_predict_init[i] = y_out_init
    active_neurons_init[i] = n_act
    
y_predict_init_reshaped = y_predict_init.reshape(Y_targ_test.shape)
active_neurons_init_reshaped = active_neurons_init.reshape(Y_targ_test.shape)

# make contour plots
zlim =14
ticks = [(i, i) for i in np.linspace(x_lim, -x_lim, 3)]
c_ticks = [(i, i) for i in np.linspace(0, zlim, 2)]

clim = (0, zlim)
bounds = (-x_lim, -x_lim, x_lim, x_lim)
nlevels = 40
cmap_neurons = 'glasbey_dark'
cmap_contour='gray'
alpha=1
    
img = hv.Image(y_predict_init_reshaped, kdims=['$x_1$', '$x_2$'], vdims='$y$', bounds=bounds).opts(
                                    cmap=cmap_contour, invert_zaxis=True, clims=clim)
levels = list(np.linspace(-zlim, zlim, nlevels))
img_contour_init = hv.operation.contours(img, group='Y', levels=levels).opts(
                                    xticks=3, yticks=3, colorbar=True, cmap=cmap_contour, 
                                    cbar_ticks=c_ticks, clim=clim, alpha=alpha, linewidth=1.5, 
                                    )

# show which neuron is active
img_nactive_init = hv.Image(active_neurons_init_reshaped, kdims=['$x_1$', '$x_2$'], vdims='$y$', bounds=bounds).opts(cmap=cmap_neurons, alpha=0.6)

# plot contours and active neurons
img_nactive_init*img_contour_init<jupyter_output><empty_output><jupyter_text>Train network parameters<jupyter_code>T = 4 # simulation time
dt = 3e-03 # time step
t_span = np.arange(0, T, dt)
num_bins = t_span.size
buffer_bins = int(1/dt)
buffer_zeros = int(buffer_bins/2)
x_sample = np.zeros((K, num_bins))

# initialize network parameters
D_weights = D_weights_init.copy()
G_weights = G_weights_init.copy()
F_weights = F_weights_init.copy()
omega = omega_init.copy()
thresholds = thresholds_init.copy()

G_weights_2 = G_weights_2_init.copy()
F_weights_2 = F_weights_2_init.copy()
omega_2 = omega_2_init.copy()
thresholds_2 = thresholds_2_init.copy()

# run supervised learning
alpha_thresh_init = 1e-03
alpha_F_init = 1e-03
leak_thresh = 0.

num_epochs = 100
thresholds_array_fit = np.zeros((N, num_epochs))
F_weights_array_fit = np.zeros((N, K, num_epochs))
decrease_learning_rate = True

for epoch in range(num_epochs):
    print ('iteration: ',epoch+1)
    data_index_list = np.arange(X_train.shape[0])
    np.random.shuffle(data_index_list)
    
    if decrease_learning_rate:
        alpha_thresh = alpha_thresh_init * np.exp(-0.0001 * (epoch + 1))
        alpha_F = alpha_F_init * np.exp(-0.0001 * (epoch + 1))

    else:
        alpha_thresh = alpha_thresh_init
        alpha_F = alpha_F_init
    
    for data_index in data_index_list:
        x_sample[:, buffer_zeros:] = X_train[data_index, :][:, None]
        y_sample = Y_train[data_index, :]

        thresholds, F_weights, thresholds_2, F_weights_2 = snn_cvx.update_weights_2(
            x_sample,
            y_sample,
            F_weights,F_weights_2,
            G_weights,G_weights_2,
            omega,omega_2,
            thresholds,thresholds_2,
            buffer_bins,
            dt,
            leak,
            leak_thresh,
            alpha_thresh,
            alpha_F,
            mu=0.,
            sigma_v=0.
        )
        
    thresholds_array_fit[:, epoch] = thresholds
    F_weights_array_fit[:, :, epoch] = F_weights<jupyter_output><empty_output><jupyter_text>After learning<jupyter_code># run snn with learnt parameters
x_sample = np.zeros((K, num_bins))

# call learnt parameters
F_weights_fit = F_weights_array_fit[:, :, -1]
thresholds_fit = thresholds_array_fit[:, -1]

y_readout = []

for data_index in range(X_train.shape[0]):
    x_sample[:, buffer_zeros:] = X_train[data_index, :][:, None]

    rates = snn_cvx.run_snn_trial(
        x_sample,
        F_weights_fit,
        omega,
        thresholds_fit,
        dt,
        leak,
    )
    
    y_readout += [np.copy(D_weights[0, :] @ rates)]
    
average_readouts_fit = np.array(y_readout)[:, buffer_zeros + 500:].mean(axis=1)
# Plot a 3D surface
fig = plt.figure(figsize=(8, 9))
ax1 = fig.add_subplot(111, projection='3d')
ax1.set_xlabel('$x_1$')
ax1.set_ylabel('$x_2$')
ax1.set_zlabel('$y$')
# Plot a 3D surface
readout_surf_snn = ax1.plot_surface(X1_train, X2_train, average_readouts_fit.reshape(Y_targ_train.shape), alpha=0.5)
fig.suptitle('SNN surface after learning', fontsize=18)
plt.show()
# plot contours and active-inactive neurons after learning parameters
active_neurons_fit = np.zeros(X_test.shape[0]) * np.nan
y_predict_fit = np.zeros(X_test.shape[0])
for i, x in enumerate(X_test):
    y_out_fit, n_act = snn_cvx.run_maxout(x, F_weights_fit, G_weights, thresholds_fit)
    y_predict_fit[i] = y_out_fit
    active_neurons_fit[i] = n_act
    
y_predict_fit_reshaped = y_predict_fit.reshape(Y_targ_test.shape)
active_neurons_fit_reshaped = active_neurons_fit.reshape(Y_targ_test.shape)

# make contour plots
zlim =14
ticks = [(i, i) for i in np.linspace(x_lim, -x_lim, 3)]
c_ticks = [(i, i) for i in np.linspace(0, zlim, 2)]

clim = (0, zlim)
bounds = (-x_lim, -x_lim, x_lim, x_lim)
nlevels = 40
cmap_neurons = 'glasbey_dark'
cmap_contour='gray'
alpha=1
    
img = hv.Image(y_predict_fit_reshaped, kdims=['$x_1$', '$x_2$'], vdims='$y$', bounds=bounds).opts(
                                    cmap=cmap_contour, invert_zaxis=True, clims=clim)
levels = list(np.linspace(-zlim, zlim, nlevels))
img_contour_fit = hv.operation.contours(img, group='Y', levels=levels).opts(
                                    xticks=3, yticks=3, colorbar=True, cmap=cmap_contour, 
                                    cbar_ticks=c_ticks, clim=clim, alpha=alpha, linewidth=1.5, 
                                    )

# show which neuron is active
img_nactive_fit = hv.Image(active_neurons_fit_reshaped, kdims=['$x_1$', '$x_2$'], vdims='$y$', bounds=bounds).opts(cmap=cmap_neurons, alpha=0.6)

# plot contours and active neurons
img_nactive_fit*img_contour_fit<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.12515976796775144, 'fraction_numerical': 0.023989774850063907, 'mean_word_length': 4.103863522328148, 'pattern_counts': {'":': 0, '<': 21, '<?xml version=': 0, '>': 21, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 8, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '5', 'dir': 'jupyter-structured-clean-dedup', 'id': '196316', 'license': 'MIT', 'n_tokens_mistral': 4171, 'n_tokens_neox': 3999, 'n_words': 772, 'path': 'experiment4_1/Convex_SNNs.ipynb', 'repo_name': 'worden1/ICTAI_snn', 'size': '10171'}