Commit 3db719e7 authored by Amaury BARRAL's avatar Amaury BARRAL

angle PDF

parent 2deff4cf
......@@ -10,6 +10,7 @@ import numpy as np
from matplotlib import pyplot as plt
from matplotlib.gridspec import GridSpec
from matplotlib.patches import Polygon
from scipy.stats.kde import gaussian_kde
import God.SaveAndLoad as SaveAndLoad
......@@ -47,6 +48,7 @@ def get_layout(option: str = "default") -> Tuple[Any, Dict[str, Any], Dict[str,
subplots["speed"] = fig.add_subplot(gs[0, 3])
subplots["angle"] = subplots["speed"].twinx()
subplots["polar"] = fig.add_subplot(gs[0, 4], projection='polar')
subplots["angle_pdf"] = fig.add_subplot(gs[0, 5])
subplots["correlations"] = fig.add_subplot(gs[1, 3])
subplots["correlation_length"] = subplots["correlations"].twinx()
subplots["evolution_group_size"] = fig.add_subplot(gs[1, 4])
......@@ -69,6 +71,7 @@ drawable_to_data = {
"avg_speed": ["avg_speed"],
"avg_angle": ["avg_angle"],
"avg_polar": ["avg_speed", "avg_angle"],
"angle_pdf": ["_simulation"],
"correlations": ["correlations"],
"correlations_fit": ["correlations_fit"],
"correlation_length": ["correlations_fit"],
......@@ -175,6 +178,7 @@ class Visualiser:
self.to_draw["avg_speed"] = "avg_speed" in to_draw
self.to_draw["avg_angle"] = "avg_angle" in to_draw
self.to_draw["avg_polar"] = "avg_polar" in to_draw
self.to_draw["angle_pdf"] = "angle_pdf" in to_draw
self.to_draw["correlations"] = "correlations" in to_draw
self.to_draw["correlations_fit"] = "correlations_fit" in to_draw
self.to_draw["correlation_length"] = "correlation_length" in to_draw
......@@ -291,6 +295,11 @@ class Visualiser:
color="red")
ax_polar.set_rmax(vel_scaling)
if self.to_draw["angle_pdf"]:
ax_angle_pdf = self.subplots["angle_pdf"]
ax_angle_pdf.set_title("Angle PDF")
self.layout_artists["angle_pdf"], = ax_angle_pdf.plot([], [], lw=2)
if self.to_draw["evolution_group_size"]:
ax_speed = self.subplots["evolution_group_size"]
ax_speed.set_xlim(self.t_start, self.t_end)
......@@ -374,6 +383,16 @@ class Visualiser:
self.layout_artists["avg_polar"].set_data(angles, speeds)
self.layout_artists["avg_polar_endpoint"].set_data([angles[-1]], [speeds[-1]])
def plot_angle_pdf(self, frame_num: int) -> None:
"""
Plot angle PDF
"""
frame = self.processed_data["_simulation"]["frames"][frame_num]
angles = np.array([np.array(bird[1]) for bird in frame])
kde = gaussian_kde(angles)
dist_space = np.linspace(min(angles), max(angles), 100)
self.layout_artists["angle_pdf"].set_data(dist_space, kde(dist_space))
def plot_correlations(self, frame_num: int) -> None:
"""
Plot correlations
......@@ -627,6 +646,9 @@ class Visualiser:
if self.to_draw["avg_polar"]:
self.plot_avg_polar(frame_number)
if self.to_draw["angle_pdf"]:
self.plot_angle_pdf(frame_number)
if self.to_draw["correlations"]:
self.plot_correlations(frame_number)
......
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment