# Understanding spiking networks through convex optimization - Figure 3¶

## Imports¶

In [18]:

    import numpy as np
    import numba as nb
    import matplotlib.pyplot as plt
    import snn_cvx
    from mpl_toolkits.mplot3d import axes3d
    %matplotlib inline

In [14]:

    import holoviews as hv
    hv.extension('matplotlib')

## Build train and test sets¶

In [3]:

    # 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]

In [4]:

    # 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]

In [22]:

    # 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()

## Network parameters¶

In [5]:

    # 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

## Before learning¶

We 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) which is a faster method than running the SNN. We also check which neuron is active by looking at the argument that maximizes the function.

In [6]:

    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()

    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)

In [23]:

    # 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()

In [24]:

    # 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 = 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

Out[24]:

## Train network parameters¶

In [27]:

    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()

    # 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 = snn_cvx.update_weights(
                x_sample,
                y_sample,
                F_weights,
                G_weights,
                omega,
                thresholds,
                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

## After learning¶

In [28]:

    # 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)

In [29]:

    # 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()

In [30]:

    # plot contours and active-inactive neurons after learning
    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 = 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

Out[30]:
In [ ]: