Adding an exponential version of nextStep_score

collab/foraging/toolkit/__init__.py

@@ -43,8 +43,8 @@
     continuous_contribution,
 )
 from .next_step_score import (  # noqa: F401
-    _generate_nextStep_score,
     generate_nextStep_score,
+    generate_nextStepExponential_score,
 )
 from .proximity import (  # noqa: F401; foragers_to_forager_distances,
     generate_proximity_predictor,

collab/foraging/toolkit/next_step_score.py

@@ -92,3 +92,84 @@ def generate_nextStep_score(foragers_object: dataObject, score_name: str):
     return _generate_nextStep_score(
         foragers_object.foragers, foragers_object.local_windows, score_name, **params
     )
+
+
+def _generate_nextStepExponential_score(
+    foragers: List[pd.DataFrame],
+    local_windows: List[List[pd.DataFrame]],
+    score_name: str,
+    decay_length: Optional[float] = 1.0,
+):
+    """
+    A function that computes a score for how far grid points are from the next position of a forager,
+    using an exponential decay.
+    If the next position of the forager is unavailable, nan values are assigned to the scores.
+
+    :param foragers: list of DataFrames containing forager positions, grouped by forager index
+    :param local_windows:  Nested list of DataFrames containing grid points to compute predictor over,
+        grouped by forager index and time
+    :param score_name: name of column to save the calculated nextStep scores under
+    :param decay_length: length scale of exponential decay
+
+    :return: Nested list of calculated scores, grouped by foragers and time
+    """
+
+    num_foragers = len(foragers)
+    num_frames = len(foragers[0])
+    score = copy.deepcopy(local_windows)
+
+    for f in range(num_foragers):
+        for t in range(num_frames - 1):
+            if score[f][t] is not None:
+                x_new = foragers[f].at[t + 1, "x"]
+                y_new = foragers[f].at[t + 1, "y"]
+                if np.isfinite(x_new) and np.isfinite(y_new):
+                    distance_to_next_step = np.sqrt(
+                        (score[f][t]["x"] - x_new) ** 2
+                        + (score[f][t]["y"] - y_new) ** 2
+                    )
+
+                    score[f][t][score_name] = np.exp(
+                        -distance_to_next_step / decay_length
+                    )
+
+                    # rescale over the grid
+                    score[f][t][score_name] = (
+                        score[f][t][score_name] - score[f][t][score_name].min()
+                    ) / (score[f][t][score_name].max() - score[f][t][score_name].min())
+                else:
+                    score[f][t][score_name] = np.nan
+
+        # save nans for last frame
+        if score[f][num_frames - 1] is not None:
+            score[f][num_frames - 1][score_name] = np.nan
+
+    return score
+
+
+def generate_nextStepExponential_score(foragers_object: dataObject, score_name: str):
+    """
+    A wrapper function that computes nextStepExponential score only taking `foragers_object` as argument,
+    and calling `_generate_nextStepExponential_score` under the hood.
+
+    The next step score computes a score for how far grid points are from the next position of a forager,
+    using an exponential decay.
+    If the next position of the forager is unavailable, nan values are assigned to the scores.
+
+    The formula for the score is:
+        next_step_score(i,t,x,y) = exp(-d_nextStep/decay_length)
+        where d_nextStep is the distance of (x,y) grid point from the position of forager i
+        at time t+1
+
+    :param foragers_object: dataObject containing positional data, local_windows, score_kwargs
+    :param score_name : name of column to save the calculated nextStep scores under
+
+    :return: Nested list of calculated scores, grouped by foragers and time
+
+    Keyword arguments:
+        :param decay_length: length scale of exponential decay
+    """
+    params = foragers_object.score_kwargs[score_name]
+    return _generate_nextStepExponential_score(
+        foragers_object.foragers, foragers_object.local_windows, score_name, **params
+    )

collab/foraging/toolkit/visualization.py

@@ -95,7 +95,9 @@ def plot_predictor(
             for c in range(len(time) % ncols, ncols):
                 fig.delaxes(axes[c])
 
+    fig.suptitle(f"Predictor: {predictor_name}")
     fig.tight_layout(pad=2)
+    fig.show()
 
 
 def animate_predictors(