logging and plot improbements

Author: FedeClaudi

.vscode/settings.json

@@ -1,4 +1,4 @@
 {
-    "python.pythonPath": "C:\\Users\\Federico\\.conda\\envs\\dev\\python.exe",
+    "python.pythonPath": "/Users/federicoclaudi/miniconda3/envs/dev/bin/python",
     "cSpell.enabled": true
 }

proj/animation/animate.py

@@ -4,14 +4,15 @@
     save_videocap_to_video,
 )
 import click
+import logging
 
 
 def animate_from_images(folder, savepath, fps):
     cap = get_cap_from_images_folder(folder, img_format="%2d.png")
     save_videocap_to_video(cap, savepath, ".mp4", fps=fps)
 
     gifpath = savepath.replace(".mp4", ".gif")
-    print(
+    logging.info(
         "To save the video as GIF, use: \n"
         + f'ffmpeg -i "{savepath}" -f gif "{gifpath}"'
     )

proj/environment/environment.py

@@ -1,4 +1,6 @@
 import numpy as np
+import logging
+
 from fcutils.maths.geometry import calc_distance_between_points_2d
 
 from proj.environment.trajectories import (
@@ -90,6 +92,7 @@ def plan(self, curr_x, g_traj, itern):
             self.moved_to_next.append(self.itern)
 
         self.curr_traj_waypoint_idx = min_idx + n_ahead
+        self.model.curr_traj_waypoint_idx = self.curr_traj_waypoint_idx
         self.current_traj_waypoint = g_traj[min_idx, :]
 
         start = min_idx + n_ahead
@@ -120,6 +123,9 @@ def isdone(self, curr_x, trajectory):
             dist = curr_x.r
 
         if dist <= self.model.trajectory["min_dist"]:
+            logging.info(
+                f"Reached the end of the threshold as we're within {self.model.trajectory['min_dist']} from the end"
+            )
             return True
         else:
             return False

proj/environment/manager.py

@@ -3,6 +3,7 @@
 import pandas as pd
 from fancylog import fancylog
 import logging
+
 from rich.logging import RichHandler
 
 from fcutils.file_io.io import save_yaml
@@ -56,6 +57,9 @@ def __init__(self, model, winstor=False):
         for fld in folders:
             fld.mkdir(exist_ok=True)
 
+        self._start_logging()
+
+    def _start_logging(self):
         # Start logging
         fancylog.start_logging(
             output_dir=str(self.datafolder),
@@ -67,16 +71,24 @@ def __init__(self, model, winstor=False):
             file_log_level="INFO",
         )
 
-        print(f"Saving data at: {self.datafolder}")
-        logging.info("\n\n\n\n")
-        logging.info(f"Saving data at: {self.datafolder}")
-
+        # log main folder
         log = logging.getLogger("rich")
+        log.setLevel(logging.INFO)
+
         log.info(
             f"[bold green] Saving data at: {self.datafolder}",
             extra={"markup": True},
         )
 
+        # log config.py
+        logging.info("=" * 20)
+        logging.info("=" * 20)
+        with open("proj/model/config.py") as f:
+            conf = f.read()
+        logging.info(conf)
+        logging.info("=" * 20)
+        logging.info("=" * 20)
+
     def _save_results(self):
         # save config
         save_yaml(

proj/environment/plotter.py

@@ -1,8 +1,10 @@
 import matplotlib.pyplot as plt
+import numpy as np
+import logging
+
 from fcutils.plotting.utils import clean_axes
 from fcutils.plotting.colors import desaturate_color
 from fcutils.plotting.plot_elements import plot_line_outlined
-import numpy as np
 
 from proj.utils import salmon
 from proj.animation import variables_colors as colors
@@ -13,7 +15,9 @@ def press(event, self):
         Deals with key press during interactive visualizatoin
     """
     if event.key == "c":
-        print("stopping")
+        logging.info(
+            "Stopping because user manually terminated simulation (presed C)"
+        )
         self.stop = True
 
 

proj/environment/trajectories.py

@@ -78,6 +78,8 @@ def compute_trajectory_stats(trajectory, duration, planning_params):
 
     # log stuff
     log = logging.getLogger("rich")
+    log.setLevel(logging.INFO)
+
     log.info(
         f"""
         n_points = {n_points}
@@ -201,19 +203,24 @@ def from_tracking(n_steps, params, planning_params, cache_fld, *args):
     speed = line_smoother(trial.body_speed) * trial.fps
     ang_speed = np.ones_like(speed)  # it will be ignored
 
+    # get start frame
+    from_start = calc_distance_from_point(np.array([x, y]), [x[0], y[0]])
+    start = np.where(from_start > params["dist_th"])[0][0]
+
+    if start >= len(x):
+        logging.error("Bad trajectory")
+        raise ValueError("Bad trajectory")
+
     # resample variables so that samples are uniformly distributed
     vars = dict(x=x, y=y, angle=angle, speed=speed, ang_speed=ang_speed)
     if params["resample"]:
         for k, v in vars.items():
-            vars[k] = _interpol(v, params["max_deg_interpol"], n_steps)
-
-    # get start frame
-    from_start = calc_distance_from_point(np.array([x, y]), [x[0], y[0]])
-    start = np.where(from_start > params["dist_th"])[0][0]
+            vars[k] = _interpol(
+                v[start:-1], params["max_deg_interpol"], n_steps
+            )
 
-    # stack and cut
+    # stack
     trajectory = np.vstack(vars.values()).T
-    trajectory = trajectory[start:-1, :]
 
     return compute_trajectory_stats(
         trajectory, len(x) / trial.fps, planning_params

proj/model/config.py

@@ -43,7 +43,7 @@ class Config:
         max_speed=100,
         min_speed=80,
         min_dist=100,  # if agent is within this distance from trajectory end the goal is considered achieved
-        dist_th=100,  # keep frames only after moved away from start location
+        dist_th=300,  # keep frames only after moved away from start location
         resample=True,  # if True when using tracking trajectory resamples it
         max_deg_interpol=8,  # if using track fit a N degree polynomial to daa to smoothen
         randomize=True,  # if true when using tracking it pulls a random trial

proj/model/model.py

@@ -84,13 +84,18 @@ def reset(self):
             tau_l=[],
             r=[],
             gamma=[],
+            trajectory_idx=[],
         )
 
     def _append_history(self):
         for ntuple in [self.curr_x, self.curr_control]:
             for k, v in ntuple._asdict().items():
                 self.history[k].append(v)
 
+        self.history["trajectory_idx"].append(
+            self.curr_traj_waypoint_idx
+        )  # this is updated by env.plan
+
     def _make_simbols(self):
         # state variables
         x, y, theta, thetadot = symbols("x, y, theta, thetadot", real=True)

proj/plotting/results.py

@@ -1,5 +1,6 @@
 import matplotlib.pyplot as plt
 import numpy as np
+import logging
 
 from fcutils.plotting.utils import clean_axes, save_figure
 from fcutils.plotting.colors import desaturate_color
@@ -102,9 +103,8 @@ def _plot_v(history, trajectory, plot_every, ax=None):
     )
 
     # plot history speed
-    x = np.linspace(0, len(trajectory), len(v))
     ax.plot(
-        x, v, color=colors["v"], lw=2, zorder=100,
+        history["trajectory_idx"], v, color=colors["v"], lw=2, zorder=100,
     )
 
 
@@ -113,9 +113,9 @@ def _plot_accel(history, ax=None):
     vdot = derivative(v)
     omegadot = derivative(omega)
 
-    plot_line_outlined(ax, vdot, lw=2, color=colors["v"], label="$v$")
+    plot_line_outlined(ax, vdot, lw=2, color=colors["v"], label="$\dot{v}$")
     plot_line_outlined(
-        ax, omegadot, lw=2, color=colors["omega"], label="$\omega$"
+        ax, omegadot, lw=2, color=colors["omega"], label="$\dot{\omega}$"
     )
     ax.legend()
 
@@ -150,3 +150,4 @@ def plot_results(results_folder, plot_every=20, save_path=None):
 
     if save_path is not None:
         save_figure(f, str(save_path))
+        logging.info(f"Saved summary figure at: {save_path}")

proj/run/runner.py

@@ -86,11 +86,11 @@ def run_experiment(
                 )
 
             # Check if we're done
-            if (
-                environment.isdone(model.curr_x, trajectory)
-                or environment.stop
-            ):
-                print(f"Reached end of trajectory after {itern} steps")
+            if environment.isdone(model.curr_x, trajectory):
+                log.info("environment says we're DONE")
+                break
+            if environment.stop:
+                log.info("environment says STOP")
                 break
 
     log.info(f"Terminated after {itern} iterations.")