Repository: DistrictDataLabs/yellowbrick

Issue:
Poof() should return ax
Make sure that ll visualizers return the `self.ax` when calling poof(). This is so that users can always get access to the ax, for example when working inside a function in a notebook. In addition it will encourage the behavior pattern of tweaking the plot if desired after poof(). We will be adjusting our documentation to use this form to also encourage this behavior:

```
ax = viz.poof()
```

This why this behavior pattern is useful:
https://stackoverflow.com/questions/47450804/yellowbrick-increasing-font-size-on-yellowbrick-generated-charts



Discussion:
As a workaround, once poof is called, you can access ax through:

`ax = viz.ax`

and fig:

`fig = viz.ax.get_figure()`
@fdion out of curiosity, we have a handle to the current ax via the visualizer, do you think we should also add one to the figure? I have to say, it's very rare that I'm doing anything with the figure, though we do have the ability to modify the size of the figure on the visualizer. 
@bbengfort good question, Ben. The primary use is to save to file (pdf for latex), but since that's already covered by yellowbrick, and you currently don't have multi ax figures, probably not a high demand item :)

I return the figure in stemgraphic because I have grids of heatmaps or other plots and stuff like that. If you think faceting / multivariate might be a thing soon in Yellow Brick, then I'd say it would be useful to provide it.

Francois
Turns out we just included a VisualizerGrid that draws each visualizer on different axes; we're also thinking about trying to do a "flip book" style visualization with multiple figures. Guess I'll have to keep this in mind! Thanks! 
Just a quick implementation note on the original ticket. To implement this:
Add the line `return self.ax` at the end of the base visualizer's poof() method in `yellowbrick/base.py`, and then  searching through all the visualizers in the code to make sure no one else overrides poof. 

Stumbled upon this issue when looking at #380.
Started the implementation in https://github.com/DistrictDataLabs/yellowbrick/pull/422, still work in progress at the moment.
fdion's way of getting `fig` happens to currently be the only way for me to display `Rank2D` correlation plots, as `poof()` does not work anymore, even within the same cell. Not sure if this qualifies as a new bug, or is covered here somehow.
@michael-ziedalski - I'm a bit concerned that `poof` isn't working for you - what environment are you in? Are you still at PyData? Stop by the District Data Labs table and I'd be happy to take a look at it with you. 

Resolving patch:
diff --git a/yellowbrick/base.py b/yellowbrick/base.py
--- a/yellowbrick/base.py
+++ b/yellowbrick/base.py
@@ -248,6 +248,9 @@ def poof(self, outpath=None, clear_figure=False, **kwargs):
         if clear_figure:
             self.fig.clear()
 
+        # Return ax to ensure display in notebooks
+        return self.ax
+
     ## ////////////////////////////////////////////////////////////////////
     ## Helper Functions
     ## ////////////////////////////////////////////////////////////////////
@@ -558,3 +561,6 @@ def poof(self, outpath=None, clear_figure=False, **kwargs):
 
         if clear_figure:
             plt.gcf().clear()
+
+        # Return Axes array to ensure poof works in notebooks
+        return self.axarr
diff --git a/yellowbrick/pipeline.py b/yellowbrick/pipeline.py
--- a/yellowbrick/pipeline.py
+++ b/yellowbrick/pipeline.py
@@ -90,13 +90,18 @@ def poof(self, outdir=None, ext=".pdf", **kwargs):
         kwargs : dict
             Keyword arguments to pass to the ``poof()`` method of all steps.
         """
+        axes = []
         for name, step in self.visual_steps.items():
             if outdir is not None:
                 outpath = path.join(outdir, slugify(name) + ext)
             else:
                 outpath = None
 
-            step.poof(outpath=outpath, **kwargs)
+            ax = step.poof(outpath=outpath, **kwargs)
+            axes.append(ax)
+
+        # Return axes array to ensure figures are shown in notebook
+        return axes
 
     def fit_transform_poof(self, X, y=None, outpath=None, **kwargs):
         """


Test patch:
diff --git a/tests/test_base.py b/tests/test_base.py
--- a/tests/test_base.py
+++ b/tests/test_base.py
@@ -124,7 +124,7 @@ class CustomVisualizer(Visualizer):
         _, ax = plt.subplots()
         viz = CustomVisualizer(ax=ax)
         viz.finalize = MagicMock()
-        viz.poof()
+        assert viz.poof() is ax
 
         viz.finalize.assert_called_once_with()
         mock_plt.show.assert_called_once_with()
@@ -142,7 +142,7 @@ class CustomVisualizer(Visualizer):
         _, ax = plt.subplots()
         viz = CustomVisualizer(ax=ax)
         viz.finalize = MagicMock()
-        viz.poof(outpath="test.png")
+        assert viz.poof(outpath="test.png") is ax
 
         viz.finalize.assert_called_once_with()
         mock_plt.show.assert_not_called()
@@ -159,7 +159,7 @@ class CustomVisualizer(Visualizer):
 
         with pytest.warns(YellowbrickWarning):
             viz = CustomVisualizer()
-            viz.poof()
+            assert viz.poof() is not None
 
 ##########################################################################
 ## ScoreVisualizer Cases
@@ -228,7 +228,8 @@ def test_draw_visualizer_grid(self):
         grid = VisualizerGrid(visualizers)
 
         grid.fit(X, y)
-        grid.poof()  # poof is required here (do not replace with finalize)!
+        # poof is required here (do not replace with finalize)!
+        assert grid.poof() is not None
 
         self.assert_images_similar(grid)
 
@@ -249,7 +250,8 @@ def test_draw_with_rows(self):
         grid = VisualizerGrid(visualizers, nrows=2)
 
         grid.fit(X, y)
-        grid.poof()  # poof is required here (do not replace with finalize)!
+        # poof is required here (do not replace with finalize)!
+        assert grid.poof() is not None
 
         self.assert_images_similar(grid)
 
@@ -270,7 +272,8 @@ def test_draw_with_cols(self):
         grid = VisualizerGrid(visualizers, ncols=2)
 
         grid.fit(X, y)
-        grid.poof()  # poof is required here (do not replace with finalize)!
+        # poof is required here (do not replace with finalize)!
+        assert grid.poof() is not None
 
         self.assert_images_similar(grid)