Skip to content

Commit

Permalink
Merge pull request #28144 from saranti/refactor_fishbone
Browse files Browse the repository at this point in the history 
DOC: Refactor code in the fishbone diagram example
  • Loading branch information
@timhoffm
timhoffm committed May 2, 2024
2 parents 2ba8545 + f393a8e commit e253aa2
Showing 1 changed file with 75 additions and 82 deletions.
157 changes: 75 additions & 82 deletions galleries/examples/specialty_plots/ishikawa_diagram.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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
Expand All @@ -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',
Expand All @@ -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
Expand All @@ -72,34 +74,33 @@ 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
-------
None.
"""
for index, cause in enumerate(data):
# First cause annotation is placed in the middle of the problems arrow
# [<x pos>, <y pos>]
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.

# [<x pos>, [<y pos top>, <y pos bottom>]]
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="->",
Expand All @@ -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
Expand Down

0 comments on commit e253aa2

Please sign in to comment.