Refactoring velocity predictors

collab2/foraging/toolkit/__init__.py

@@ -55,9 +55,6 @@
     update_rewards,
 )
 from .velocity import (  # noqa: F401
-    _add_velocity,
-    _generic_velocity_predictor,
-    _velocity_predictor_contribution,
     generate_pairwiseCopying_predictor,
     generate_vicsek_predictor,
 )

collab2/foraging/toolkit/velocity.py

@@ -40,7 +40,9 @@ def _add_velocity(
             v_x = df["x"].diff(periods=dt) / dt
             v_y = df["y"].diff(periods=dt) / dt
             df[v_ID] = np.sqrt(v_x**2 + v_y**2)
-            df[theta_ID] = np.arctan2(v_y, v_x)
+            df[theta_ID] = np.arctan2(
+                v_y, v_x
+            )  # note that this returns 0 (instead of NaN) when vx=0,vy=0!
 
     return foragers_processed, pd.concat(foragers_processed)
 
@@ -71,16 +73,25 @@ def _velocity_predictor_contribution(
     v_implied = np.sqrt((grid["x"] - x) ** 2 + (grid["y"] - y) ** 2)
     theta_implied = np.arctan2(grid["y"] - y, grid["x"] - x)
     P_v = norm.pdf(x=v_implied, loc=v_pref, scale=sigma_v)
-    # there is a discontinuity when taking the difference of angles (2pi \equiv 0 !),
-    # so always choose the smaller difference
-    d_theta = theta_implied - theta_pref
-    d_theta[d_theta > np.pi] += -2 * np.pi
-    d_theta[d_theta < -np.pi] += 2 * np.pi
-    P_theta = norm.pdf(x=d_theta, loc=0, scale=sigma_t)
+
+    # note that np.arctan2(0,0) = 0 (instead of NaN)
+    # however, when |v|=0, the direction is meaningless, and it does not make sense to
+    # have angular preference in the predictor.
+    # therefore, we check for that condition and return an isotropic gaussian if true
+
+    if v_pref == 0 or np.isnan(theta_pref):  # then theta_pref does not matter!
+        P_theta = 1
+    else:
+        # there is a discontinuity when taking the difference of angles (2pi \equiv 0 !),
+        # so always choose the smaller difference
+        d_theta = theta_implied - theta_pref
+        d_theta[d_theta > np.pi] += -2 * np.pi
+        d_theta[d_theta < -np.pi] += 2 * np.pi
+        P_theta = norm.pdf(x=d_theta, loc=0, scale=sigma_t)
     return P_v * P_theta
 
 
-def _generic_velocity_predictor(
+def _generate_pairwiseCopying_predictor(
     foragers: List[pd.DataFrame],
     foragersDF: pd.DataFrame,
     local_windows: List[List[pd.DataFrame]],
@@ -89,33 +100,29 @@ def _generic_velocity_predictor(
     dt: int,
     sigma_v: float,
     sigma_t: float,
-    transformation_function: Callable[[pd.DataFrame], pd.DataFrame],
     interaction_constraint: Optional[
         Callable[[List[int], int, int, pd.DataFrame], List[int]]
     ] = None,
     interaction_constraint_params: dict[str, Any] = {},
 ) -> List[List[pd.DataFrame]]:
     """
-    A function that calculates predictor scores for arbitrary velocity alignment mechanisms, as specified by
-    `transformation_function`. This function takes the velocities of interaction partners and outputs a transformation
-    (eg, average, identity, max) as required by the corresponding mechanism
-    Predictors are not calculated for frames where interaction partners have missing velocities.
-    In this case, fraction of dropped frames is reported.
-    Predictors are normalized by dividing by their max value for each forager & frame.
+    A function that calculates predictor scores for a pairwise copying interaction, using
+    `_velocity_predictor_contribution`. Predictors are not calculated for frames where
+    interaction partners have missing velocities.
+
+    Predictors are rescaled by dividing by their max value for each forager & frame.
 
     :param foragers : List of DataFrames containing forager positions and velocities grouped by forager index
     :param foragersDF : Flattened DataFrame of forager positions and velocities
     :param local_windows : Nested list of DataFrames containing grid points to compute predictor over,
             grouped by forager index and time
     :param predictor_name : Name given to column containing predictor scores in `predictor`
-    :param nteraction_length : Maximum inter-forager distance for velocity copying interaction
+    :param interaction_length : Maximum inter-forager distance for velocity copying interaction
     :param dt : frames skipped in calculation of velocities
             Note: This function requires `foragers` and `foragersDF` to contain
             columns "v_dt={dt}", "theta_dt={dt}"
     :param sigma_v : standard deviation of Gaussian for velocity magnitude
     :param sigma_t : standard deviation of Gaussian for velocity direction
-    :param transformation_function : Function that implements a transformation of velocities of interaction partners,
-            as stipulated by the chosen velocity alignment mechanism
     :param interaction_constraint : Optional function to model other interaction constraints
     :param interaction_constraint_params : Optional kwargs to be passed to `interaction_constraint`
     :return: Nested list of calculated predictor scores, grouped by foragers and time
@@ -150,7 +157,6 @@ def _generic_velocity_predictor(
                 ]
 
                 if v_values.notna().all(axis=None):
-                    v_values = transformation_function(v_values)
                     x = foragers[f].loc[t, "x"]
                     y = foragers[f].loc[t, "y"]
                     # additively combine the influence of all confocals
@@ -178,7 +184,7 @@ def generate_pairwiseCopying_predictor(
 ):
     """
     A function that calculates the predictor scores associated with random, pairwise velocity copying,
-    by specifying an identity transformation to `_generic_velocity_predictor`.
+    by calling `_generate_pairwiseCopying_predictor`.
     The necessary parameters from `foragers_object`. Thus, `foragers_object` must contain as attribute
     `predictor_kwargs` : dict, with `predictor_name` as a valid key.
 
@@ -187,10 +193,6 @@ def generate_pairwiseCopying_predictor(
     :return: Nested list of calculated predictor scores, grouped by foragers and time
     """
 
-    # define transformation function
-    def transformation_pairwiseCopying(v_values):
-        return v_values
-
     # grab relevant parameters from foragers_object
     params = foragers_object.predictor_kwargs[predictor_name]
 
@@ -200,37 +202,138 @@ def transformation_pairwiseCopying(v_values):
     )
 
     # calculate predictor values
-    predictor = _generic_velocity_predictor(
+    predictor = _generate_pairwiseCopying_predictor(
         foragers_object.foragers,
         foragers_object.foragersDF,
         foragers_object.local_windows,
         predictor_name,
-        transformation_function=transformation_pairwiseCopying,
         **params,
     )
 
     return predictor
 
 
+def _generate_vicsek_predictor(
+    foragers: List[pd.DataFrame],
+    foragersDF: pd.DataFrame,
+    local_windows: List[List[pd.DataFrame]],
+    predictor_name: str,
+    interaction_length: float,
+    dt: int,
+    sigma_v: float,
+    sigma_t: float,
+    interaction_constraint: Optional[
+        Callable[[List[int], int, int, pd.DataFrame], List[int]]
+    ] = None,
+    interaction_constraint_params: dict[str, Any] = {},
+) -> List[List[pd.DataFrame]]:
+    """
+    A function that calculates predictor scores for a vicsek alignment interaction, using
+    `_velocity_predictor_contribution`. Predictors are not calculated for frames where
+    interaction partners have missing velocities.
+
+    Predictors are rescaled by dividing by their max value for each forager & frame.
+
+    :param foragers : List of DataFrames containing forager positions and velocities grouped by forager index
+    :param foragersDF : Flattened DataFrame of forager positions and velocities
+    :param local_windows : Nested list of DataFrames containing grid points to compute predictor over,
+            grouped by forager index and time
+    :param predictor_name : Name given to column containing predictor scores in `predictor`
+    :param interaction_length : Maximum inter-forager distance for velocity copying interaction
+    :param dt : frames skipped in calculation of velocities
+            Note: This function requires `foragers` and `foragersDF` to contain
+            columns "v_dt={dt}", "theta_dt={dt}"
+    :param sigma_v : standard deviation of Gaussian for velocity magnitude
+    :param sigma_t : standard deviation of Gaussian for velocity direction
+    :param interaction_constraint : Optional function to model other interaction constraints
+    :param interaction_constraint_params : Optional kwargs to be passed to `interaction_constraint`
+    :return: Nested list of calculated predictor scores, grouped by foragers and time
+    """
+
+    num_foragers = len(foragers)
+    num_frames = len(foragers[0])
+    predictor = copy.deepcopy(local_windows)
+
+    for f in range(num_foragers):
+        for t in range(num_frames):
+            if predictor[f][t] is not None:
+                # add column for predictor_ID
+                predictor[f][t][predictor_name] = 0
+                # find confocals within interaction length
+                interaction_partners = filter_by_distance(
+                    foragersDF,
+                    f,
+                    t,
+                    interaction_length,
+                    interaction_constraint=interaction_constraint,
+                    **interaction_constraint_params,
+                )
+
+                # we want to have two different behaviors here.
+                # if no interaction partners, predictor=0 everywhere.
+                # if interaction partneres exist, but velocity values are missing,
+                # predictor = nan everywhere
+
+                if len(interaction_partners):
+                    # check if all interaction partners have valid velocity values before computing predictor
+                    v_values = foragersDF.loc[
+                        np.logical_and(
+                            foragersDF["forager"].isin(interaction_partners),
+                            foragersDF["time"] == t,
+                        ),
+                        [f"v_dt={dt}", f"theta_dt={dt}"],
+                    ]
+                    if v_values.notna().all(axis=None):
+                        x = foragers[f].loc[t, "x"]
+                        y = foragers[f].loc[t, "y"]
+                        v_pref = foragers[f].loc[t, f"v_dt={dt}"]
+
+                        # calculate theta_pref from avg of neighbors
+                        v_x = np.mean(v_values.iloc[:, 0] * np.cos(v_values.iloc[:, 1]))
+                        v_y = np.mean(v_values.iloc[:, 0] * np.sin(v_values.iloc[:, 1]))
+
+                        if v_x**2 + v_y**2:
+                            theta_pref = np.arctan2(v_y, v_x)
+                        else:
+                            theta_pref = np.nan
+
+                        predictor[f][t][predictor_name] = (
+                            _velocity_predictor_contribution(
+                                v_pref,
+                                theta_pref,
+                                x,
+                                y,
+                                predictor[f][t],
+                                sigma_v,
+                                sigma_t,
+                            )
+                        )
+
+                    else:
+                        predictor[f][t][predictor_name] = np.nan
+
+                    # normalize predictor by dividing by max
+                    max_val = predictor[f][t][predictor_name].abs().max()
+                    if max_val > 0:
+                        predictor[f][t][predictor_name] = (
+                            predictor[f][t][predictor_name] / max_val
+                        )
+
+    return predictor
+
+
 def generate_vicsek_predictor(foragers_object: dataObject, predictor_name: str):
     """
     A function that calculates the predictor scores associated with vicsek flocking,
-    by specifying an averaging transformation to `_generic_velocity_predictor`.
-    The necessary parameters from `foragers_object`. Thus, `foragers_object` must contain as attribute
+    by calling `_generate_vicsek_predictor`. The necessary parameters are taken from
+    `foragers_object`. Thus, `foragers_object` must contain as attribute
     `predictor_kwargs` : dict, with `predictor_name` as a valid key.
 
     :param foragers_object : dataObject containing positional data and necessary kwargs
     :param predictorID : Name given to column containing predictor scores in `predictor`
     :return: Nested list of calculated predictor scores, grouped by foragers and time
     """
 
-    # define transformation function
-    def transformation_vicsek(v_values):
-        v_x = np.mean(v_values.iloc[:, 0] * np.cos(v_values.iloc[:, 1]))
-        v_y = np.mean(v_values.iloc[:, 0] * np.sin(v_values.iloc[:, 1]))
-        v_transformed = pd.DataFrame([[np.sqrt(v_x**2 + v_y**2), np.arctan2(v_y, v_x)]])
-        return v_transformed
-
     # grab relevant parameters from foragers_object
     params = foragers_object.predictor_kwargs[predictor_name]
 
@@ -240,12 +343,11 @@ def transformation_vicsek(v_values):
     )
 
     # calculate predictor values
-    predictor = _generic_velocity_predictor(
+    predictor = _generate_vicsek_predictor(
         foragers_object.foragers,
         foragers_object.foragersDF,
         foragers_object.local_windows,
         predictor_name,
-        transformation_function=transformation_vicsek,
         **params,
     )
 

collab2/foraging/toolkit/visualization.py

@@ -124,6 +124,8 @@ def animate_predictors(
     Returns :
         - ani : animation
     """
+    trajectory_data = foragersDF[foragersDF["forager"].isin(forager_position_indices)]
+
     num_foragers = foragersDF["forager"].nunique()
     num_frames = foragersDF["time"].nunique()
 
@@ -146,9 +148,9 @@ def animate_predictors(
     # TODO potentially expand with forager legend
     # ax.legend()
 
-    ax.set_xticks([])
-    ax.set_yticks([])
-    ax.axis("off")
+    ax.set_xticks([0, grid_size])
+    ax.set_yticks([0, grid_size])
+    # ax.axis("off")
 
     # Initialize function to set up the background of each frame
     def init():
@@ -167,19 +169,19 @@ def init():
     # Update function for each frame
     def update(frame):
 
-        filtered_foragers = foragersDF[
-            foragersDF["forager"].isin(forager_position_indices)
-        ]
-
         # Update positions of the particles
-        current_positions = filtered_foragers.loc[
-            foragersDF["time"] == frame, ["x", "y"]
+        current_positions = trajectory_data.loc[
+            trajectory_data["time"] == frame, ["x", "y"]
         ].values
         foragers_scat.set_offsets(current_positions)
 
         # Update face and edge colors of the particles
-        foragers_scat.set_facecolor(colors)  # Set face colors directly
-        foragers_scat.set_edgecolor(colors)  # Set edge colors directly
+        foragers_scat.set_facecolor(
+            colors[forager_position_indices]
+        )  # Set face colors directly
+        foragers_scat.set_edgecolor(
+            colors[forager_position_indices]
+        )  # Set edge colors directly
 
         # Update predictor
         for i, f in enumerate(forager_predictor_indices):
@@ -201,7 +203,6 @@ def update(frame):
         return foragers_scat, *predictors_scat_list
 
     # Create the animation
-    print(num_frames)
     ani = animation.FuncAnimation(
         fig,
         update,