refactor central park birds in collab 2

.gitignore

@@ -150,7 +150,6 @@ tmp_files/
 # folder to temporarily store files
 data/user_data/*
 
-
 **/xmemo/
 
 data/communicators/communicators_strong/*
@@ -171,3 +170,7 @@ data/foraging/locust/ds/priors_sam30_15EQ20191202_s0_e10.pkl
 docs/foraging/random-hungry-followers/sim_data/hungry_sim.dill
 docs/foraging/random-hungry-followers/sim_data/hungry_foragers_samples.dill
 docs/foraging/random-hungry-followers/sim_data/hungry_foragers_samples.dill
+docs/foraging/central-park-birds/ducks_proximity_single_optimizer.pkl
+docs/foraging/central-park-birds/ducks_proximity_optimizer.pkl
+docs/foraging/central-park-birds/sparrows_proximity_single_optimizer.pkl
+docs/foraging/central-park-birds/ducks_access_optimizer.pkl

collab2/foraging/toolkit/__init__.py

@@ -25,7 +25,12 @@
 from .filtering import constraint_filter_nearest, filter_by_distance  # noqa: F401
 from .food import generate_food_predictor  # noqa: F401
 from .inference import run_svi_inference  # type: ignore # noqa: F401
-from .inference import get_samples, prep_data_for_inference, summary  # noqa: F401
+from .inference import (  # noqa: F401
+    get_samples,
+    prep_data_for_inference,
+    prep_DF_data_for_inference,
+    summary,
+)
 from .local_windows import (  # noqa: F401
     _generate_local_windows,
     _get_grid,
@@ -69,4 +74,7 @@
 )
 from .visualization import animate_predictors, plot_predictor  # noqa: F401
 
+# from .waic import compute_waic  # noqa: F401
+
+
 # from .trace import rewards_to_trace, rewards_trace  # noqa: F401

collab2/foraging/toolkit/animate_foragers.py

@@ -11,7 +11,7 @@
 warnings.simplefilter(action="ignore", category=FutureWarning)
 
 
-def plot_trajectories(df, title):
+def plot_trajectories(df, title, legend=True):
     unique_foragers = df["forager"].unique()
     plt.figure()
 
@@ -41,7 +41,8 @@ def plot_trajectories(df, title):
     plt.axis("equal")
     plt.gca().invert_yaxis()
     plt.axis("off")
-    plt.legend()
+    if legend:
+        plt.legend()
     plt.suptitle(f"Trajectories: {title}", fontsize=16)
     return plt
 

collab2/foraging/toolkit/inference.py

@@ -40,6 +40,44 @@ def prep_data_for_inference(
     return predictor_tensors, outcome_tensors
 
 
+def prep_DF_data_for_inference(
+    DF, predictors: List[str], outcome_vars: str, subsample_rate: float = 1.0
+) -> Tuple[Dict[str, torch.Tensor], Dict[str, torch.Tensor]]:
+
+    if isinstance(outcome_vars, str):
+        outcome_list = [outcome_vars]
+    else:
+        outcome_list = outcome_vars
+
+    df = DF[predictors + outcome_list].copy()
+
+    # assert no NaNs in df
+    assert df.notna().all().all(), "Dataframe contains NaN values"
+
+    # Apply subsampling
+    if subsample_rate < 1.0:
+        df = df.sample(frac=subsample_rate).reset_index(drop=True)
+
+    # Apply subsampling
+    if subsample_rate < 1.0:
+        df = df.sample(frac=subsample_rate).reset_index(drop=True)
+
+    predictor_tensors = {
+        key: torch.tensor(df[key].values, dtype=torch.float32) for key in predictors
+    }
+    outcome_tensors = {
+        key: torch.tensor(df[key].values, dtype=torch.float32) for key in outcome_list
+    }
+
+    # print size
+    print("Sample size:", len(df))
+
+    # print size
+    print("Sample size:", len(df))
+
+    return predictor_tensors, outcome_tensors
+
+
 def summary(samples, sites):
     site_stats = {}
     for site_name, values in samples.items():
@@ -79,7 +117,7 @@ def run_svi_inference(
         loss.backward()
         losses.append(loss.item())
         adam.step()
-        if (step % 50 == 0) or (step == 1) & verbose:
+        if (step % 200 == 0) or (step == 1) & verbose:
             print("[iteration %04d] loss: %.4f" % (step, loss))
 
     if plot:
@@ -97,13 +135,20 @@ def get_samples(
     outcome,
     num_svi_iters,
     num_samples,
+    plot=True,
+    verbose=True,
 ):
 
     logging.info(f"Starting SVI inference with {num_svi_iters} iterations.")
     start_time = time.time()
     pyro.clear_param_store()
     guide = run_svi_inference(
-        model, n_steps=num_svi_iters, predictors=predictors, outcome=outcome
+        model,
+        n_steps=num_svi_iters,
+        predictors=predictors,
+        outcome=outcome,
+        plot=plot,
+        verbose=verbose,
     )
     end_time = time.time()
     elapsed_time = end_time - start_time
@@ -118,18 +163,21 @@ def get_samples(
         for k, v in predictive(predictors, outcome).items()
         if k != "obs"
     }
-    print(samples.keys())
 
     sites = [
         key
         for key in samples.keys()
         if (key.startswith("weight") and not key.endswith("sigma"))
     ]
-    print(sites)
 
     print("Coefficient marginals:")
     for site, values in summary(samples, sites).items():
         print("Site: {}".format(site))
         print(values, "\n")
 
-    return {"samples": samples, "guide": guide, "predictive": predictive}
+    return {
+        "samples": samples,
+        "guide": guide,
+        "predictive": predictive,
+        "summaries": summary(samples, sites),
+    }

collab2/foraging/toolkit/next_step_score.py

@@ -54,9 +54,13 @@ def _generate_nextStep_score(
                     score[f][t]["distance_to_next_step"] = np.nan
                     score[f][t][score_name] = np.nan
 
-        # save nans for last frame
-        score[f][num_frames - 1]["distance_to_next_step"] = np.nan
-        score[f][num_frames - 1][score_name] = np.nan
+        # save nans for last fram
+        if (
+            score[f][num_frames - 1] is not None
+        ):  # conditioning needed as last frame is not always present
+            # as in cp birds data
+            score[f][num_frames - 1]["distance_to_next_step"] = np.nan
+            score[f][num_frames - 1][score_name] = np.nan
 
     return score
 

collab2/foraging/toolkit/proximity.py

@@ -42,7 +42,10 @@ def _piecewise_proximity_function(
 
     result = np.where(
         cond1,
-        np.sin(np.pi / (2 * repulsion_radius) * (distance + 3 * repulsion_radius)),
+        np.sin(
+            np.pi / ((2 * repulsion_radius) * (distance + 3 * repulsion_radius))
+            + 0.0001
+        ),  # division by zero errors
         np.where(
             cond2,
             np.sin(

collab2/foraging/toolkit/utils.py

@@ -94,6 +94,21 @@ def __init__(
         self.derived_quantities: dict[str, List[List[pd.DataFrame]]] = {}
         self.derivedDF: pd.DataFrame
 
+    def calculate_step_distances(self):
+        step_distances = []
+        for f in range(len(self.foragers)):
+            data = self.foragers[f].dropna()
+            unique_t = data["time"].unique()
+            for t in unique_t:
+                if (t - 1) in unique_t:
+                    step_ago = data[data["time"] == t - 1]
+                    xdiff = data[data["time"] == t]["x"].values - step_ago["x"].values
+                    ydiff = data[data["time"] == t]["y"].values - step_ago["y"].values
+                    step_distance = np.sqrt(xdiff**2 + ydiff**2)
+                    step_distances.append(step_distance)
+
+        return np.concatenate(step_distances)
+
     def calculate_step_size_max(self):
         step_maxes = []
 
@@ -203,6 +218,8 @@ def distances_and_peaks(distances, bins=40, x_min=None, x_max=None):
             color="red",
         )
 
+    plt.show()
+
 
 # remove rewards eaten by foragers in proximity
 def update_rewards(sim, rewards, foragers, start=1, end=None):

collab2/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(

collab2/foraging/toolkit/waic.py

@@ -0,0 +1,69 @@
+from typing import Any, Callable, Dict, Optional
+
+import pyro
+import torch
+from pyro.infer.enum import get_importance_trace
+
+
+def compute_waic(
+    model: Callable[..., Any],
+    guide: Callable[..., Any],
+    num_particles: int,
+    max_plate_nesting: int,
+    sites: Optional[list[str]] = None,
+    *args: Any,
+    **kwargs: Any
+) -> Dict[str, Any]:
+
+    def vectorize(fn: Callable[..., Any]) -> Callable[..., Any]:
+        def _fn(*args: Any, **kwargs: Any) -> Any:
+            with pyro.plate(
+                "num_particles_vectorized", num_particles, dim=-max_plate_nesting
+            ):
+                return fn(*args, **kwargs)
+
+        return _fn
+
+    model_trace, _ = get_importance_trace(
+        "flat", max_plate_nesting, vectorize(model), vectorize(guide), args, kwargs
+    )
+
+    def site_filter_is_observed(site_name: str) -> bool:
+        return model_trace.nodes[site_name]["is_observed"]
+
+    def site_filter_in_sites(site_name: str) -> bool:
+        return sites is not None and site_name in sites
+
+    if sites is None:
+        site_filter = site_filter_is_observed
+    else:
+        site_filter = site_filter_in_sites
+
+    observed_nodes = {
+        name: node for name, node in model_trace.nodes.items() if site_filter(name)
+    }
+
+    log_p_post = {
+        key: observed_nodes[key]["log_prob"].mean(dim=0)  # sum(axis = 0)/num_particles
+        for key in observed_nodes.keys()
+    }
+
+    lppd = torch.stack([log_p_post[key] for key in log_p_post.keys()]).sum()
+
+    var_log_p_post = {
+        key: (observed_nodes[key]["log_prob"]).var(axis=0)
+        for key in observed_nodes.keys()
+    }
+
+    p_waic = torch.stack([var_log_p_post[key] for key in var_log_p_post.keys()]).sum()
+
+    waic = -2 * (lppd - p_waic)
+
+    return {
+        "waic": waic,
+        "nodes": observed_nodes,
+        "log_p_post": log_p_post,
+        "var_log_p_post": var_log_p_post,
+        "lppd": lppd,
+        "p_waic": p_waic,
+    }