flexible head direction

.gitignore

@@ -0,0 +1,2 @@
+**__pycache__**
+figures/*

ratinabox/Agent.py

@@ -6,6 +6,7 @@
 import os
 import matplotlib
 from matplotlib import pyplot as plt
+import warnings
 
 
 from ratinabox import utils
@@ -41,6 +42,7 @@ class Agent:
             "speed_std": 0.08,
             "rotational_velocity_coherence_time": 0.08,
             "rotational_velocity_std": 120 * (np.pi / 180),
+            "head_direction_smoothing_timescale" : 0.0,
             "thigmotaxis": 0.5,
             "wall_repel_distance": 0.1,
             "walls_repel": True,
@@ -62,6 +64,7 @@ class Agent:
         "rotational_velocity_std": (
             120 * (np.pi / 180)
         ),  # std of rotational speed, σ_w wall following parameter
+        "head_direction_smoothing_timescale" : 0.0, # timescale over which head direction is smoothed (head dir = smoothed velocity vector) 
         "thigmotaxis": 0.5,  # tendency for agents to linger near walls [0 = not at all, 1 = max]
         "wall_repel_distance": 0.1,
         "walls_repel": True,  # whether or not the walls repel
@@ -84,12 +87,20 @@ def __init__(self, Environment, params={}):
         utils.update_class_params(self, self.params, get_all_defaults=True)
         utils.check_params(self, params.keys())
 
+        # check if dt < coherence times and warn if not
+        if self.head_direction_smoothing_timescale!= 0 and self.dt >= self.head_direction_smoothing_timescale:
+            warnings.warn("WARNING: dt >= head_direction_smoothing_timescale. This will break the head direction smoothing. \
+                          Set head_direction_smoothing_timescale to 0 to disable head direction smoothing OR \
+                          set dt < head_direction_smoothing_timescale")
+
+
         # initialise history dataframes
         self.history = {}
         self.history["t"] = []
         self.history["pos"] = []
         self.history["vel"] = []
         self.history["rot_vel"] = []
+        self.history["head_direction"] = []
 
         self.Neurons = []  # each new Neurons class belonging to this Agent will append itself to this list
 
@@ -117,9 +128,14 @@ def __init__(self, Environment, params={}):
             self.velocity = np.array([self.speed_mean])
             if self.Environment.boundary_conditions == "solid":
                 if self.speed_mean != 0:
-                    print(
+                    warnings.warn(
                         "Warning: You have solid 1D boundary conditions and non-zero speed mean."
                     )
+
+        # normally this will just be a low pass filter over the velocity vector
+        # this is done to smooth out head turning and make it more realistic 
+        # (potentially stablize behaviours associated with head direction)
+        self.head_direction = self.velocity
 
         if ratinabox.verbose is True:
             print(
@@ -416,6 +432,8 @@ def update(self, dt=None, drift_velocity=None, drift_to_random_strength_ratio=1)
                     self.times
                 )
 
+        self.update_head_direction(dt=dt)
+
         if len(self.history["pos"]) >= 1:
             self.distance_travelled += np.linalg.norm(
                 self.Environment.get_vectors_between___accounting_for_environment(
@@ -435,12 +453,40 @@ def update(self, dt=None, drift_velocity=None, drift_to_random_strength_ratio=1)
 
         return
 
+    def update_head_direction(self, dt):
+        """
+        This function updates the head direction of the agent. #
+
+        Args: 
+            dt: the time step (float)
+
+        The head direction is updated by a low pass filter of the the current velocity vector.
+        """
+
+        tau_head_direction = self.head_direction_smoothing_timescale
+        immediate_head_direction = self.velocity / np.linalg.norm(self.velocity)
+
+        if dt < tau_head_direction:
+            warnings.warn("dt >= head_direction_smoothing_timescale. This will break the head direction smoothing.")
+
+        if self.head_direction is None:
+            self.head_direction = self.velocity
+
+        if tau_head_direction == 0: 
+            self.head_direction = immediate_head_direction
+        else:
+            self.head_direction = self.head_direction * ( 1 - dt / tau_head_direction ) + dt / tau_head_direction * immediate_head_direction
+
+        # normalize the head direction
+        self.head_direction = self.head_direction / np.linalg.norm(self.head_direction)
+
     def save_to_history(self):
         self.history["t"].append(self.t)
         self.history["pos"].append(list(self.pos))
         self.history["vel"].append(list(self.save_velocity))
+        self.history["head_direction"].append(self.head_direction)
         if self.Environment.dimensionality == "2D":
-            self.history["rot_vel"].append(self.rotational_velocity)
+            self.history["rot_vel"].append(self.rotational_velocity)     
         return
 
     def reset_history(self):

ratinabox/Neurons.py

@@ -8,6 +8,7 @@
 from matplotlib import pyplot as plt
 import scipy
 from scipy import stats as stats
+import warnings
 
 from ratinabox import utils
 
@@ -1242,14 +1243,20 @@ def get_state(self, evaluate_at="agent", **kwargs):
         # if egocentric references frame shift angle into coordinate from of heading direction of agent
         if self.reference_frame == "egocentric":
             if evaluate_at == "agent":
-                vel = self.Agent.velocity
+                head_direction = self.Agent.head_direction
+            elif "head_direction" in kwargs.keys():
+                head_direction = kwargs["head_direction"]
             elif "vel" in kwargs.keys():
-                vel = kwargs["vel"]
+                # just to make backwards compatible
+                warning.warn("'vel' kwarg deprecated in favour of 'head_direction'")
+                head_direction = kwargs["vel"]
             else:
-                vel = np.array([1, 0])
-            vel = np.array(vel)
-            head_direction_angle = utils.get_angle(vel)
-            test_angles = test_angles - head_direction_angle
+                head_direction = np.array([1, 0])
+                warnings.warn(
+                    "BVCs in egocentric plane require a head direction vector but none was passed. Using [1,0]"
+                )
+            head_bearing = utils.get_angle(head_direction)
+            test_angles -= head_bearing  # account for head direction
 
         tuning_angles = np.tile(
             np.expand_dims(np.expand_dims(self.tuning_angles, axis=-1), axis=-1),
@@ -1529,15 +1536,19 @@ def get_state(self, evaluate_at="agent", **kwargs):
         )  # (N_pos,N_objects) #vectors go from pos2 to pos1 so must do subtract pi from bearing
         if self.reference_frame == "egocentric":
             if evaluate_at == "agent":
-                vel = self.Agent.velocity
+                head_direction = self.Agent.head_direction
+            elif "head_direction" in kwargs.keys():
+                head_direction = kwargs["head_direction"]
             elif "vel" in kwargs.keys():
-                vel = kwargs["vel"]
+                # just to make backwards compatible
+                warning.warn("'vel' kwarg deprecated in favour of 'head_direction'")
+                head_direction = kwargs["vel"]
             else:
-                vel = np.array([1, 0])
-                print(
-                    "Field of view OVCs require a velocity vector but none was passed. Using [1,0]"
+                head_direction = np.array([1, 0])
+                warnings.warn(
+                    "OVCs in egocentric plane require a head direction vector but none was passed. Using [1,0]"
                 )
-            head_bearing = utils.get_angle(vel)
+            head_bearing = utils.get_angle(head_direction)
             bearings_to_objects -= head_bearing  # account for head direction
 
         tuning_distances = np.tile(
@@ -1648,24 +1659,27 @@ def get_state(self, evaluate_at="agent", **kwargs):
         """In 2D n head direction cells encode the head direction of the animal. By default velocity (which determines head direction) is taken from the agent but this can also be passed as a kwarg 'vel'"""
 
         if evaluate_at == "agent":
-            vel = self.Agent.history["vel"][-1]
+            head_direction = self.Agent.head_direction
+        elif "head_direction" in kwargs.keys():
+            head_direction = kwargs["head_direction"]
         elif "vel" in kwargs.keys():
-            vel = np.array(kwargs["vel"])
+            head_direction = np.array(kwargs["vel"])
+            warning.warn("'vel' kwarg deprecated in favour of 'head_direction'")
         else:
-            print("HeadDirection cells need a velocity but not was given, taking...")
+            print("HeadDirection cells need a head direction but not was given, taking...")
             if self.Agent.Environment.dimensionality == "2D":
-                vel = np.array([1, 0])
+                head_direction = np.array([1, 0])
                 print("...[1,0] as default")
             if self.Agent.Environment.dimensionality == "1D":
-                vel = np.array([1])
+                head_direction = np.array([1])
                 print("...[1] as default")
 
         if self.Agent.Environment.dimensionality == "1D":
-            hdleft_fr = max(0, np.sign(vel[0]))
-            hdright_fr = max(0, -np.sign(vel[0]))
+            hdleft_fr = max(0, np.sign(head_direction[0]))
+            hdright_fr = max(0, -np.sign(head_direction[0]))
             firingrate = np.array([hdleft_fr, hdright_fr])
         if self.Agent.Environment.dimensionality == "2D":
-            current_angle = utils.get_angle(vel)
+            current_angle = utils.get_angle(head_direction)
             firingrate = utils.von_mises(
                 current_angle, self.preferred_angles, self.angular_tunings, norm=1
             )

ratinabox/contribs/FieldOfViewNeurons.py

@@ -141,10 +141,10 @@ def display_manifold(self, fig=None, ax=None, t=None, object_type=0, **kwargs):
             fig, ax = self.Agent.plot_trajectory(t_start=t - 10, t_end=t, **kwargs)
 
         pos = self.Agent.history["pos"][t_id]
-        vel = self.Agent.history["vel"][t_id]
+        head_direction = self.Agent.history["head_direction"][t_id]
         fr = self.history["firingrate"][t_id]
-        ego_y = vel / np.linalg.norm(
-            vel
+        ego_y = head_direction / np.linalg.norm(
+            head_direction
         )  # this is the "y" direction" in egocentric coords
         ego_x = utils.rotate(ego_y, -np.pi / 2)
         facecolor = "C1"