def plot_distributions(questions: List['LikertQuestion'],
order: Optional[List[str]] = None,
sort_by: Union[str, Callable, list] = None,
ax: Optional[Axes] = None) -> Axes:
"""
Plot a stacked bar chart of Likert values for each question.
:param questions: Questions to plot distributions of
:param order: Optional ordering for the Likert values
:param sort_by: Column name(s) and/or lambda function(s) to sort by,
e.g. lambda d: d['1'] + d['2']
:param ax: Optional matplotlib axes to plot on
"""
data = [question.data for question in questions]
texts = [question.text for question in questions]
common_leading = find_common_leading(texts)
texts = trim_common_leading(texts)
texts = trim_common_trailing(texts)
for q, question in enumerate(questions):
data[q].name = texts[q]
ax = ax or new_axes()
ax = plot_categorical_distributions(data=data,
order=order,
sort_by=sort_by,
ax=ax)
if common_leading:
ax.set_title(common_leading)
ax.set_ylabel('# Respondents')
return ax
def plot_parameters():
ax = new_axes()
dist.prior().plot(x=x, color='r', ax=ax)
dist.posterior().plot(x=x, color='g', ax=ax)
ax.legend()
plt.show()
def draw(self, ax: Optional[Axes]) -> Axes:
"""
Draw the network.
:param ax: Optional matplotlib Axes.
"""
ax = ax or new_axes()
pos = self._get_layout()
draw_networkx_nodes(
G=self._graph,
pos=pos,
ax=ax,
nodelist=self._graph.nodes(),
alpha=0.5,
)
draw_networkx_labels(
G=self._graph,
pos=pos,
ax=ax,
labels=self._task_name_dict,
font_size=10,
)
draw_networkx_edges(G=self._graph,
pos=pos,
label='weight',
ax=ax,
edge_color='gray')
draw_networkx_edge_labels(G=self._graph,
pos=pos,
edge_labels=self._edge_weight_dict())
return ax
def plot_wikipedia_pmfs():
"""
https://en.wikipedia.org/wiki/Poisson_distribution#/media/
File:Poisson_pmf.svg
"""
ax = new_axes(width=10, height=10)
Poisson(lambda_=1).plot(k=k_wikipedia,
kind='line',
color='gray',
mfc='orange',
marker='o',
ax=ax)
Poisson(lambda_=4).plot(k=k_wikipedia,
kind='line',
color='gray',
mfc='purple',
marker='o',
ax=ax)
Poisson(lambda_=10).plot(k=k_wikipedia,
kind='line',
color='gray',
mfc='lightblue',
marker='o',
ax=ax)
ax.set_ylim(0, 0.4)
ax.set_title('Probability mass function')
ax.legend(loc='upper right')
plt.show()
def plot(self,
k: Union[Iterable[Iterable], ndarray],
ax: Optional[Axes] = None,
**kwargs) -> Axes:
"""
Plot the most probable values of the function.
:param k: Range of values of x to plot p(x) over.
:param ax: Optional matplotlib axes to plot on.
:param kwargs: Additional arguments for bar plot.
"""
k = array(k)
if k.ndim != 2:
raise ValueError('k must have 2 dimensions: [num_points, K]')
data = self.at(k)
# filter to "top" most likely
data = data.rename('p').reset_index()
new_data = []
for name in self._parent.names:
new_frame = data[[name, 'p']].rename(columns={name: 'x'})
new_frame['$x_k$'] = name
new_data.append(new_frame)
data = concat(new_data)
ax = ax or new_axes()
lineplot(data=data, x='x', y='p', hue='$x_k$', ax=ax, **kwargs)
ax.set_xlabel('X')
ax.set_ylabel(f'{self._name}(X)')
ax.legend(loc='upper right')
return ax
def plot_2d(self,
x1: Union[Iterable, ndarray],
x2: Union[Iterable, ndarray],
color_map: str = 'viridis', ax: Optional[Axes] = None,
**kwargs) -> Axes:
"""
Plot a 2-dimensional function as a grid heat-map.
N.B. don't use for distributions where calculating the function for the
full range of x1 and x2 values would cause an error e.g. for a Dirichlet
where x1 + x2 must equal 1.
:param x1: Range of values of x1 to plot p(x1, x2) over.
:param x2: Range of values of x2 to plot p(x1, x2) over.
:param color_map: Optional colormap for the heat-map.
:param ax: Optional matplotlib axes to plot on.
:param kwargs: Additional arguments for contourf method.
"""
x1_grid, x2_grid = meshgrid(x1, x2)
x1_x2 = dstack((x1_grid, x2_grid))
f = self._method(x1_x2)
ax = ax or new_axes()
ax.contourf(x1_grid, x2_grid, f, cmap=color_map, **kwargs)
ax.set_xlabel('x1')
ax.set_ylabel('x2')
return ax
def plot_categorical_distributions(data: List[Series], order: List[str] = None,
sort_by: Union[str, Callable, list] = None,
ax: Optional[Axes] = None) -> Axes:
"""
Plot a comparison of several categorical distributions
e.g. a list of LikertQuestion responses.
:param data: List of response data to Categorical Questions
:param order: Optional ordering for the Categorical values
:param sort_by: Column name(s) and/or lambda function(s) to sort by
e.g. lambda d: d['1'] + d['2']
:param ax: Optional matplotlib axes to plot on.
"""
joined = concat(data, axis=1)
counts = joined.apply(value_counts).fillna(0).T
if order:
counts = counts[order]
if sort_by is None:
sort_by = order
counts = sort_data_frame_values(counts, sort_by, ascending=False)
counts.index = wrap_text(counts.index)
ax = ax or new_axes()
counts.plot(kind='bar', stacked=True, ax=ax)
percentages = counts.astype(float)
percentages = 100 * (percentages.T / percentages.T.sum()).T
for label_x in range(percentages.shape[0]):
y_base = 0
for count_y in range(percentages.shape[1]):
y = y_base + counts.iloc[label_x, count_y] / 2
text = f'{percentages.iloc[label_x, count_y]:.1f}%'
ax.text(x=label_x, y=y, s=text, ha='center', va='center')
y_base += counts.iloc[label_x, count_y]
return ax
def plot(
self,
x: Iterable,
kind: str = 'line',
colors: Optional[List[str]] = None,
ax: Optional[Axes] = None,
**kwargs
) -> Axes:
"""
Plot the marginal distribution of each component.
:param x: Range of values of x to plot p(x) over.
:param kind: Kind of plot e.g. 'bar', 'line'.
:param colors: Optional list of colors for each series.
:param ax: Optional matplotlib axes to plot on.
:param kwargs: Additional arguments for the matplotlib plot function.
"""
parent = self._parent
if colors is None:
colors = [f'C{i}' for i in range(len(parent.names))]
if len(colors) != len(parent.names):
raise ValueError(f'Pass 0 or {len(parent.names)} colors.')
ax = ax or new_axes()
for k, color in zip(parent.names, colors):
data = getattr(parent[k], self._method_name)().at(x)
data.plot(x=x, kind=kind, color=color,
ax=ax, label=f'{k}', **kwargs)
ax.legend()
ax.set_xlabel(parent.x_label)
ax.set_ylabel(f'{self._name}({parent.x_label})')
return ax
def plot_distribution(self,
item: Union[str, RespondentAttribute, Question],
transpose: bool = False,
ax: Optional[Axes] = None) -> Axes:
"""
Plot distribution of answers to a question or values of an attribute.
:param item: The question or attribute to plot the distribution of.
:param transpose: Whether to transpose the labels to the y-axis.
:param ax: Optional matplotlib axes to plot on.
"""
item = self._find_item(item)
ax = ax or new_axes()
item.plot_distribution(transpose=transpose, ax=ax)
# find the item if the name was passed
if isinstance(item, str):
item = self._find_item(item)
# plot the distribution if the item is valid
if (isinstance(item, RespondentAttribute)
or isinstance(item, Question)):
item.plot_distribution(transpose=transpose, ax=ax)
return ax
else:
raise TypeError('Cannot plot this kind of question or attribute.')
def plot_distribution(self,
data: Optional[Series] = None,
transpose: bool = False,
bins: Optional[Bins] = None,
color: str = 'C0',
pct_size: int = None,
grid: bool = False,
title: Optional[str] = None,
x_label: Optional[str] = None,
y_label: Optional[str] = None,
ax: Optional[Axes] = None) -> Axes:
"""
Plot a histogram of the distribution of the response data.
:param data: The answers given by Respondents to the Question.
:param transpose: True to plot horizontally.
:param bins: Value for hist bins. Leave as None for integer bins.
:param color: Color or list of colors for the bars.
:param pct_size: Font size for the percent markers.
:param grid: Whether to show a plot grid or not.
:param title: Optional title for the plot.
:param x_label: Label for the x-axis.
:param y_label: Label for the y-axis.
:param ax: Optional matplotlib axes to plot on.
"""
ax = ax or new_axes()
data = data if data is not None else self._data
if data is None:
raise ValueError('No data!')
orientation = 'horizontal' if transpose else 'vertical'
bins = bins or self._default_hist(data)
data.plot(kind='hist',
ax=ax,
bins=bins,
orientation=orientation,
color=color)
# add percentages
hist, edges = histogram(data, bins=bins)
item_counts = Series(index=0.5 * (edges[1:] + edges[:-1]), data=hist)
item_pcts = 100 * item_counts / item_counts.sum()
label_bar_plot_pcts(item_counts=item_counts,
item_pcts=item_pcts,
ax=ax,
transpose=transpose,
font_size=pct_size)
# add titles and grid
ax.set_title(self.text)
AxesFormatter(ax).set_text(
title=title, x_label=x_label,
y_label=y_label).set_axis_below(True).grid(grid)
if transpose and not x_label:
ax.set_xlabel('# Respondents')
elif not transpose and not y_label:
ax.set_ylabel('# Respondents')
return ax
def plot_wikipedia_cdfs():
ax = new_axes()
InverseGamma(alpha=1, beta=1).cdf().plot(x=x, color='red', ax=ax)
InverseGamma(alpha=2, beta=1).cdf().plot(x=x, color='green', ax=ax)
InverseGamma(alpha=3, beta=1).cdf().plot(x=x, color='blue', ax=ax)
InverseGamma(alpha=3, beta=0.5).cdf().plot(x=x, color='cyan', ax=ax)
ax.legend()
plt.show()
def plot_predictions():
ax = new_axes()
predicted = dist.rvs(100000)
ax.hist(predicted, bins=100, density=True, label='PPD samples')
x_actual = arange(predicted.min(), predicted.max(), 0.01)
actual = norm(loc=mu, scale=sigma).pdf(x_actual)
ax.plot(x_actual, actual, label='True Distribution')
ax.legend()
plt.show()
def plot_distribution(self,
data: Optional[Series] = None,
transpose: bool = False,
ax: Optional[Axes] = None,
rug: Optional[bool] = True,
kde: Optional[bool] = True,
hist: Optional[bool] = False,
max_pct: float = 1.0,
**kwargs) -> Axes:
"""
Plot the distribution of answers to the Question.
:param data: Optional response data. Required if the Question does not
already have any.
:param transpose: True to plot vertically.
:param ax: Optional matplotlib axes to plot on.
"""
data = data if data is not None else self._data
if data is None:
raise ValueError('No data!')
ax = ax or new_axes()
a = data.dropna()
if max_pct < 1:
a = a.loc[a <= a.quantile(max_pct)]
distplot(a=a,
rug=rug,
kde=kde,
hist=hist,
vertical=transpose,
ax=ax,
**kwargs)
min_val = a.min()
max_val = a.max()
# add titles
ax.set_title(self.text)
if transpose:
ax.set_xlabel('# Respondents')
ax.set_ylabel(self.name)
else:
ax.set_xlabel(self.name)
ax.set_ylabel('# Respondents')
# format axes
if transpose:
ax.set_xticks([])
if min_val <= max_val / 10:
ax.set_ylim(0, ax.get_ylim()[1])
else:
ax.set_yticks([])
if min_val <= max_val / 10:
ax.set_xlim(0, ax.get_xlim()[1])
return ax
def plot_wikipedia_cdfs():
"""
https://en.wikipedia.org/wiki/Lomax_distribution#/media/File:LomaxCDF.png
"""
ax = new_axes()
Lomax(lambda_=1, alpha=2).cdf().plot(x=x, color='blue', ax=ax)
Lomax(lambda_=2, alpha=2).cdf().plot(x=x, color='green', ax=ax)
Lomax(lambda_=4, alpha=1).cdf().plot(x=x, color='red', ax=ax)
Lomax(lambda_=6, alpha=1).cdf().plot(x=x, color='orange', ax=ax)
ax.legend()
plt.show()
def plot_wikipedia_pdfs():
"""
https://en.wikipedia.org/wiki/Laplace_distribution#/media/
File:Laplace_pdf_mod.svg
"""
ax = new_axes()
Laplace(mu=0, b=1).plot(x=x, color='red', ax=ax)
Laplace(mu=0, b=2).plot(x=x, color='black', ax=ax)
Laplace(mu=0, b=4).plot(x=x, color='blue', ax=ax)
Laplace(mu=-5, b=4).plot(x=x, color='green', ax=ax)
ax.legend(loc='upper right')
plt.show()
def plot_wikipedia_pdfs():
"""
https://en.wikipedia.org/wiki/Exponential_distribution#/media/
File:Exponential_probability_density.svg
"""
ax = new_axes()
Exponential(lambda_=0.5).plot(x=x, color='orange', ax=ax)
Exponential(lambda_=1).plot(x=x, color='purple', ax=ax)
Exponential(lambda_=1.5).plot(x=x, color='lightblue', ax=ax)
ax.legend()
ax.set_ylim(0, 1.5)
plt.show()
def plot_wikipedia_cdfs():
"""
https://en.wikipedia.org/wiki/Student%27s_t-distribution#/media/
File:Student_t_cdf.svg
"""
ax = new_axes()
StudentsT(nu=1).cdf().plot(x=x_wikipedia_1, color='orange', ax=ax)
StudentsT(nu=2).cdf().plot(x=x_wikipedia_1, color='purple', ax=ax)
StudentsT(nu=5).cdf().plot(x=x_wikipedia_1, color='lightblue', ax=ax)
StudentsT(nu=1e9).cdf().plot(x=x_wikipedia_1, color='black', ax=ax)
ax.legend(loc='lower right')
plt.show()
def plot_wikipedia_pdfs():
"""
https://en.wikipedia.org/wiki/Normal_distribution#/media/
File:Normal_Distribution_PDF.svg
"""
ax = new_axes()
Normal(mu=0, sigma_sq=0.2).plot(x=x, color='blue', ax=ax)
Normal(mu=0, sigma_sq=1.0).plot(x=x, color='red', ax=ax)
Normal(mu=0, sigma_sq=5.0).plot(x=x, color='orange', ax=ax)
Normal(mu=-2, sigma_sq=0.5).plot(x=x, color='green', ax=ax)
ax.legend(loc='upper right')
plt.show()
def plot_wikipedia_cdfs():
"""
https://en.wikipedia.org/wiki/PERT_distribution#/media/
File:PERT_cdf_examples.jpg
"""
ax = new_axes(width=10, height=10)
PERT(0, 10, 100).cdf().plot(x=x, color='blue', ax=ax)
PERT(0, 50, 100).cdf().plot(x=x, color='orange', ax=ax)
PERT(0, 70, 100).cdf().plot(x=x, color='gray', ax=ax)
ax.set_ylim(0, 1)
ax.set_title('Cumulative density function')
ax.legend(loc='upper center')
plt.show()
def plot_wikipedia_cdfs():
"""
https://en.wikipedia.org/wiki/Beta_distribution#/media/
File:Beta_distribution_cdf.svg
"""
ax = new_axes(width=10, height=10)
Beta(0.5, 0.5).cdf().plot(x=x, color='red', ax=ax)
Beta(5, 1).cdf().plot(x=x, color='blue', ax=ax)
Beta(1, 3).cdf().plot(x=x, color='green', ax=ax)
Beta(2, 2).cdf().plot(x=x, color='purple', ax=ax)
Beta(2, 5).cdf().plot(x=x, color='orange', ax=ax)
ax.set_title('Cumulative distribution function')
ax.legend(loc='upper left')
plt.show()
def plot_cpd(survey_data: DataFrame, probability: str, condition: str,
transpose: bool = False,
set_title: bool = True, legend: bool = True,
x_label: bool = True, y_label: bool = True,
x_tick_labels: bool = True, y_tick_labels: bool = True,
ax: Optional[Axes] = None) -> Axes:
"""
Plot a conditional probability distribution as a kde for each condition
category.
:param survey_data: The DataFrame containing the survey data.
:param probability: The question or attribute to find probability of.
:param condition: The question or attribute to condition on.
:param transpose: Set to True to put values along y-axis.
:param set_title: Whether to add a title to the plot.
:param legend: Whether to show a legend or not.
:param x_label: Whether to show the default label on the x-axis or not,
or string of text for the label.
:param y_label: Whether to show the default label on the y-axis or not,
or string of text for the label.
:param x_tick_labels: Whether to show labels on the ticks on the x-axis.
:param y_tick_labels: Whether to show labels on the ticks on the y-axis.
:param ax: Optional matplotlib axes to plot on.
"""
condition_data = survey_data[condition]
prob_data = survey_data[probability].dropna()
ax = ax or new_axes()
for category in condition_data.unique():
distplot(
a=prob_data.loc[condition_data == category],
kde=True, rug=True, hist=False,
ax=ax, label=category, vertical=transpose
)
ax.legend(title=condition).set_visible(legend)
# label axes
set_cpd_axes_labels(ax=ax, x_label=x_label, y_label=y_label,
prob_name=probability, transpose=transpose)
set_cp_tick_labels(ax, x_tick_labels, y_tick_labels)
# axes limits
if transpose:
if prob_data.min() <= prob_data.max() / 10:
ax.set_ylim(0, ax.get_ylim()[1])
else:
if prob_data.min() <= prob_data.max() / 10:
ax.set_xlim(0, ax.get_xlim()[1])
# title
if set_title:
ax.set_title(f'p({probability}|{condition})')
return ax
def plot_wikipedia_pdfs():
"""
https://en.wikipedia.org/wiki/Beta_distribution#/media/
File:Beta_distribution_pdf.svg
"""
ax = new_axes(width=10, height=10)
Beta(0.5, 0.5).pdf().plot(x=x, color='red', mean=True, std=True, ax=ax)
Beta(5, 1).pdf().plot(x=x, color='blue', mean=True, ax=ax)
Beta(1, 3).pdf().plot(x=x, color='green', mean=True, ax=ax)
Beta(2, 2).pdf().plot(x=x, color='purple', mean=True, ax=ax)
Beta(2, 5).pdf().plot(x=x, color='orange', mean=True, ax=ax)
ax.set_ylim(0, 2.5)
ax.set_title('Probability density function')
ax.legend(loc='upper center')
plt.show()
def plot_simplex(self, num_contours: int = 100, num_sub_div: int = 8,
color_map: str = 'viridis', border: bool = True,
ax: Optional[Axes] = None,
**kwargs) -> Axes:
"""
Plot a 3-dimensional functions as a simplex heat-map.
:param num_contours: The number of levels of contours to plot.
:param num_sub_div: Number of recursive subdivisions to create.
:param color_map: Optional colormap for the plot.
:param border: Whether to plot a border around the simplex heat-map.
:param ax: Optional matplotlib axes to plot on.
:param kwargs: Additional arguments for tricontourf method.
"""
corners = array([[0, 0], [1, 0], [0.5, 0.75 ** 0.5]])
triangle = Triangulation(corners[:, 0], corners[:, 1])
mid_points = [
(corners[(i + 1) % 3] + corners[(i + 2) % 3]) / 2
for i in range(3)
]
def to_barycentric(cartesian):
"""
Converts 2D Cartesian to barycentric coordinates.
:param cartesian: A length-2 sequence containing the x and y value.
"""
s = [(corners[i] - mid_points[i]).dot(
cartesian - mid_points[i]
) / 0.75
for i in range(3)]
s_clipped = clip(a=s, a_min=0, a_max=1)
return s_clipped / norm(s_clipped, ord=1)
refiner = UniformTriRefiner(triangle)
tri_mesh = refiner.refine_triangulation(subdiv=num_sub_div)
f = [self._method(to_barycentric(xy))
for xy in zip(tri_mesh.x, tri_mesh.y)]
ax = ax or new_axes()
ax.tricontourf(tri_mesh, f, num_contours, cmap=color_map, **kwargs)
ax.set_aspect('equal')
ax.set_xlim(0, 1)
ax.set_ylim(0, 0.75 ** 0.5)
ax.set_axis_off()
if border:
ax.triplot(triangle, linewidth=1)
return ax
def plot_distribution(self,
data: Optional[Series] = None,
transpose: bool = False,
top: int = 25,
title: bool = True,
x_label: bool = True,
y_label: bool = True,
ax: Optional[Axes] = None) -> Axes:
"""
Plot the distribution of top words in answers to the Question.
:param data: The answers given by Respondents to the Question.
:param transpose: Whether to transpose the labels to the y-axis.
:param top: Number of most frequent words to plot.
:param title: Whether to add a title to the plot.
:param x_label: Whether to add a label to the x-axis.
:param y_label: Whether to add a label to the y-axis.
:param ax: Optional matplotlib axes to plot on.
"""
data = data if data is not None else self._data
if data is None:
raise ValueError('No data!')
words = pre_process_text_series(data)
value_counts = Series(words).value_counts()[:top]
plot_type = 'barh' if transpose else 'bar'
ax = ax or new_axes()
value_counts.index = wrap_text(value_counts.index)
value_counts.plot(kind=plot_type, ax=ax)
if title:
ax.set_title(self.text)
if transpose:
x_label_value = '# Respondents'
y_label_value = data.name
else:
x_label_value = data.name
y_label_value = '# Respondents'
if x_label:
ax.set_xlabel(x_label_value)
else:
ax.set_xlabel('')
if y_label:
ax.set_ylabel(y_label_value)
else:
ax.set_ylabel('')
return ax
def plot_comparison(self, ax: Axes = None, **kwargs) -> Axes:
"""
Plot a comparison between the different question ratings.
:param ax: Optional matplotlib axes.
"""
ax = ax or new_axes()
if 'cmap' not in kwargs:
kwargs['cmap'] = 'Blues'
data = DataFrame(
{k: q.value_counts()
for k, q in self.item_dict.items()})
heatmap(data=data, ax=ax, annot=True, fmt='d', **kwargs)
AxesFormatter(ax).set_text(x_label='Question',
y_label='Rating').invert_y_axis()
draw_vertical_dividers(ax)
return ax
def plot_wikipedia_cdfs():
"""
https://en.wikipedia.org/wiki/Binomial_distribution#/media/
File:Binomial_distribution_cdf.svg
"""
ax = new_axes()
Binomial(n=20, p=0.5).cdf().plot(
k=k_wiki, kind='line', color='blue', ax=ax
)
Binomial(n=20, p=0.7).cdf().plot(
k=k_wiki, kind='line', color='lightgreen', ax=ax
)
Binomial(n=40, p=0.5).cdf().plot(
k=k_wiki, kind='line', color='red', ax=ax
)
ax.set_ylim(0, 1.0)
ax.set_title('Cumulative distribution function')
ax.legend(loc='upper right')
plt.show()
def plot_top(self, top: int, ax: Optional[Axes] = None, **kwargs) -> Axes:
"""
Plot the probability of the most likely values of the distribution.
:param top: The number of top values to plot probability for.
:param ax: Axes to plot on.
:param kwargs: Additional arguments to pass to bar plot method.
"""
k = array(self._parent.permutations())
if k.ndim != 2:
raise ValueError('k must have 2 dimensions: [num_points, K]')
data = self.at(k)
data = data.sort_values(ascending=False).head(top)
ax = ax or new_axes()
data.plot.bar(label=str(self._parent), ax=ax, **kwargs)
ax.set_xlabel('K')
ax.set_ylabel(f'{self._name}(K)')
ax.legend(loc='upper right')
return ax
def plot_wikipedia_pmfs():
"""
https://en.wikipedia.org/wiki/Binomial_distribution#/media/
File:Binomial_distribution_pmf.svg
"""
ax = new_axes()
Binomial(n=20, p=0.5).plot(
k=k_wiki, kind='line', color='blue', ax=ax, ls='', marker='d',
mean=True
)
Binomial(n=20, p=0.7).plot(
k=k_wiki, kind='line', color='lightgreen', ax=ax, ls='', marker='s'
)
Binomial(n=40, p=0.5).plot(
k=k_wiki, kind='line', color='red', ax=ax, ls='', marker='o'
)
ax.set_ylim(0, 0.25)
ax.set_title('Probability mass function')
ax.legend(loc='upper right')
plt.show()
def draw(self,
node_labels: Optional[str] = None,
ax: Optional[Axes] = None) -> Axes:
"""
Draw the DecisionTree.
:param node_labels: One of {'name', 'amount', None}
:param ax: Optional matplotlib axes.
"""
ax = ax or new_axes()
pos = self._get_layout()
for nodes, node_shape, node_color in zip(
(self.decision_nodes(), self.chance_nodes(), self.amount_nodes()),
('s', 'o', 'H'), ('r', 'b', 'g')):
draw_networkx_nodes(
G=self._graph,
pos=pos,
ax=ax,
nodelist=nodes,
node_shape=node_shape,
node_color=node_color,
alpha=0.5,
)
if node_labels is not None:
if node_labels == 'name':
labels = self.node_names_dict()
elif node_labels == 'amount':
labels = self.node_amounts_dict()
else:
raise ValueError()
draw_networkx_labels(
G=self._graph,
pos=pos,
ax=ax,
labels=labels,
font_size=10,
)
draw_networkx_edges(G=self._graph, pos=pos, ax=ax, edge_color='gray')
return ax
def plot_results(
self, x: str = 'effect_mean', y: str = 'p_superior',
split_by: str = 'answers_given', ax: Optional[Axes] = None
) -> Axes:
"""
Create a scatter plot of the experimental results.
:param x: Name of the attribute for the x-axis.
:param y: Name of the attribute for the y-axis.
:param split_by: Name of the attribute to split into different series.
:param ax: Optional matplotlib axes to plot on.
"""
data = self.results_data(group_values=True)
ax = ax or new_axes()
for split_val in data[split_by].unique():
split_data = data.loc[data[split_by] == split_val]
ax.scatter(split_data[x], split_data[y], label=split_val)
ax.set_xlabel(x)
ax.set_ylabel(y)
ax.legend()
return ax
