framework: allow recording network outputs with the 'record' parameter

Author: nilsceberg

lib/pim/models/new/stone/__init__.py

@@ -15,12 +15,13 @@
 from . import bee_simulator
 
 class StoneResults(ExperimentResults):
-    def __init__(self, name: str, parameters: dict, headings, velocities, log, cpu4_snapshot) -> None:
+    def __init__(self, name: str, parameters: dict, headings, velocities, log, cpu4_snapshot, recordings: dict) -> None:
         super().__init__(name, parameters)
         self.headings = headings
         self.velocities = velocities
         self.log = log
         self.cpu4_snapshot = cpu4_snapshot
+        self.recordings = recordings
 
     def report(self):
         logger.info("plotting route")
@@ -40,7 +41,7 @@ def serialize(self):
         return {
             "headings": self.headings.tolist(),
             "velocities": self.headings.tolist(),
-
+            "recordings": {layer: [entry.tolist() for entry in entries] for layer, entries in self.recordings.items()}
             # annoying to serialize:
             #"log": self.log,
             #"cpu4_snapshot": self.cpu4_snapshot,
@@ -66,20 +67,20 @@ def __init__(self, parameters: dict) -> None:
         noise = self.parameters["noise"]
         cx_type = self.parameters["cx"]
 
-        phi = self.parameters["phi"]
-        beta = self.parameters["beta"]
-        T_half = self.parameters["T_half"]
-        epsilon = self.parameters["epsilon"]
-        length = self.parameters["length"]
-        c_tot = self.parameters["c_tot"]
-
         if cx_type == "basic":
             cx = basic.CXBasic()
         elif cx_type == "rate":
             cx = rate.CXRate(noise = noise)
         elif cx_type == "pontine":
             cx = rate.CXRatePontine(noise = noise)
         elif cx_type == "dye":
+            phi = self.parameters["phi"]
+            beta = self.parameters["beta"]
+            T_half = self.parameters["T_half"]
+            epsilon = self.parameters["epsilon"]
+            length = self.parameters["length"]
+            c_tot = self.parameters["c_tot"]
+
             cx = dye.CXDye(
                 noise=noise,
                 phi=phi,
@@ -103,6 +104,9 @@ def run(self, name: str) -> ExperimentResults:
         cx_type = self.parameters["cx"]
         time_subdivision = self.parameters["time_subdivision"] if "time_subdivision" in self.parameters else 1
 
+        layers_to_record = self.parameters["record"] if "record" in self.parameters else []
+        recordings = {layer: [] for layer in layers_to_record}
+
         logger.info("initializing central complex")
 
         headings = np.zeros(T_outbound + T_inbound)
@@ -120,6 +124,8 @@ def run(self, name: str) -> ExperimentResults:
         for heading, velocity in zip(headings[0:T_outbound], velocities[0:T_outbound, :]):
             for ts in range(time_subdivision):
                 self.cx.update(dt, heading, velocity)
+            for layer in layers_to_record:
+                recordings[layer].append(self.cx.network.output(layer))
 
         for t in range(T_outbound, T_outbound + T_inbound):
             heading = headings[t-1]
@@ -139,6 +145,8 @@ def run(self, name: str) -> ExperimentResults:
                 )
 
             headings[t], velocities[t,:] = heading, velocity
+            for layer in layers_to_record:
+                recordings[layer].append(self.cx.network.output(layer))
 
-        return StoneResults(name, self.parameters, headings, velocities, log = None, cpu4_snapshot = None)
+        return StoneResults(name, self.parameters, headings, velocities, log = None, cpu4_snapshot = None, recordings = recordings)
 

lib/pim/models/new/stone/dye.py

@@ -1,6 +1,6 @@
 from abc import abstractmethod
 from cgi import print_arguments
-from pim.models.new.stone.rate import CXRatePontine, noisy_sigmoid
+from pim.models.new.stone.rate import CXRatePontine, MemorylessCPU4Layer, noisy_sigmoid
 
 from ...network import Network, RecurrentForwardNetwork, Layer, FunctionLayer, Output
 from .constants import *
@@ -9,33 +9,6 @@
 from scipy.interpolate import interp1d
 from scipy.integrate import solve_ivp
 
-class DyeCPU4Layer(Layer):
-    def __init__(self, TB1, TN1, TN2, W_TN, W_TB1, gain, slope, bias, noise, background_activity = 1.0):
-        self.TB1 = TB1
-        self.TN1 = TN1
-        self.TN2 = TN2
-
-        self.W_TN = W_TN
-        self.W_TB1 = W_TB1
-
-        self.gain = gain
-        self.slope = slope
-        self.bias = bias
-
-        self.noise = noise
-        self.background_activity = background_activity
-
-        super().__init__(initial = np.ones(N_CPU4) * self.background_activity)
-
-    def output(self, network: Network) -> Output:
-        tb1 = network.output(self.TB1)
-        tn1 = network.output(self.TN1)
-        tn2 = network.output(self.TN2)
-
-        mem_update = np.dot(self.W_TN, tn2)
-        mem_update -= np.dot(self.W_TB1, tb1)
-        return noisy_sigmoid(mem_update, self.slope, self.bias, self.noise) + self.background_activity
-
 class DyeLayer(Layer):
     def __init__(self):
         self.weights = np.ones(N_CPU4) * 0.5
@@ -153,7 +126,7 @@ def __init__(self, phi=1.0, beta=1.0, T_half=1.0, epsilon=1.0, length=10e-4, c_t
         self.beta = beta
         self.c_tot = c_tot
 
-        super().__init__(DyeCPU4Layer, **kwargs)
+        super().__init__(MemorylessCPU4Layer, **kwargs)
 
     def build_cpu4_layer(self) -> Layer:
         return self.CPU4LayerClass(

lib/pim/models/new/stone/rate.py

@@ -34,6 +34,33 @@ def noisify_weights(W, noise=0.01):
     return W + N_nonzero
 
 
+class MemorylessCPU4Layer(Layer):
+    def __init__(self, TB1, TN1, TN2, W_TN, W_TB1, gain, slope, bias, noise, background_activity = 1.0):
+        self.TB1 = TB1
+        self.TN1 = TN1
+        self.TN2 = TN2
+
+        self.W_TN = W_TN
+        self.W_TB1 = W_TB1
+
+        self.gain = gain
+        self.slope = slope
+        self.bias = bias
+
+        self.noise = noise
+        self.background_activity = background_activity
+
+        super().__init__(initial = np.ones(N_CPU4) * self.background_activity)
+
+    def output(self, network: Network) -> Output:
+        tb1 = network.output(self.TB1)
+        tn1 = network.output(self.TN1)
+        tn2 = network.output(self.TN2)
+
+        mem_update = np.dot(self.W_TN, tn2)
+        mem_update -= np.dot(self.W_TB1, tb1)
+        return noisy_sigmoid(mem_update, self.slope, self.bias, self.noise) + self.background_activity
+
 class CPU4Layer(Layer):
     def __init__(self, TB1, TN1, TN2, W_TN, W_TB1, gain, slope, bias, noise):
         self.TB1 = TB1
@@ -335,6 +362,15 @@ def build_network(self) -> Network:
                 function = self.tn2_output,
                 initial = np.zeros(N_TN2),
             ),
+            "CPU4-no-mem": MemorylessCPU4Layer(
+                "TB1", "TN1", "TN2",
+                self.W_TN_CPU4,
+                self.W_TB1_CPU4,
+                self.cpu4_mem_gain,
+                self.cpu4_slope,
+                self.cpu4_bias,
+                self.noise,
+            ),
             "CPU4": self.build_cpu4_layer(),
             "Pontine": FunctionLayer(
                 inputs = ["CPU4"],