diff --git a/src/plotting/tracker_plot.py b/src/plotting/tracker_plot.py index e6502ea87dafe3dca1d79ea66cfe2879e2198aa4..e63e2988c5b3ed216cc73ddb359d9e7eb7618932 100644 --- a/src/plotting/tracker_plot.py +++ b/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())