added functions, option to save plots and added names as param

setup.py

@@ -16,7 +16,7 @@ setup(
     version="0.1",
     author="",
     author_email="",
-    install_requires=["h5py>=3", "numpy"],
+    install_requires=["h5py>=3", "numpy", "seaborn", "pandas"],
     classifiers=[
         "Development Status :: 3 - Alpha",
         "Intended Audience :: Science/Research",

src/robofish/evaluate/evaluate.py

@@ -5,7 +5,7 @@ import numpy as np
 import pandas as pd
 
 
-def evaluate_speed(paths):
+def evaluate_speed(paths, names=None, save_path=None):
 
     files_per_path = [robofish.io.read_multiple_files(p) for p in paths]
     speeds = []
@@ -19,14 +19,10 @@ def evaluate_speed(paths):
 
         speeds.append(path_speeds)
 
-    print(len(speeds))
-    max_len = max([len(speed) for speed in speeds])
+    if names is None:
+        names = paths
 
-    print(max_len)
-    print([len(speed) for speed in speeds])
-    # speeds = [speed[:max_len] for speed in speeds]
-
-    plt.hist(speeds, bins=20, label=paths, density=True, range=[0, 1])
+    plt.hist(speeds, bins=20, label=names, density=True, range=[0, 1])
     # sns.displot(speeds, label=paths, stat="probability", kind="hist")
     # sns.displot(speeds[1], label=paths[1], stat="probability", kind="hist")
     plt.title("Agent speeds")
@@ -36,10 +32,14 @@ def evaluate_speed(paths):
     # plt.xscale("log", nonpositive="clip")
     plt.legend()
     plt.tight_layout()
+
+    if save_path is None:
         plt.show()
+    else:
+        plt.savefig(save_path)
 
 
-def evaluate_turn(paths):
+def evaluate_turn(paths, names=None, save_path=None):
     files_per_path = [robofish.io.read_multiple_files(p) for p in paths]
     turns = []
     for files in files_per_path:
@@ -53,31 +53,109 @@ def evaluate_turn(paths):
                 path_turns.extend(e_turns)
         turns.append(path_turns)
 
-    plt.hist(turns, bins=40, label=paths, density=True, range=[-np.pi, np.pi])
+    if names is None:
+        names = paths
+
+    plt.hist(turns, bins=40, label=names, density=True, range=[-np.pi, np.pi])
     plt.title("Agent turns")
     plt.xlabel("Change in orientation (radians)")
     plt.ylabel("Frequency")
     plt.ticklabel_format(useOffset=False)
     plt.legend()
     # plt.tight_layout()
+
+    if save_path is None:
+        plt.show()
+    else:
+        plt.savefig(save_path)
+
+
+def evaluate_orientation(paths, names=None, save_path=None):
+    files_per_path = [robofish.io.read_multiple_files(p) for p in paths]
+    orientations = []
+    for files in files_per_path:
+        path_orientations = []
+        for p, file in files.items():
+            poses = file.get_poses_array()
+            for e_poses in poses:
+                # convert ori_x, ori_y to radians
+                ori_rad = np.arctan2(e_poses[:, 2], e_poses[:, 3])
+                path_orientations.extend(ori_rad)
+        orientations.append(path_orientations)
+
+    if names is None:
+        names = paths
+
+    plt.hist(orientations, bins=40, label=names, density=True, range=[-np.pi, np.pi])
+    plt.title("Absolute orientation")
+    plt.xlabel("orientation in radians")
+    plt.ylabel("Frequency")
+    plt.ticklabel_format(useOffset=False)
+    plt.legend()
+    # plt.tight_layout()
+
+    if save_path is None:
+        plt.show()
+    else:
+        plt.savefig(save_path)
+
+
+def evaluate_relativeOrientation(paths, names=None, save_path=None):
+    files_per_path = [robofish.io.read_multiple_files(p) for p in paths]
+    orientations = []
+    for files in files_per_path:
+        path_orientations = []
+        for p, file in files.items():
+            poses = file.get_poses_array()
+            for i in range(len(poses)):
+                for j in range(len(poses)):
+                    if i != j:
+                        ori_f1 = np.arctan2(poses[i, :, 2], poses[i, :, 3])
+                        ori_f2 = np.arctan2(poses[j, :, 2], poses[j, :, 3])
+                        ori_diff = ori_f1 - ori_f2
+                        path_orientations.extend(ori_diff)
+        orientations.append(path_orientations)
+
+    if names is None:
+        names = paths
+
+    plt.hist(orientations, bins=40, label=names, density=True, range=[-np.pi, np.pi])
+    plt.title("Relative orientation")
+    plt.xlabel("orientation in radians")
+    plt.ylabel("Frequency")
+    plt.ticklabel_format(useOffset=False)
+    plt.legend()
+    # plt.tight_layout()
+
+    if save_path is None:
         plt.show()
+    else:
+        plt.savefig(save_path)
 
 
-def evaluate_distanceToWall(paths):
+def evaluate_distanceToWall(paths, names=None, save_path=None):
+    """
+    only works for rectangular tanks
+    """
     files_per_path = [robofish.io.read_multiple_files(p) for p in paths]
     distances = []
-    # TODO: make world_bounds customizable
-    world_bounds = [-0.5, -0.5, 0.5, 0.5]
+    for files in files_per_path:
+        path_distances = []
+        for p, file in files.items():
+            poses = file.get_poses_array()
+            world_size = file.attrs["world size"]
+            world_bounds = [
+                -world_size[0] / 2,
+                -world_size[1] / 2,
+                world_size[0] / 2,
+                world_size[1] / 2,
+            ]
             wall_lines = [
                 (world_bounds[0], world_bounds[1], world_bounds[0], world_bounds[3]),
                 (world_bounds[0], world_bounds[1], world_bounds[2], world_bounds[1]),
                 (world_bounds[2], world_bounds[3], world_bounds[2], world_bounds[1]),
                 (world_bounds[2], world_bounds[3], world_bounds[0], world_bounds[3]),
             ]
-    for files in files_per_path:
-        path_distances = []
-        for p, file in files.items():
-            poses = file.get_poses_array()
             for e_poses in poses:
                 dist = []
                 for wall in wall_lines:
@@ -91,17 +169,30 @@ def evaluate_distanceToWall(paths):
                 path_distances.extend(dist)
         distances.append(path_distances)
 
-    plt.hist(distances, bins=20, label=paths, density=True, range=[0, 0.5])
+    if names is None:
+        names = paths
+
+    plt.hist(
+        distances,
+        bins=20,
+        label=names,
+        density=True,
+        range=[0, max(world_size[0] / 2, world_size[1] / 2)],
+    )
     plt.title("Distance to closest wall")
     plt.xlabel("Distance in cm")
     plt.ylabel("Frequency")
     plt.ticklabel_format(useOffset=False)
     plt.legend()
     # plt.tight_layout()
+
+    if save_path is None:
         plt.show()
+    else:
+        plt.savefig(save_path)
 
 
-def evaluate_tankpositions(paths):
+def evaluate_tankpositions(paths, save_path=None):
     """
     Heatmap of fishpositions
     By Moritz Maxeiner
@@ -109,35 +200,43 @@ def evaluate_tankpositions(paths):
     # should the plots be on the same figure, if multiple paths are given??
     files_per_path = [robofish.io.read_multiple_files(p) for p in paths]
     x_pos, y_pos = [], []
+    world_bounds = []
     for files in files_per_path:
         path_x_pos, path_y_pos = [], []
         for p, file in files.items():
             poses = file.get_poses_array()
+            world_bounds.append(file.attrs["world size"])
             for e_poses in poses:
                 path_x_pos.extend(e_poses[:, 0])
                 path_y_pos.extend(e_poses[:, 1])
         x_pos.append(np.array(path_x_pos))
         y_pos.append(np.array(path_y_pos))
 
-    for i in range(len(x_pos)):
     fig, ax = plt.subplots(figsize=(7, 7))
     ax.set_title("Tankpositions")
-        ax.set_xlim(-0.5, 0.5)
-        ax.set_ylim(-0.5, 0.5)
+    ax.set_xlim(-world_bounds[0][0] / 2, world_bounds[0][0] / 2)
+    ax.set_ylim(-world_bounds[0][1] / 2, world_bounds[0][1] / 2)
+    for i in range(len(x_pos)):
         sns.kdeplot(x_pos[i], y_pos[i] * (-1), n_levels=25, shade=True, ax=ax)
+
+    if save_path is None:
         plt.show()
+    else:
+        fig.savefig(save_path)
 
 
-def evaluate_trajectories(paths):
+def evaluate_trajectories(paths, save_path=None):
     """
     trajectories of fishes
     By Moritz Maxeiner
     """
     files_per_path = [robofish.io.read_multiple_files(p) for p in paths]
     pos = []
+    world_bounds = []
     for files in files_per_path:
         for p, file in files.items():
             poses = file.get_poses_array()
+            world_bounds.append(file.attrs["world size"])
             path_pos = {
                 fish: pd.DataFrame(
                     {
@@ -161,8 +260,8 @@ def evaluate_trajectories(paths):
         sns.scatterplot(
             x="x", y="y", hue="Agent", linewidth=0, s=4, data=pos[i][1], ax=ax
         )
-        ax.set_xlim(-0.5, 0.5)
-        ax.set_ylim(-0.5, 0.5)
+        ax.set_xlim(-world_bounds[i][0] / 2, world_bounds[i][0] / 2)
+        ax.set_ylim(-world_bounds[i][1] / 2, world_bounds[i][1] / 2)
         ax.invert_yaxis()
         ax.xaxis.set_ticks_position("top")
         ax.xaxis.set_label_position("top")
@@ -185,10 +284,14 @@ def evaluate_trajectories(paths):
             label="End",
         )
         ax.legend()
+
+        if save_path is None:
             plt.show()
+        else:
+            fig.savefig(save_path[i])
 
 
-def evaluate_positionVec(paths):
+def evaluate_positionVec(paths, save_path=None):
     files_per_path = [robofish.io.read_multiple_files(p) for p in paths]
     posVec = []
     for files in files_per_path:
@@ -210,13 +313,95 @@ def evaluate_positionVec(paths):
             posVec.append(np.concatenate(path_posVec, axis=0))
 
     for i in range(len(posVec)):
-        # TODO: maybe find better way to plot?
+        # TODO: maybe find better way to plot? maybe limit x and y axis to some value?
         # fig, ax = plt.subplots(figsize=(8, 8))
         # sns.kdeplot(posVec[i][:, 0], posVec[i][:, 1], n_levels=25, shade=True, ax=ax)
         # plt.hist2d(posVec[i][:, 0], posVec[i][:, 1], bins=100)
         df = pd.DataFrame({"x": posVec[i][:, 0], "y": posVec[i][:, 1]})
         sns.displot(df, x="x", y="y", binwidth=(0.03, 0.03), cbar=True)
+
+        if save_path is None:
             plt.show()
+        else:
+            plt.savefig(save_path[i])
+
+
+def evaluate_follow_iid(paths, save_path=None):
+    files_per_path = [robofish.io.read_multiple_files(p) for p in paths]
+    follow, iid = [], []
+    for files in files_per_path:
+        for p, file in files.items():
+            poses = file.get_poses_array()
+            for i in range(len(poses)):
+                for j in range(i + 1, len(poses)):
+                    iid.append(calculate_iid(poses[i, :-1, 0:2], poses[j, :-1, 0:2]))
+                    iid.append(iid[-1])
+
+                    follow.append(calculate_follow(poses[i, :, 0:2], poses[j, :, 0:2]))
+                    follow.append(calculate_follow(poses[j, :, 0:2], poses[i, :, 0:2]))
+
+    follow_iid_data = pd.DataFrame(
+        {
+            "IID [m]": np.concatenate(iid, axis=0),
+            "Follow": np.concatenate(follow, axis=0),
+        }
+    )
+
+    grid = sns.jointplot(
+        x="IID [m]", y="Follow", data=follow_iid_data, linewidth=0, s=1, kind="scatter"
+    )
+    grid.fig.set_figwidth(9)
+    grid.fig.set_figheight(6)
+    grid.fig.subplots_adjust(top=0.9)
+
+    if save_path is None:
+        plt.show()
+    else:
+        grid.fig.savefig(save_path)
+
+
+def evaluate_all(paths, names=None, save_folder=None):
+    evaluate_speed(paths, names, save_folder + "speed.png")
+    evaluate_turn(paths, names, save_folder + "turn.png")
+    evaluate_orientation(paths, names, save_folder + "orientation.png")
+    evaluate_relativeOrientation(paths, names, save_folder + "relativeOrientation.png")
+    evaluate_distanceToWall(paths, names, save_folder + "distanceToWall.png")
+    evaluate_follow_iid(paths, save_folder + "follow_iid.png")
+    evaluate_tankpositions(paths, save_folder + "tankpositions.png")
+    evaluate_trajectories(
+        paths,
+        [save_folder + "trajectories_" + str(i) + ".png" for i in range(len(paths))],
+    )
+    evaluate_positionVec(
+        paths,
+        [save_folder + "positionVec_" + str(i) + ".png" for i in range(len(paths))],
+    )
+
+
+def calculate_follow(a, b):
+    """
+    Given two series of poses - with X and Y coordinates of their positions as the first two elements -
+    return the follow metric from the first to the second series.
+    """
+    a_v = a[1:, :2] - a[:-1, :2]
+    b_p = normalize_series(b[:-1, :2] - a[:-1, :2])
+    return (a_v * b_p).sum(axis=-1)
+
+
+def calculate_iid(a, b):
+    """
+    Given two series of poses - with X and Y coordinates of their positions as the first two elements -
+    return the inter-individual distance (between the positions).
+    """
+    return np.linalg.norm(b[:, :2] - a[:, :2], axis=-1)
+
+
+def normalize_series(x):
+    """
+    Given a series of vectors, return a series of normalized vectors.
+    Null vectors are mapped to `NaN` vectors.
+    """
+    return (x.T / np.linalg.norm(x, axis=-1)).T
 
 
 def calculate_posVec(data, angle):
@@ -262,7 +447,7 @@ def calculate_distLinePoint(x1, y1, x2, y2, points):
     return dist
 
 
-evaluate_positionVec(
+evaluate_distanceToWall(
     [
         [
             "I:/Code/BachelorThesis/BachelorThesis/Fish/Guppy/io/DQN_22_12_2020_02_1models_det.hdf5"