Merge pull request #17 from fjhheras/exp_social_context

Exposing/renaming social context

tests/socialcontext_test.py

@@ -1,9 +1,17 @@
 import numpy as np
 import pytest
+import random
 
 import trajectorytools as tt
 import trajectorytools.constants as cons
-from trajectorytools.socialcontext import in_alpha_border, in_convex_hull
+from trajectorytools.socialcontext import (
+    in_alpha_border,
+    in_convex_hull,
+    neighbour_indices,
+    adjacency_matrix,
+    neighbour_indices_in_frame,
+    adjacency_matrix_in_frame,
+)
 
 
 def test_convex_hull_vs_alpha_border():
@@ -16,3 +24,64 @@ def test_convex_hull_vs_alpha_border():
         np.logical_not(alpha_border), convex_hull
     )
     assert not np.any(in_alpha_border_not_in_convex_hull)
+
+
+@pytest.mark.parametrize("num_neighbours", [1, 3, 15])
+def test_neighbour_indices_vs_adjacency_matrix(num_neighbours):
+    t = np.load(cons.test_raw_trajectories_path, allow_pickle=True)
+    tt.interpolate_nans(t)
+
+    num_frames = t.shape[0]
+    num_individuals = t.shape[1]
+
+    nb_indices = neighbour_indices(t, num_neighbours=num_neighbours)
+    assert nb_indices.shape == tuple(
+        [num_frames, num_individuals, num_neighbours + 1]
+    )
+    adj_matrix = adjacency_matrix(t, num_neighbours=num_neighbours)
+    assert adj_matrix.shape == (num_frames, num_individuals, num_individuals)
+
+    # When there is an index in neighbour_indices output, the
+    # corresponding elment in the adjacency_matrix must be True
+    for _ in range(5):
+        frame = random.randrange(0, num_frames)
+        individual = random.randrange(0, num_individuals)
+
+        indices_neighbours = nb_indices[frame, individual, :]
+        indices_no_neighbours = [
+            i for i in range(num_individuals) if i not in indices_neighbours
+        ]
+        assert np.all(adj_matrix[frame, individual, indices_neighbours])
+        assert not np.any(adj_matrix[frame, individual, indices_no_neighbours])
+
+
+@pytest.mark.parametrize("num_neighbours", [1, 3, 15])
+def test_neighbour_indices_vs_adjacency_matrix_in_frame(num_neighbours):
+    t = np.load(cons.test_raw_trajectories_path, allow_pickle=True)
+    tt.interpolate_nans(t)
+
+    num_frames = t.shape[0]
+    num_individuals = t.shape[1]
+
+    frame = t[random.randrange(0, num_frames)]
+
+    nb_indices = neighbour_indices_in_frame(
+        frame, num_neighbours=num_neighbours
+    )
+    assert nb_indices.shape == (num_individuals, num_neighbours + 1)
+    adj_matrix = adjacency_matrix_in_frame(
+        frame, num_neighbours=num_neighbours
+    )
+    assert adj_matrix.shape == (num_individuals,) * 2
+
+    # When there is an index in neighbour_indices output, the
+    # corresponding elment in the adjacency_matrix must be True
+    for _ in range(5):
+        individual = random.randrange(0, num_individuals)
+
+        indices_neighbours = nb_indices[individual, :]
+        indices_no_neighbours = [
+            i for i in range(num_individuals) if i not in indices_neighbours
+        ]
+        assert np.all(adj_matrix[individual, indices_neighbours])
+        assert not np.any(adj_matrix[individual, indices_no_neighbours])

trajectorytools/examples/follow_animate.py

@@ -23,7 +23,7 @@ def animate_and_follow(individual=0, num_neighbours=15):
     v = tt.fixed_to_comoving(v_, e_)
     center = tt.fixed_to_comoving(-s_[:, [individual], :], e_)
 
-    indices = ttsocial.give_indices(s, num_neighbours)
+    indices = ttsocial.neighbour_indices(s, num_neighbours)
     sn = ttsocial.restrict(s, indices, individual)
     vn = ttsocial.restrict(v, indices, individual)
 

trajectorytools/socialcontext/leadership.py

@@ -63,7 +63,7 @@ def sweep_delays(data, indices, max_delay, individual=None):
 def sweep_delayed_orientation_with_neighbours(orientation, indices, max_delay):
     # Orientation: time x num_individuals x 2
     # Indices: assumed to be exclusive of own
-    # i.e. if they come from ttsocial.give_indices
+    # i.e. if they come from ttsocial.neighbour_indices
     # then the first row must be removed
     total_time_steps = orientation.shape[0] - max_delay
     sweep_delay_e = sweep_delays(orientation, indices, max_delay)

trajectorytools/socialcontext/socialcontext.py

@@ -2,6 +2,7 @@
 import scipy.spatial
 import scipy.spatial.distance as spdist
 from sklearn.neighbors import NearestNeighbors
+import warnings
 
 
 def _in_convex_hull(positions):
@@ -73,31 +74,68 @@ def in_alpha_border(positions, alpha=5):
 # LOCAL NEIGHBOURS
 
 
-def _neighbours_indices_in_frame(
-    positions, num_neighbours, adjacency=False, mode="connectivity"
-):
+def neighbour_indices_in_frame(
+    positions: np.ndarray, num_neighbours: int,
+) -> np.ndarray:
+    """ Calculate the indices of the nearest neighbours
+
+    :param positions: array of locations with dimensions
+    (individual x coordinates)
+    :param num_neighbours: number of closest neighbours requested
+    :return: output dime (individual x num_neighbours + 1)
+    """
     nbrs = NearestNeighbors(
         n_neighbors=num_neighbours + 1, algorithm="ball_tree"
     ).fit(positions)
-    if adjacency:
-        return nbrs.kneighbors_graph(positions, mode=mode).toarray()
-    else:
-        return nbrs.kneighbors(positions, return_distance=False)
+    return nbrs.kneighbors(positions, return_distance=False)
 
 
 def give_indices(positions, num_neighbours):
+    warnings.warn(
+        "give_indices to be deprecated. Use neighbour_indices instead"
+    )
+    return neighbour_indices(positions, num_neighbours)
+
+
+def neighbour_indices(
+    positions: np.ndarray, num_neighbours: int
+) -> np.ndarray:
+    """ Calculates the indices of the nearest neighbours
+
+    :param positions: array of locations with dimensions
+    (time x individual x coordinates)
+    :param num_neighbours: number of closest neighbours (does not
+    include the focal, e.g. max is individuals - 1)
+    :return: array with dims (time x individual x num_neighbours + 1)
+    """
     total_time_steps = positions.shape[0]
     individuals = positions.shape[1]
     next_neighbours = np.empty(
         [total_time_steps, individuals, num_neighbours + 1], dtype=np.int
     )
     for frame in range(total_time_steps):
-        next_neighbours[frame, ...] = _neighbours_indices_in_frame(
+        next_neighbours[frame, ...] = neighbour_indices_in_frame(
             positions[frame], num_neighbours
         )
     return next_neighbours
 
 
+def adjacency_matrix_in_frame(
+    positions: np.ndarray, num_neighbours: int, mode: str = "connectivity",
+) -> np.ndarray:
+    """
+    :param positions: array of locations with dimensions
+    (individual x coordinates)
+    :param num_neighbours: number of closest neighbours requested
+    :param mode: adjacency mode "connectivity" or "distance".     
+    :return: output has dimension (individual x individual)
+    """
+    nbrs = NearestNeighbors(
+        n_neighbors=num_neighbours + 1, algorithm="ball_tree"
+    ).fit(positions)
+    return nbrs.kneighbors_graph(positions, mode=mode).toarray()
+
+
 def adjacency_matrix(
     positions,
     num_neighbours=None,
@@ -129,8 +167,8 @@ def adjacency_matrix(
                 )
     else:
         for frame in range(total_time_steps):
-            adjacency_m[frame, ...] = _neighbours_indices_in_frame(
-                positions[frame], num_neighbours, adjacency=True, mode=mode
+            adjacency_m[frame, ...] = adjacency_matrix_in_frame(
+                positions[frame], num_neighbours, mode=mode
             )
     return adjacency_m