simulation: update bee physics to allow for variable time steps, and use that in StoneExperiment

Author: nilsceberg

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

@@ -45,23 +45,26 @@ def serialize(self):
             #"cpu4_snapshot": self.cpu4_snapshot,
         }
 
+    def reconstruct_path(self):
+        position = np.zeros(2)
+        positions = [position]
+        for velocity in self.velocities:
+            position = position + velocity
+            positions.append(position)
+        return positions
+
+    def closest_position(self):
+        path = self.reconstruct_path()
+        return min(path[self.parameters["T_outbound"]:], key = np.linalg.norm)
+
 class StoneExperiment(Experiment):
     def __init__(self, parameters: dict) -> None:
         super().__init__()
         self.parameters = parameters
 
-    def run(self, name: str) -> ExperimentResults:
-        # extract some parameters
-        T_outbound = self.parameters["T_outbound"]
-        T_inbound = self.parameters["T_inbound"]
         noise = self.parameters["noise"]
         cx_type = self.parameters["cx"]
 
-        logger.info(f"generating outbound route")
-        headings, velocities = trials.generate_route(T = T_outbound, vary_speed = True)
-
-        logger.info("initializing central complex")
-
         if cx_type == "basic":
             cx = basic.CXBasic()
         elif cx_type == "rate":
@@ -71,7 +74,23 @@ def run(self, name: str) -> ExperimentResults:
         else:
             raise RuntimeError("unknown cx type: " + cx_type)
 
-        cx.setup()
+        self.cx = cx
+        self.cx.setup()
+
+    def run(self, name: str) -> ExperimentResults:
+        # extract some parameters
+        T_outbound = self.parameters["T_outbound"]
+        T_inbound = self.parameters["T_inbound"]
+        noise = self.parameters["noise"]
+        cx_type = self.parameters["cx"]
+        time_subdivision = self.parameters["time_subdivision"] if "time_subdivision" in self.parameters else 1
+
+        logger.info(f"generating outbound route")
+        headings, velocities = trials.generate_route(T = T_outbound, vary_speed = True)
+        #headings = np.repeat(headings, time_subdivision)
+        #headings = np.repeat(headings, time_subdivision)
+
+        logger.info("initializing central complex")
 
         headings = np.zeros(T_outbound + T_inbound)
         velocities = np.zeros((T_outbound + T_inbound, 2))
@@ -84,23 +103,29 @@ def run(self, name: str) -> ExperimentResults:
         )
 
         logger.info("simulating outbound path")
+        dt = 1.0 / time_subdivision
         for heading, velocity in zip(headings[0:T_outbound], velocities[0:T_outbound, :]):
-            cx.update(1.0, heading, velocity)
+            for ts in range(time_subdivision):
+                self.cx.update(dt, heading, velocity)
 
         for t in range(T_outbound, T_outbound + T_inbound):
             heading = headings[t-1]
             velocity = velocities[t-1,:]
 
-            motor = cx.update(1.0, heading, velocity)
-            rotation = motor
+            for ts in range(time_subdivision):
+                motor = self.cx.update(dt, heading, velocity)
+                rotation = motor
 
-            headings[t], velocities[t,:] = bee_simulator.get_next_state(
-                velocity=velocity,
-                heading=heading,
-                acceleration=0.1,
-                drag=trials.default_drag,
-                rotation=rotation,
-            )
+                heading, velocity = bee_simulator.get_next_state(
+                    dt=dt,
+                    velocity=velocity,
+                    heading=heading,
+                    acceleration=0.1,
+                    drag=trials.default_drag,
+                    rotation=rotation,
+                )
+
+            headings[t], velocities[t,:] = heading, velocity
 
         return StoneResults(name, self.parameters, headings, velocities, log = None, cpu4_snapshot = None)
 

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

@@ -81,7 +81,6 @@ def step(self, network: Network, dt: float):
         mem_update -= 0.5 * (0.5 - tn1.reshape(2, 1))
 
         # Constant purely to visualise same as rate-based model
-        # multiply with dt here or multiply incoming TN?
         mem_reshaped += self.gain * mem_update * dt
 
         self.memory = np.clip(mem_reshaped.reshape(-1), 0.0, 1.0)

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

@@ -62,26 +62,26 @@ def get_flow(heading, velocity, pref_angle=np.pi/4):
     return np.dot(A, velocity)
 
 
-def rotate(theta, r):
+def rotate(dt, theta, r):
     """Return new heading after a rotation around Z axis."""
-    return (theta + r + np.pi) % (2.0 * np.pi) - np.pi
+    return (theta + r * dt + np.pi) % (2.0 * np.pi) - np.pi
 
 
-def thrust(theta, acceleration):
+def thrust(dt, theta, acceleration):
     """Thrust vector from current heading and acceleration
 
     theta: clockwise radians around z-axis, where 0 is forward
     acceleration: float where max speed is ....?!?
     """
-    return np.array([np.sin(theta), np.cos(theta)]) * acceleration
+    return np.array([np.sin(theta), np.cos(theta)]) * acceleration * dt
 
 
-def get_next_state(heading, velocity, rotation, acceleration, drag=0.5):
+def get_next_state(dt, heading, velocity, rotation, acceleration, drag=0.5):
     """Get new heading and velocity, based on relative rotation and
     acceleration and linear drag."""
-    theta = rotate(heading, rotation)
-    v = velocity + thrust(theta, acceleration)
-    v -= drag * v
+    theta = rotate(dt, heading, rotation)
+    v = velocity + thrust(dt, theta, acceleration)
+    v *= (1.0 - drag)**dt
     return theta, v
 
 

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

@@ -88,12 +88,12 @@ def step(self, network: Network, dt: float):
         cpu4[0-7] store optic flow peaking at left 45 deg
         cpu[8-15] store optic flow peaking at right 45 deg."""
         tb1 = network.output(self.TB1)
-        tn1 = network.output(self.TN1) * dt
-        tn2 = network.output(self.TN2) * dt
+        tn1 = network.output(self.TN1)
+        tn2 = network.output(self.TN2)
 
         self.memory += (np.clip(np.dot(self.W_TN, 0.5-tn1), 0, 1) *
-                     self.gain * np.dot(self.W_TB1, 1.0-tb1))
-        self.memory -= self.gain * 0.25 * np.dot(self.W_TN, tn2)
+                     self.gain * np.dot(self.W_TB1, 1.0-tb1)) * dt
+        self.memory -= self.gain * 0.25 * np.dot(self.W_TN, tn2) * dt
         self.memory = np.clip(self.memory, 0.0, 1.0)
 
     def output(self, network: Network) -> Output:
@@ -107,14 +107,14 @@ def step(self, network: Network, dt: float):
         cpu4[0-7] store optic flow peaking at left 45 deg
         cpu[8-15] store optic flow peaking at right 45 deg."""
         tb1 = network.output(self.TB1)
-        tn1 = network.output(self.TN1) * dt
-        tn2 = network.output(self.TN2) * dt
+        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)
         mem_update = np.clip(mem_update, 0, 1)
         mem_update *= self.gain
-        self.memory += mem_update
+        self.memory += mem_update * dt
         self.memory -= 0.125 * self.gain * dt
         self.memory = np.clip(self.memory, 0.0, 1.0)
 
@@ -149,15 +149,7 @@ def build_network(self) -> Network:
                 function = self.tn2_output,
                 initial = np.zeros(N_TN2),
             ),
-            "CPU4": CPU4Layer(
-                "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(),
             "CPU1": FunctionLayer(
                 inputs = ["TB1", "CPU4"],
                 function = self.cpu1_output,
@@ -169,10 +161,22 @@ def build_network(self) -> Network:
             )
         })
 
+    def build_cpu4_layer(self) -> Layer:
+        return self.CPU4LayerClass(
+            "TB1", "TN1", "TN2",
+            self.W_TN_CPU4,
+            self.W_TB1_CPU4,
+            self.cpu4_mem_gain,
+            self.cpu4_slope,
+            self.cpu4_bias,
+            self.noise,
+        )
+
     """Class to keep a set of parameters for a model together.
     No state is held in the class currently."""
     def __init__(self,
                  noise=0.1,
+                 CPU4LayerClass = CPU4Layer,
                  tl2_slope=tl2_slope_tuned,
                  tl2_bias=tl2_bias_tuned,
                  tl2_prefs=np.tile(np.linspace(0, 2*np.pi, N_TB1,
@@ -190,6 +194,7 @@ def __init__(self,
                  weight_noise=0.0,
                  ):
         super().__init__()
+        self.CPU4LayerClass = CPU4LayerClass
 
         # Default noise used by the model for all layers
         self.noise = noise
@@ -361,15 +366,7 @@ def build_network(self) -> Network:
                 function = self.tn2_output,
                 initial = np.zeros(N_TN2),
             ),
-            "CPU4": CPU4PontinLayer(
-                "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(),
             "Pontin": FunctionLayer(
                 inputs = ["CPU4"],
                 function = self.pontin_output,
@@ -386,8 +383,8 @@ def build_network(self) -> Network:
             )
         })
 
-    def __init__(self, **kwargs):
-        super().__init__(**kwargs)
+    def __init__(self, CPU4LayerClass = CPU4PontinLayer, **kwargs):
+        super().__init__(CPU4LayerClass = CPU4LayerClass, **kwargs)
 
         self.cpu4_mem_gain *= 0.5
         self.cpu1_bias = -1.0

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

@@ -44,6 +44,7 @@ def generate_route(T=1500, mean_acc=default_acc, drag=default_drag,
 
     for t in range(1, T):
         headings[t], velocity[t, :] = bee_simulator.get_next_state(
+            dt=1.0,
             heading=headings[t-1], velocity=velocity[t-1, :],
             rotation=rotation[t], acceleration=acceleration[t], drag=drag)
     return headings, velocity