import numpy as np

def poisson_homogenous(mu, timesteps):
    """
    Generate a spiketrain for a single neuron
    """
    size = (1, timesteps)
    spiketrain = np.random.poisson(lam=mu, size=size)
    spiketrain = np.array(spiketrain, dtype=bool)

    return spiketrain

def poisson_inhomogenous(mus, timesteps):
    """
    Generate a spiketrain for a single neuron
    using an inhomogenous poisson distribution
    :param mus: List or Tuple of the Lambdas
    """

    if timesteps % len(mus) != 0:
        raise Exception("Cannot divide the %d mu's on the %d timesteps equally." % (len(mus), timesteps))

    spiketrain = np.zeros((1, timesteps), dtype=bool)

    bucketsize = timesteps // len(mus)

    for i, mu in enumerate(mus):
        startindex = i*bucketsize
        spiketrain[0, startindex:startindex+bucketsize] = np.random.poisson(lam=mu, size=bucketsize)

    # print("Generated inhomogenous spiketrain", spiketrain)

    return spiketrain

def sound(timesteps, midpoint, maximum, variance):
    mu = np.arange(timesteps)
    mu = maximum * np.exp(-((mu - midpoint ) ** 2) / variance ** 2)
    import matplotlib.pyplot as plt
    plt.plot(mu)
    s = poisson_inhomogenous(mu, timesteps)
    return s

def plot(spiketrain):

    import matplotlib.pyplot as plt

    plt.title("Spiketrain plot")
    plt.ylabel("Neuron")
    plt.xlabel("Time")

    r, c = np.where(spiketrain)
    plt.scatter(c, r, c='r', s=20, marker="x")

    r, c = np.where(spiketrain == False)
    plt.scatter(c, r, c='b', s=5, marker="_")
    plt.show()


def plot_psth(spiketrain, binsize=20):
    """ Plot a Peri-Stimulus Time Histogram """

    import matplotlib.pyplot as plt

    neurons, timesteps = spiketrain.shape

    bins=range(0, timesteps + binsize, binsize)
    print ("Bins:", bins)

    fig, axes = plt.subplots(nrows=neurons)

    for i, axis in enumerate(axes):
        # Histogramm
        times = np.where(spiketrain[i])[0]
        axis.hist(times, bins, histtype='bar', stacked=True, fill=True, facecolor='green', alpha=0.5, zorder=0)

        # Scatter (Spikes)
        y = np.ones(times.shape)
        axis.scatter(times, y, c='r', s=20, marker="x", zorder=1)

        axis.set_title('Neuron %d' % i)
        axis.set_ylim(bottom=0)
        axis.set_xlim(0, timesteps)

    plt.tight_layout()
    plt.show()


if __name__ == "__main__":
    s = np.zeros((5, 280), dtype=bool)
    for i in range(5):
        s[i] = sound(280, i*50, 0.4, 50)

    print(s)
    plot_psth(s, binsize=25)
@@ -13,16 +13,116 @@
  as np%0A%0A
+import matplotlib.pyplot as plt%0Aimport matplotlib.animation as animation%0Afrom matplotlib import cm%0A%0A
 def pois
@@ -1342,45 +1342,8 @@
 ):%0A%0A
-    import matplotlib.pyplot as plt%0A%0A
     
@@ -1604,16 +1604,619 @@
 how()%0A%0A%0A
+class WeightPlotter:%0A    def __init__(self, save_interval):%0A        self.save_interval = save_interval%0A        self.weights = %5B%5D%0A%0A    def add(self, weights):%0A        self.weights.append(weights.copy())%0A%0A    def plot_weights(self):%0A%0A        fig = plt.figure()%0A%0A        ims = %5B%5D%0A        for weight in self.weights:%0A            im = plt.imshow(weight, vmin=-1, vmax=1, interpolation='none', cmap=cm.coolwarm)%0A            ims.append(%5Bim%5D)%0A%0A        ani = animation.ArtistAnimation(fig, ims, interval=20, repeat_delay=3000)%0A#        ani.save('weightchange.mp4', writer='mencoder', fps=15)%0A        plt.show()%0A%0A
 def plot
@@ -2243,16 +2243,37 @@
 nsize=20
+, neuron_indices=None
 ):%0A    %22
@@ -2467,24 +2467,189 @@
 s:%22, bins)%0A%0A
+    if not neuron_indices:%0A        print (%22Printing all neurons!!%22)%0A        neuron_indices = range(neurons)%0A%0A    n_plots = len(neuron_indices) # Number of subplots%0A%0A
     fig, axe
@@ -2672,21 +2672,21 @@
 (nrows=n
-euron
+_plot
 s)%0A%0A    
@@ -2713,24 +2713,65 @@
 rate(axes):%0A
+        neuron_index = neuron_indices%5Bi%5D%0A
         # Hi
@@ -2815,17 +2815,28 @@
 ketrain%5B
-i
+neuron_index
 %5D)%5B0%5D%0A  
@@ -3110,17 +3110,28 @@
 n %25d' %25 
-i
+neuron_index
 )%0A      

Add weightplotter, change plot_psth so that one can define neurons to print out