Skip to content

Commit 4bc44de

Browse files
committed
framework: allow recording network outputs with the 'record' parameter
1 parent 200e444 commit 4bc44de

File tree

3 files changed

+56
-39
lines changed

3 files changed

+56
-39
lines changed

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

Lines changed: 18 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -15,12 +15,13 @@
1515
from . import bee_simulator
1616

1717
class StoneResults(ExperimentResults):
18-
def __init__(self, name: str, parameters: dict, headings, velocities, log, cpu4_snapshot) -> None:
18+
def __init__(self, name: str, parameters: dict, headings, velocities, log, cpu4_snapshot, recordings: dict) -> None:
1919
super().__init__(name, parameters)
2020
self.headings = headings
2121
self.velocities = velocities
2222
self.log = log
2323
self.cpu4_snapshot = cpu4_snapshot
24+
self.recordings = recordings
2425

2526
def report(self):
2627
logger.info("plotting route")
@@ -40,7 +41,7 @@ def serialize(self):
4041
return {
4142
"headings": self.headings.tolist(),
4243
"velocities": self.headings.tolist(),
43-
44+
"recordings": {layer: [entry.tolist() for entry in entries] for layer, entries in self.recordings.items()}
4445
# annoying to serialize:
4546
#"log": self.log,
4647
#"cpu4_snapshot": self.cpu4_snapshot,
@@ -66,20 +67,20 @@ def __init__(self, parameters: dict) -> None:
6667
noise = self.parameters["noise"]
6768
cx_type = self.parameters["cx"]
6869

69-
phi = self.parameters["phi"]
70-
beta = self.parameters["beta"]
71-
T_half = self.parameters["T_half"]
72-
epsilon = self.parameters["epsilon"]
73-
length = self.parameters["length"]
74-
c_tot = self.parameters["c_tot"]
75-
7670
if cx_type == "basic":
7771
cx = basic.CXBasic()
7872
elif cx_type == "rate":
7973
cx = rate.CXRate(noise = noise)
8074
elif cx_type == "pontine":
8175
cx = rate.CXRatePontine(noise = noise)
8276
elif cx_type == "dye":
77+
phi = self.parameters["phi"]
78+
beta = self.parameters["beta"]
79+
T_half = self.parameters["T_half"]
80+
epsilon = self.parameters["epsilon"]
81+
length = self.parameters["length"]
82+
c_tot = self.parameters["c_tot"]
83+
8384
cx = dye.CXDye(
8485
noise=noise,
8586
phi=phi,
@@ -103,6 +104,9 @@ def run(self, name: str) -> ExperimentResults:
103104
cx_type = self.parameters["cx"]
104105
time_subdivision = self.parameters["time_subdivision"] if "time_subdivision" in self.parameters else 1
105106

107+
layers_to_record = self.parameters["record"] if "record" in self.parameters else []
108+
recordings = {layer: [] for layer in layers_to_record}
109+
106110
logger.info("initializing central complex")
107111

108112
headings = np.zeros(T_outbound + T_inbound)
@@ -120,6 +124,8 @@ def run(self, name: str) -> ExperimentResults:
120124
for heading, velocity in zip(headings[0:T_outbound], velocities[0:T_outbound, :]):
121125
for ts in range(time_subdivision):
122126
self.cx.update(dt, heading, velocity)
127+
for layer in layers_to_record:
128+
recordings[layer].append(self.cx.network.output(layer))
123129

124130
for t in range(T_outbound, T_outbound + T_inbound):
125131
heading = headings[t-1]
@@ -139,6 +145,8 @@ def run(self, name: str) -> ExperimentResults:
139145
)
140146

141147
headings[t], velocities[t,:] = heading, velocity
148+
for layer in layers_to_record:
149+
recordings[layer].append(self.cx.network.output(layer))
142150

143-
return StoneResults(name, self.parameters, headings, velocities, log = None, cpu4_snapshot = None)
151+
return StoneResults(name, self.parameters, headings, velocities, log = None, cpu4_snapshot = None, recordings = recordings)
144152

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

Lines changed: 2 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
from abc import abstractmethod
22
from cgi import print_arguments
3-
from pim.models.new.stone.rate import CXRatePontine, noisy_sigmoid
3+
from pim.models.new.stone.rate import CXRatePontine, MemorylessCPU4Layer, noisy_sigmoid
44

55
from ...network import Network, RecurrentForwardNetwork, Layer, FunctionLayer, Output
66
from .constants import *
@@ -9,33 +9,6 @@
99
from scipy.interpolate import interp1d
1010
from scipy.integrate import solve_ivp
1111

12-
class DyeCPU4Layer(Layer):
13-
def __init__(self, TB1, TN1, TN2, W_TN, W_TB1, gain, slope, bias, noise, background_activity = 1.0):
14-
self.TB1 = TB1
15-
self.TN1 = TN1
16-
self.TN2 = TN2
17-
18-
self.W_TN = W_TN
19-
self.W_TB1 = W_TB1
20-
21-
self.gain = gain
22-
self.slope = slope
23-
self.bias = bias
24-
25-
self.noise = noise
26-
self.background_activity = background_activity
27-
28-
super().__init__(initial = np.ones(N_CPU4) * self.background_activity)
29-
30-
def output(self, network: Network) -> Output:
31-
tb1 = network.output(self.TB1)
32-
tn1 = network.output(self.TN1)
33-
tn2 = network.output(self.TN2)
34-
35-
mem_update = np.dot(self.W_TN, tn2)
36-
mem_update -= np.dot(self.W_TB1, tb1)
37-
return noisy_sigmoid(mem_update, self.slope, self.bias, self.noise) + self.background_activity
38-
3912
class DyeLayer(Layer):
4013
def __init__(self):
4114
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
153126
self.beta = beta
154127
self.c_tot = c_tot
155128

156-
super().__init__(DyeCPU4Layer, **kwargs)
129+
super().__init__(MemorylessCPU4Layer, **kwargs)
157130

158131
def build_cpu4_layer(self) -> Layer:
159132
return self.CPU4LayerClass(

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

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,33 @@ def noisify_weights(W, noise=0.01):
3434
return W + N_nonzero
3535

3636

37+
class MemorylessCPU4Layer(Layer):
38+
def __init__(self, TB1, TN1, TN2, W_TN, W_TB1, gain, slope, bias, noise, background_activity = 1.0):
39+
self.TB1 = TB1
40+
self.TN1 = TN1
41+
self.TN2 = TN2
42+
43+
self.W_TN = W_TN
44+
self.W_TB1 = W_TB1
45+
46+
self.gain = gain
47+
self.slope = slope
48+
self.bias = bias
49+
50+
self.noise = noise
51+
self.background_activity = background_activity
52+
53+
super().__init__(initial = np.ones(N_CPU4) * self.background_activity)
54+
55+
def output(self, network: Network) -> Output:
56+
tb1 = network.output(self.TB1)
57+
tn1 = network.output(self.TN1)
58+
tn2 = network.output(self.TN2)
59+
60+
mem_update = np.dot(self.W_TN, tn2)
61+
mem_update -= np.dot(self.W_TB1, tb1)
62+
return noisy_sigmoid(mem_update, self.slope, self.bias, self.noise) + self.background_activity
63+
3764
class CPU4Layer(Layer):
3865
def __init__(self, TB1, TN1, TN2, W_TN, W_TB1, gain, slope, bias, noise):
3966
self.TB1 = TB1
@@ -335,6 +362,15 @@ def build_network(self) -> Network:
335362
function = self.tn2_output,
336363
initial = np.zeros(N_TN2),
337364
),
365+
"CPU4-no-mem": MemorylessCPU4Layer(
366+
"TB1", "TN1", "TN2",
367+
self.W_TN_CPU4,
368+
self.W_TB1_CPU4,
369+
self.cpu4_mem_gain,
370+
self.cpu4_slope,
371+
self.cpu4_bias,
372+
self.noise,
373+
),
338374
"CPU4": self.build_cpu4_layer(),
339375
"Pontine": FunctionLayer(
340376
inputs = ["CPU4"],

0 commit comments

Comments
 (0)