some docs and refac on TrackPlot

src/plotting/tracker_plot.py

@@ -226,12 +226,14 @@ class TrackPlot:
         self.x_max = None
 
         track_chain_obj = TrackChain(tracker_list)
-        scopes = track_chain_obj.scopes
-        dims = track_chain_obj.dims
         track_chain_dict = track_chain_obj.create_track_chain()
-        self.set_ypos_anchor(scopes, dims)
+        self.set_ypos_anchor(track_chain_obj.scopes, track_chain_obj.dims)
         self.fig, self.ax = plt.subplots(figsize=(len(tracker_list) * 2, (self.anchor.max() - self.anchor.min()) / 3))
-        stages, v_lines = self.create_track_chain_plot(track_chain_dict, sparse_conn_mode=sparse_conn_mode, skip_run_env=skip_run_env)
+        self._plot(track_chain_dict, sparse_conn_mode, skip_run_env, plot_folder)
+
+    def _plot(self, track_chain_dict, sparse_conn_mode, skip_run_env, plot_folder):
+        stages, v_lines = self.create_track_chain_plot(track_chain_dict, sparse_conn_mode=sparse_conn_mode,
+                                                       skip_run_env=skip_run_env)
         self.set_lims()
         self.add_variable_names()
         self.add_stages(v_lines, stages)
@@ -240,40 +242,43 @@ class TrackPlot:
         plt.savefig(plot_name, dpi=600)
 
     def line(self, start_x, end_x, y, color="darkgrey"):
+        """Draw grey horizontal connection line from start_x to end_x on y-pos."""
+        # draw white border line
         l = mlines.Line2D([start_x + self.width, end_x], [y + self.height / 2, y + self.height / 2], color="white",
                           linewidth=2.5)
         self.ax.add_line(l)
+        # draw grey line
         l = mlines.Line2D([start_x + self.width, end_x], [y + self.height / 2, y + self.height / 2], color=color,
                           linewidth=1.4)
         self.ax.add_line(l)
 
     def step(self, start_x, end_x, start_y, end_y, color="black"):
+        """Draw black connection step line from start_xy to end_xy. Step is taken shortly before end position."""
+        # adjust start and end by width height
         start_x += self.width
         start_y += self.height / 2
         end_y += self.height / 2
-        step_x = end_x - (self.space_intern_x) / 2
+        step_x = end_x - (self.space_intern_x) / 2  # step is taken shortly before end
         pos_x = [start_x, step_x, step_x, end_x]
-        # pos_x = [start_x, (start_x + end_x) / 2, (start_x + end_x) / 2, end_x]
         pos_y = [start_y, start_y, end_y, end_y]
+        # draw white border line
         l = mlines.Line2D(pos_x, pos_y, color="white", linewidth=2.5)
         self.ax.add_line(l)
+        # draw black line
         l = mlines.Line2D(pos_x, pos_y, color=color, linewidth=1.4)
         self.ax.add_line(l)
 
     def rect(self, x, y, method="get"):
-
-        if method == "get":
-            color = "orange"
-        else:
-            color = "lightblue"
+        """Draw rectangle with lower left at (x,y), size equal to width/height and label/color according to method."""
+        # draw rectangle
+        color = {"get": "orange"}.get(method, "lightblue")
         r = Rectangle((x, y), self.width, self.height, color=color, label=color)
-
         self.ax.add_artist(r)
+        # add label
         rx, ry = r.get_xy()
         cx = rx + r.get_width() / 2.0
         cy = ry + r.get_height() / 2.0
-        self.ax.annotate(method, (cx, cy), color='w', weight='bold',
-                    fontsize=6, ha='center', va='center')
+        self.ax.annotate(method, (cx, cy), color='w', weight='bold', fontsize=6, ha='center', va='center')
 
     def set_ypos_anchor(self, scopes, dims):
         anchor = sum(dims.values())