]> git.armaanb.net Git - norepinephrine_wm.git/blobdiff - model.py
update
[norepinephrine_wm.git] / model.py
index cd6002207c20f0c1f5b6246163964ce490c9e937..7feeee2be94f4f04b1bdade82760b5d2bc71710e 100644 (file)
--- a/model.py
+++ b/model.py
@@ -1,11 +1,14 @@
 from datetime import datetime
 from os import mkdir
 import logging
+import pickle
 
 import matplotlib.pyplot as plt
 import matplotlib.ticker as mtick
 import nengo
 import numpy as np
+import pandas as pd
+from tqdm import tqdm
 
 exec(open("conf.py").read())
 
@@ -35,30 +38,37 @@ def time_function(t):
 
 
 def decision_function(x):
-    return 1.0 if x[0] + x[1] > 0.0 else -1.0
+    output = 0.0
+    value = x[0] + x[1]
+    if value > 0.0:
+        output = 1.0
+    elif value < 0.0:
+        output = -1.0
+    return output
+    # return 1.0 if x[0] + x[1] > 0.0 else -1.0
 
 
-class Alpha(object):
+class Alpha():
     """
     Base class for alpha receptors. Not to be used directly.
     """
 
     def __init__(self):
-        self.x = np.logspace(0, 3, steps)
+        self.x = steps
         self.y = 1 / (1 + (999 * np.exp(-0.1233 * (self.x / self.offset))))
 
         self.gains = []
         self.biass = []
 
-        for i in range(steps):
-            y = self.y[i]
-            self.gains.append(1 + self.gaind * y)
-            self.biass.append(1 + self.biasd * y)
+        for i in range(len(steps)):
+            self.gains.append(self.gaind * self.y[i] + 1)
+            self.biass.append(self.biasd * self.y[i] + 1)
 
     def plot(self):
         out = f"./out/{self.__class__.__name__}"
 
-        title = "Norepinepherine Concentration vs Neuron Activity in " + self.pretty
+        title = "Norepinepherine Concentration vs Neuron Activity in " + \
+            self.pretty
         logging.info("Plotting " + title)
         plt.figure()
         plt.plot(self.x, self.y)
@@ -140,13 +150,51 @@ class Alpha2(Alpha):
         super().__init__()
 
 
-def simulate(a1, a2):
-    for i in range(steps):
-        gain = a1.gains[i] + a2.gains[i] - 1
-        bias = a1.biass[i] + a2.biass[i] - 1
-        logging.info(f"gain: {fmt_num(gain)}, bias: {fmt_num(bias)}")
+class Simulation():
+    def __init__(self):
+        self.a1 = Alpha1()
+        self.a2 = Alpha2()
+        self.num_spikes = np.ones(len(steps))
+        self.out = np.ones(3)
+        self.trial = 0
+
+        # correctly perceived (not necessarily remembered) cues
+        self.perceived = np.ones(3)
+        rng = np.random.RandomState(seed=seed)
+        # whether the cues is on the left or right
+        self.cues = 2 * rng.randint(2, size=3)-1
+        for n in range(len(self.perceived)):
+            if rng.rand() < misperceive:
+                self.perceived[n] = 0
+
+    def plot(self):
+        title = "Norepinephrine Concentration vs Spiking Rate"
+        logging.info("Plotting " + title)
+        plt.figure()
+        plt.plot(steps, self.num_spikes)
+
+        #plt.xscale("log")
+
+        plt.xlabel("Norepinephrine concentration (nM)")
+        plt.ylabel("Spiking rate (spikes/time step)")
+        plt.title(title)
+
+        plt.draw()
+        plt.savefig("./out/concentration-spiking.png")
+
+    def cue_function(self, t):
+        if t < t_cue and self.perceived[self.trial] != 0:
+            return cue_scale * self.cues[self.trial]
+        else:
+            return 0
+
+    def run(self):
+        self.a1.plot()
+        self.a2.plot()
+
         with nengo.Network() as net:
             # Nodes
+            cue_node = nengo.Node(output=self.cue_function)
             time_node = nengo.Node(output=time_function)
             noise_wm_node = nengo.Node(output=noise_bias_function)
             noise_decision_node = nengo.Node(
@@ -154,13 +202,12 @@ def simulate(a1, a2):
 
             # Ensembles
             wm = nengo.Ensemble(neurons_wm, 2)
-            wm.gain = np.full(wm.n_neurons, gain)
-            wm.bias = np.full(wm.n_neurons, bias)
             decision = nengo.Ensemble(neurons_decide, 2)
             inputs = nengo.Ensemble(neurons_inputs, 2)
             output = nengo.Ensemble(neurons_decide, 1)
 
             # Connections
+            nengo.Connection(cue_node, inputs[0], synapse=None)
             nengo.Connection(time_node, inputs[1], synapse=None)
             nengo.Connection(inputs, wm, synapse=tau_wm,
                              function=inputs_function)
@@ -174,25 +221,55 @@ def simulate(a1, a2):
             nengo.Connection(decision, output, function=decision_function)
 
             # Probes
-            probes_wm = nengo.Probe(wm[0], synapse=0.01)
-            probe_output = nengo.Probe(output, synapse=None)
-
-        # Run simulation
-        with nengo.Simulator(net, dt=dt, progress_bar=False) as sim:
-            sim.run(t_cue + t_delay)
+            probes_wm = nengo.Probe(wm[0], synapse=0.01, sample_every=probe_dt)
+            probe_spikes = nengo.Probe(wm.neurons, sample_every=probe_dt)
+            probe_output = nengo.Probe(
+                output, synapse=None, sample_every=probe_dt)
+
+            # Run simulation
+            with nengo.Simulator(net, dt=dt, progress_bar=False) as sim:
+                for i, _ in tqdm(enumerate(steps), total=len(steps)):
+                    wm.gain = (self.a1.gains[i] + self.a2.gains[i]) * sim.data[wm].gain
+                    wm.bias = (self.a1.biass[i] + self.a2.biass[i]) * sim.data[wm].bias
+                    wm_recurrent.solver = MySolver(
+                        sim.model.params[wm_recurrent].weights)
+                    wm_to_decision.solver = MySolver(
+                        sim.model.params[wm_to_decision].weights)
+                    sim = nengo.Simulator(net, dt=dt, progress_bar=False)
+                    for self.trial in range(3):
+                        logging.info(
+                            f"Simulating: trial: {self.trial}, gain: {fmt_num(wm.gain)}, bias: {fmt_num(wm.bias)}")
+                        sim.run(t_cue + t_delay)
+                        self.out[self.trial] = np.count_nonzero(
+                            sim.data[probe_spikes])
+                    self.num_spikes[i] = np.average(self.out)
+
+        with open(f"out/{datetime.now().isoformat()}-spikes.pkl", "wb") as pout:
+            pickle.dump(self, pout)
+
+        self.plot()
+
+
+class MySolver(nengo.solvers.Solver):
+    def __init__(self, weights):
+        self.weights = False
+        self.my_weights = weights
+        self._paramdict = dict()
+
+    def __call__(self, A, Y, rng=None, E=None):
+        return self.my_weights.T, dict()
 
 
 def main():
     logging.info("Initializing simulation")
     plt.style.use("ggplot")  # Nice looking and familiar style
 
-    a1 = Alpha1()
-    a1.plot()
-
-    a2 = Alpha2()
-    a2.plot()
-
-    simulate(a1, a2)
+    try:
+        data = open("simulation.pkl", "rb")
+    except FileNotFoundError:
+        Simulation().run()
+    else:
+        pickle.load(data).plot()
 
 
 if __name__ == "__main__":
@@ -202,9 +279,6 @@ if __name__ == "__main__":
         pass
 
     logging.basicConfig(filename=f"out/{datetime.now().isoformat()}.log",
-                        level=logging.DEBUG)
-    console = logging.StreamHandler()
-    console.setLevel(logging.INFO)
-    logging.getLogger("").addHandler(console)
+                        level=logging.INFO)
 
     main()