From f393a8e7986a72b73dfa10e2ad9a1e00e5b8ccf7 Mon Sep 17 00:00:00 2001 From: saranti Date: Sat, 27 Apr 2024 15:38:44 +1000 Subject: [PATCH] refactor fishbone diagram --- .../specialty_plots/ishikawa_diagram.py | 157 +++++++++--------- 1 file changed, 75 insertions(+), 82 deletions(-) diff --git a/galleries/examples/specialty_plots/ishikawa_diagram.py b/galleries/examples/specialty_plots/ishikawa_diagram.py index 18761ca36043..072d7b463c00 100644 --- a/galleries/examples/specialty_plots/ishikawa_diagram.py +++ b/galleries/examples/specialty_plots/ishikawa_diagram.py @@ -9,11 +9,12 @@ Source: https://en.wikipedia.org/wiki/Ishikawa_diagram """ +import math + import matplotlib.pyplot as plt from matplotlib.patches import Polygon, Wedge -# Create the fishbone diagram fig, ax = plt.subplots(figsize=(10, 6), layout='constrained') ax.set_xlim(-5, 5) ax.set_ylim(-5, 5) @@ -22,18 +23,18 @@ def problems(data: str, problem_x: float, problem_y: float, - prob_angle_x: float, prob_angle_y: float): + angle_x: float, angle_y: float): """ Draw each problem section of the Ishikawa plot. Parameters ---------- data : str - The category name. + The name of the problem category. problem_x, problem_y : float, optional The `X` and `Y` positions of the problem arrows (`Y` defaults to zero). - prob_angle_x, prob_angle_y : float, optional - The angle of the problem annotations. They are angled towards + angle_x, angle_y : float, optional + The angle of the problem annotations. They are always angled towards the tail of the plot. Returns @@ -42,8 +43,8 @@ def problems(data: str, """ ax.annotate(str.upper(data), xy=(problem_x, problem_y), - xytext=(prob_angle_x, prob_angle_y), - fontsize='10', + xytext=(angle_x, angle_y), + fontsize=10, color='white', weight='bold', xycoords='data', @@ -56,7 +57,8 @@ def problems(data: str, pad=0.8)) -def causes(data: list, cause_x: float, cause_y: float, +def causes(data: list, + cause_x: float, cause_y: float, cause_xytext=(-9, -0.3), top: bool = True): """ Place each cause to a position relative to the problems @@ -72,7 +74,9 @@ def causes(data: list, cause_x: float, cause_y: float, cause_xytext : tuple, optional Adjust to set the distance of the cause text from the problem arrow in fontsize units. - top : bool + top : bool, default: True + Determines whether the next cause annotation will be + plotted above or below the previous one. Returns ------- @@ -80,26 +84,23 @@ def causes(data: list, cause_x: float, cause_y: float, """ for index, cause in enumerate(data): - # First cause annotation is placed in the middle of the problems arrow + # [, ] + coords = [[0.02, 0], + [0.23, 0.5], + [-0.46, -1], + [0.69, 1.5], + [-0.92, -2], + [1.15, 2.5]] + + # First 'cause' annotation is placed in the middle of the 'problems' arrow # and each subsequent cause is plotted above or below it in succession. - - # [, [, ]] - coords = [[0, [0, 0]], - [0.23, [0.5, -0.5]], - [-0.46, [-1, 1]], - [0.69, [1.5, -1.5]], - [-0.92, [-2, 2]], - [1.15, [2.5, -2.5]]] - if top: - cause_y += coords[index][1][0] - else: - cause_y += coords[index][1][1] cause_x -= coords[index][0] + cause_y += coords[index][1] if top else -coords[index][1] ax.annotate(cause, xy=(cause_x, cause_y), horizontalalignment='center', xytext=cause_xytext, - fontsize='9', + fontsize=9, xycoords='data', textcoords='offset fontsize', arrowprops=dict(arrowstyle="->", @@ -108,82 +109,74 @@ def causes(data: list, cause_x: float, cause_y: float, def draw_body(data: dict): """ - Place each section in its correct place by changing + Place each problem section in its correct place by changing the coordinates on each loop. Parameters ---------- data : dict - The input data (can be list or tuple). ValueError is - raised if more than six arguments are passed. + The input data (can be a dict of lists or tuples). ValueError + is raised if more than six arguments are passed. Returns ------- None. """ - second_sections = [] - third_sections = [] - # Resize diagram to automatically scale in response to the number - # of problems in the input data. - if len(data) == 1 or len(data) == 2: - spine_length = (-2.1, 2) - head_pos = (2, 0) - tail_pos = ((-2.8, 0.8), (-2.8, -0.8), (-2.0, -0.01)) - first_section = [1.6, 0.8] - elif len(data) == 3 or len(data) == 4: - spine_length = (-3.1, 3) - head_pos = (3, 0) - tail_pos = ((-3.8, 0.8), (-3.8, -0.8), (-3.0, -0.01)) - first_section = [2.6, 1.8] - second_sections = [-0.4, -1.2] - else: # len(data) == 5 or 6 - spine_length = (-4.1, 4) - head_pos = (4, 0) - tail_pos = ((-4.8, 0.8), (-4.8, -0.8), (-4.0, -0.01)) - first_section = [3.5, 2.7] - second_sections = [1, 0.2] - third_sections = [-1.5, -2.3] - - # Change the coordinates of the annotations on each loop. + # Set the length of the spine according to the number of 'problem' categories. + length = (math.ceil(len(data) / 2)) - 1 + draw_spine(-2 - length, 2 + length) + + # Change the coordinates of the 'problem' annotations after each one is rendered. + offset = 0 + prob_section = [1.55, 0.8] for index, problem in enumerate(data.values()): - top_row = True - cause_arrow_y = 1.7 - if index % 2 != 0: # Plot problems below the spine. - top_row = False - y_prob_angle = -16 - cause_arrow_y = -1.7 - else: # Plot problems above the spine. - y_prob_angle = 16 - # Plot the 3 sections in pairs along the main spine. - if index in (0, 1): - prob_arrow_x = first_section[0] - cause_arrow_x = first_section[1] - elif index in (2, 3): - prob_arrow_x = second_sections[0] - cause_arrow_x = second_sections[1] - else: - prob_arrow_x = third_sections[0] - cause_arrow_x = third_sections[1] + plot_above = index % 2 == 0 + cause_arrow_y = 1.7 if plot_above else -1.7 + y_prob_angle = 16 if plot_above else -16 + + # Plot each section in pairs along the main spine. + prob_arrow_x = prob_section[0] + length + offset + cause_arrow_x = prob_section[1] + length + offset + if not plot_above: + offset -= 2.5 if index > 5: raise ValueError(f'Maximum number of problems is 6, you have entered ' f'{len(data)}') - # draw main spine - ax.plot(spine_length, [0, 0], color='tab:blue', linewidth=2) - # draw fish head - ax.text(head_pos[0] + 0.1, head_pos[1] - 0.05, 'PROBLEM', fontsize=10, - weight='bold', color='white') - semicircle = Wedge(head_pos, 1, 270, 90, fc='tab:blue') - ax.add_patch(semicircle) - # draw fishtail - triangle = Polygon(tail_pos, fc='tab:blue') - ax.add_patch(triangle) - # Pass each category name to the problems function as a string on each loop. problems(list(data.keys())[index], prob_arrow_x, 0, -12, y_prob_angle) - # Start the cause function with the first annotation being plotted at - # the cause_arrow_x, cause_arrow_y coordinates. - causes(problem, cause_arrow_x, cause_arrow_y, top=top_row) + causes(problem, cause_arrow_x, cause_arrow_y, top=plot_above) + + +def draw_spine(xmin: int, xmax: int): + """ + Draw main spine, head and tail. + + Parameters + ---------- + xmin : int + The default position of the head of the spine's + x-coordinate. + xmax : int + The default position of the tail of the spine's + x-coordinate. + + Returns + ------- + None. + + """ + # draw main spine + ax.plot([xmin - 0.1, xmax], [0, 0], color='tab:blue', linewidth=2) + # draw fish head + ax.text(xmax + 0.1, - 0.05, 'PROBLEM', fontsize=10, + weight='bold', color='white') + semicircle = Wedge((xmax, 0), 1, 270, 90, fc='tab:blue') + ax.add_patch(semicircle) + # draw fish tail + tail_pos = [[xmin - 0.8, 0.8], [xmin - 0.8, -0.8], [xmin, -0.01]] + triangle = Polygon(tail_pos, fc='tab:blue') + ax.add_patch(triangle) # Input data