Improve axes layout in plot_decomposition example (#4278)

* Fix axis layout and padding for decomposition fig.

* More vertical aspect ratio.

* Remove unnecessary imports.

1 files changed, 12 insertions, 10 deletions

diff --git a/examples/algorithms/plot_decomposition.py b/examples/algorithms/plot_decomposition.py
index 7b83a718..d0f6cc6e 100644
--- a/examples/algorithms/plot_decomposition.py
+++ b/examples/algorithms/plot_decomposition.py
@@ -7,8 +7,6 @@ Example of creating a junction tree from a directed graph.
 """
 
 import networkx as nx
-from networkx.algorithms import moral
-from networkx.algorithms.tree.decomposition import junction_tree
 from networkx.drawing.nx_agraph import graphviz_layout as layout
 import matplotlib.pyplot as plt
 
@@ -20,19 +18,23 @@ B.add_edges_from(
 
 options = {"with_labels": True, "node_color": "white", "edgecolors": "blue"}
 
+fig = plt.figure(figsize=(6, 9))
+axgrid = fig.add_gridspec(3, 2)
+
 bayes_pos = layout(B, prog="neato")
-ax1 = plt.subplot(1, 3, 1)
-plt.title("Bayesian Network")
+ax1 = fig.add_subplot(axgrid[0, 0])
+ax1.set_title("Bayesian Network")
 nx.draw_networkx(B, pos=bayes_pos, **options)
 
-mg = moral.moral_graph(B)
-plt.subplot(1, 3, 2, sharex=ax1, sharey=ax1)
-plt.title("Moralized Graph")
+mg = nx.moral_graph(B)
+ax2 = fig.add_subplot(axgrid[0, 1], sharex=ax1, sharey=ax1)
+ax2.set_title("Moralized Graph")
 nx.draw_networkx(mg, pos=bayes_pos, **options)
 
-jt = junction_tree(B)
-plt.subplot(1, 3, 3)
-plt.title("Junction Tree")
+jt = nx.junction_tree(B)
+ax3 = fig.add_subplot(axgrid[1:, :])
+ax3.set_title("Junction Tree")
+ax3.margins(0.15, 0.25)
 nsize = [2000 * len(n) for n in list(jt.nodes())]
 nx.draw_networkx(jt, pos=layout(jt, prog="neato"), node_size=nsize, **options)