Merge pull request #119 from VincentRouvreau/persistence_graphical_tool_improvements

Persistence graphical tool improvements

14 files changed, 146 insertions, 101 deletions

diff --git a/src/common/include/gudhi/reader_utils.h b/src/common/include/gudhi/reader_utils.h
index 98335552..db31bf5c 100644
--- a/src/common/include/gudhi/reader_utils.h
+++ b/src/common/include/gudhi/reader_utils.h
@@ -293,6 +293,9 @@ Note: the function does not check that birth <= death.
 **/
 template <typename OutputIterator>
 void read_persistence_intervals_and_dimension(std::string const& filename, OutputIterator out) {
+#ifdef DEBUG_TRACES
+  std::cout << "read_persistence_intervals_and_dimension - " << filename << std::endl;
+#endif  // DEBUG_TRACES
   std::ifstream in(filename);
   if (!in.is_open()) {
     std::string error_str("read_persistence_intervals_and_dimension - Unable to open file ");
@@ -307,6 +310,13 @@ void read_persistence_intervals_and_dimension(std::string const& filename, Outpu
     if (line.length() != 0 && line[0] != '#') {
       double numbers[4];
       int n = sscanf(line.c_str(), "%lf %lf %lf %lf", &numbers[0], &numbers[1], &numbers[2], &numbers[3]);
+#ifdef DEBUG_TRACES
+      std::cout << "[" << n << "] = ";
+      for (int i = 0; i < n; i++) {
+        std::cout << numbers[i] << ",";
+      }
+      std::cout << std::endl;
+#endif  // DEBUG_TRACES
       if (n >= 2) {
         int dim = (n >= 3 ? static_cast<int>(numbers[n - 3]) : -1);
         *out++ = std::make_tuple(dim, numbers[n - 2], numbers[n - 1]);
diff --git a/src/python/doc/persistence_graphical_tools_user.rst b/src/python/doc/persistence_graphical_tools_user.rst
index b2124fdd..2de99252 100644
--- a/src/python/doc/persistence_graphical_tools_user.rst
+++ b/src/python/doc/persistence_graphical_tools_user.rst
@@ -20,6 +20,7 @@ This function can display the persistence result as a barcode:
 .. plot::
    :include-source:
 
+    import matplotlib.pyplot as plot
     import gudhi
 
     off_file = gudhi.__root_source_dir__ + '/data/points/tore3D_300.off'
@@ -29,7 +30,7 @@ This function can display the persistence result as a barcode:
     simplex_tree = rips_complex.create_simplex_tree(max_dimension=3)
     diag = simplex_tree.persistence(min_persistence=0.4)
 
-    plot = gudhi.plot_persistence_barcode(diag)
+    gudhi.plot_persistence_barcode(diag)
     plot.show()
 
 Show persistence as a diagram
@@ -43,14 +44,15 @@ This function can display the persistence result as a diagram:
 .. plot::
    :include-source:
 
+    import matplotlib.pyplot as plot
     import gudhi
 
     # rips_on_tore3D_1307.pers obtained from write_persistence_diagram method
     persistence_file=gudhi.__root_source_dir__ + \
         '/data/persistence_diagram/rips_on_tore3D_1307.pers'
-    plt = gudhi.plot_persistence_diagram(persistence_file=persistence_file,
+    gudhi.plot_persistence_diagram(persistence_file=persistence_file,
         legend=True)
-    plt.show()
+    plot.show()
 
 Persistence density
 -------------------
@@ -63,11 +65,12 @@ If you want more information on a specific dimension, for instance:
 .. plot::
    :include-source:
 
+    import matplotlib.pyplot as plot
     import gudhi
 
     # rips_on_tore3D_1307.pers obtained from write_persistence_diagram method
     persistence_file=gudhi.__root_source_dir__ + \
         '/data/persistence_diagram/rips_on_tore3D_1307.pers'
-    plt = gudhi.plot_persistence_density(persistence_file=persistence_file,
+    gudhi.plot_persistence_density(persistence_file=persistence_file,
         max_intervals=0, dimension=1, legend=True)
-    plt.show()
+    plot.show()
diff --git a/src/python/example/alpha_complex_diagram_persistence_from_off_file_example.py b/src/python/example/alpha_complex_diagram_persistence_from_off_file_example.py
index b8f283b3..4079a469 100755
--- a/src/python/example/alpha_complex_diagram_persistence_from_off_file_example.py
+++ b/src/python/example/alpha_complex_diagram_persistence_from_off_file_example.py
@@ -1,7 +1,8 @@
 #!/usr/bin/env python
 
-import gudhi
 import argparse
+import matplotlib.pyplot as plot
+import gudhi
 
 """ This file is part of the Gudhi Library - https://gudhi.inria.fr/ - which is released under MIT.
     See file LICENSE or go to https://gudhi.inria.fr/licensing/ for full license details.
@@ -60,8 +61,8 @@ with open(args.file, "r") as f:
         print(simplex_tree.betti_numbers())
 
         if args.no_diagram == False:
-            pplot = gudhi.plot_persistence_diagram(diag, band=args.band)
-            pplot.show()
+            gudhi.plot_persistence_diagram(diag, band=args.band)
+            plot.show()
     else:
         print(args.file, "is not a valid OFF file")
 
diff --git a/src/python/example/euclidean_strong_witness_complex_diagram_persistence_from_off_file_example.py b/src/python/example/euclidean_strong_witness_complex_diagram_persistence_from_off_file_example.py
index 610ba44f..0eedd140 100755
--- a/src/python/example/euclidean_strong_witness_complex_diagram_persistence_from_off_file_example.py
+++ b/src/python/example/euclidean_strong_witness_complex_diagram_persistence_from_off_file_example.py
@@ -1,7 +1,8 @@
 #!/usr/bin/env python
 
-import gudhi
 import argparse
+import matplotlib.pyplot as plot
+import gudhi
 
 """ This file is part of the Gudhi Library - https://gudhi.inria.fr/ - which is released under MIT.
     See file LICENSE or go to https://gudhi.inria.fr/licensing/ for full license details.
@@ -75,8 +76,8 @@ with open(args.file, "r") as f:
         print(simplex_tree.betti_numbers())
 
         if args.no_diagram == False:
-            pplot = gudhi.plot_persistence_diagram(diag, band=args.band)
-            pplot.show()
+            gudhi.plot_persistence_diagram(diag, band=args.band)
+            plot.show()
     else:
         print(args.file, "is not a valid OFF file")
 
diff --git a/src/python/example/euclidean_witness_complex_diagram_persistence_from_off_file_example.py b/src/python/example/euclidean_witness_complex_diagram_persistence_from_off_file_example.py
index 7587b732..1fe55737 100755
--- a/src/python/example/euclidean_witness_complex_diagram_persistence_from_off_file_example.py
+++ b/src/python/example/euclidean_witness_complex_diagram_persistence_from_off_file_example.py
@@ -1,7 +1,8 @@
 #!/usr/bin/env python
 
-import gudhi
 import argparse
+import matplotlib.pyplot as plot
+import gudhi
 
 """ This file is part of the Gudhi Library - https://gudhi.inria.fr/ - which is released under MIT.
     See file LICENSE or go to https://gudhi.inria.fr/licensing/ for full license details.
@@ -74,8 +75,8 @@ with open(args.file, "r") as f:
         print(simplex_tree.betti_numbers())
 
         if args.no_diagram == False:
-            pplot = gudhi.plot_persistence_diagram(diag, band=args.band)
-            pplot.show()
+            gudhi.plot_persistence_diagram(diag, band=args.band)
+            plot.show()
     else:
         print(args.file, "is not a valid OFF file")
 
diff --git a/src/python/example/gudhi_graphical_tools_example.py b/src/python/example/gudhi_graphical_tools_example.py
index 3b0ca54d..37ecbf53 100755
--- a/src/python/example/gudhi_graphical_tools_example.py
+++ b/src/python/example/gudhi_graphical_tools_example.py
@@ -1,5 +1,6 @@
 #!/usr/bin/env python
 
+import matplotlib.pyplot as plot
 import gudhi
 
 """ This file is part of the Gudhi Library - https://gudhi.inria.fr/ - which is released under MIT.
@@ -29,15 +30,24 @@ persistence = [
     (0, (0.0, 1.0)),
 ]
 gudhi.plot_persistence_barcode(persistence)
+plot.show()
 
 print("#####################################################################")
 print("Show diagram persistence example")
 
-pplot = gudhi.plot_persistence_diagram(persistence)
-pplot.show()
+gudhi.plot_persistence_diagram(persistence)
+plot.show()
 
 print("#####################################################################")
 print("Show diagram persistence example with a confidence band")
 
-pplot = gudhi.plot_persistence_diagram(persistence, band=0.2)
-pplot.show()
+gudhi.plot_persistence_diagram(persistence, band=0.2)
+plot.show()
+
+print("#####################################################################")
+print("Show barcode and diagram persistence side by side example")
+fig, axes = plot.subplots(nrows=1, ncols=2)
+gudhi.plot_persistence_barcode(persistence, axes = axes[0])
+gudhi.plot_persistence_diagram(persistence, axes = axes[1])
+fig.suptitle("barcode versus diagram")
+plot.show()
diff --git a/src/python/example/periodic_cubical_complex_barcode_persistence_from_perseus_file_example.py b/src/python/example/periodic_cubical_complex_barcode_persistence_from_perseus_file_example.py
index 9cb855cd..c692e66f 100755
--- a/src/python/example/periodic_cubical_complex_barcode_persistence_from_perseus_file_example.py
+++ b/src/python/example/periodic_cubical_complex_barcode_persistence_from_perseus_file_example.py
@@ -1,7 +1,8 @@
 #!/usr/bin/env python
 
-import gudhi
 import argparse
+import matplotlib.pyplot as plot
+import gudhi
 
 """ This file is part of the Gudhi Library - https://gudhi.inria.fr/ - which is released under MIT.
     See file LICENSE or go to https://gudhi.inria.fr/licensing/ for full license details.
@@ -70,5 +71,6 @@ if is_file_perseus(args.file):
     print(periodic_cubical_complex.betti_numbers())
     if args.no_barcode == False:
         gudhi.plot_persistence_barcode(diag)
+        plot.show()
 else:
     print(args.file, "is not a valid perseus style file")
diff --git a/src/python/example/rips_complex_diagram_persistence_from_correlation_matrix_file_example.py b/src/python/example/rips_complex_diagram_persistence_from_correlation_matrix_file_example.py
index 3571580b..1acb187c 100755
--- a/src/python/example/rips_complex_diagram_persistence_from_correlation_matrix_file_example.py
+++ b/src/python/example/rips_complex_diagram_persistence_from_correlation_matrix_file_example.py
@@ -1,8 +1,9 @@
 #!/usr/bin/env python
 
-import gudhi
 import sys
 import argparse
+import matplotlib.pyplot as plot
+import gudhi
 
 """ This file is part of the Gudhi Library - https://gudhi.inria.fr/ - which is released under MIT.
     See file LICENSE or go to https://gudhi.inria.fr/licensing/ for full license details.
@@ -83,5 +84,5 @@ invert_diag = [
 ]
 
 if args.no_diagram == False:
-    pplot = gudhi.plot_persistence_diagram(invert_diag, band=args.band)
-    pplot.show()
+    gudhi.plot_persistence_diagram(invert_diag, band=args.band)
+    plot.show()
diff --git a/src/python/example/rips_complex_diagram_persistence_from_distance_matrix_file_example.py b/src/python/example/rips_complex_diagram_persistence_from_distance_matrix_file_example.py
index 0b9a9ba9..79ccca96 100755
--- a/src/python/example/rips_complex_diagram_persistence_from_distance_matrix_file_example.py
+++ b/src/python/example/rips_complex_diagram_persistence_from_distance_matrix_file_example.py
@@ -1,7 +1,8 @@
 #!/usr/bin/env python
 
-import gudhi
 import argparse
+import matplotlib.pyplot as plot
+import gudhi
 
 """ This file is part of the Gudhi Library - https://gudhi.inria.fr/ - which is released under MIT.
     See file LICENSE or go to https://gudhi.inria.fr/licensing/ for full license details.
@@ -59,5 +60,5 @@ print("betti_numbers()=")
 print(simplex_tree.betti_numbers())
 
 if args.no_diagram == False:
-    pplot = gudhi.plot_persistence_diagram(diag, band=args.band)
-    pplot.show()
+    gudhi.plot_persistence_diagram(diag, band=args.band)
+    plot.show()
diff --git a/src/python/example/rips_complex_diagram_persistence_from_off_file_example.py b/src/python/example/rips_complex_diagram_persistence_from_off_file_example.py
index 2b335bba..b9074cf9 100755
--- a/src/python/example/rips_complex_diagram_persistence_from_off_file_example.py
+++ b/src/python/example/rips_complex_diagram_persistence_from_off_file_example.py
@@ -1,7 +1,8 @@
 #!/usr/bin/env python
 
-import gudhi
 import argparse
+import matplotlib.pyplot as plot
+import gudhi
 
 """ This file is part of the Gudhi Library - https://gudhi.inria.fr/ - which is released under MIT.
     See file LICENSE or go to https://gudhi.inria.fr/licensing/ for full license details.
@@ -64,8 +65,8 @@ with open(args.file, "r") as f:
         print(simplex_tree.betti_numbers())
 
         if args.no_diagram == False:
-            pplot = gudhi.plot_persistence_diagram(diag, band=args.band)
-            pplot.show()
+            gudhi.plot_persistence_diagram(diag, band=args.band)
+            plot.show()
     else:
         print(args.file, "is not a valid OFF file")
 
diff --git a/src/python/example/rips_persistence_diagram.py b/src/python/example/rips_persistence_diagram.py
index f5897d7b..2a90b4bc 100755
--- a/src/python/example/rips_persistence_diagram.py
+++ b/src/python/example/rips_persistence_diagram.py
@@ -1,5 +1,6 @@
 #!/usr/bin/env python
 
+import matplotlib.pyplot as plot
 import gudhi
 
 """ This file is part of the Gudhi Library - https://gudhi.inria.fr/ - which is released under MIT.
@@ -26,5 +27,5 @@ simplex_tree = rips.create_simplex_tree(max_dimension=1)
 diag = simplex_tree.persistence(homology_coeff_field=2, min_persistence=0)
 print("diag=", diag)
 
-pplot = gudhi.plot_persistence_diagram(diag)
-pplot.show()
+gudhi.plot_persistence_diagram(diag)
+plot.show()
diff --git a/src/python/example/sparse_rips_persistence_diagram.py b/src/python/example/sparse_rips_persistence_diagram.py
index 671d5e34..410a6a86 100755
--- a/src/python/example/sparse_rips_persistence_diagram.py
+++ b/src/python/example/sparse_rips_persistence_diagram.py
@@ -1,5 +1,6 @@
 #!/usr/bin/env python
 
+import matplotlib.pyplot as plot
 import gudhi
 
 """ This file is part of the Gudhi Library - https://gudhi.inria.fr/ - which is released under MIT.
@@ -28,5 +29,5 @@ simplex_tree = rips.create_simplex_tree(max_dimension=2)
 diag = simplex_tree.persistence(homology_coeff_field=2, min_persistence=0)
 print("diag=", diag)
 
-pplot = gudhi.plot_persistence_diagram(diag)
-pplot.show()
+gudhi.plot_persistence_diagram(diag)
+plot.show()
diff --git a/src/python/example/tangential_complex_plain_homology_from_off_file_example.py b/src/python/example/tangential_complex_plain_homology_from_off_file_example.py
index 456bc9eb..f0df2189 100755
--- a/src/python/example/tangential_complex_plain_homology_from_off_file_example.py
+++ b/src/python/example/tangential_complex_plain_homology_from_off_file_example.py
@@ -1,7 +1,8 @@
 #!/usr/bin/env python
 
-import gudhi
 import argparse
+import matplotlib.pyplot as plot
+import gudhi
 
 """ This file is part of the Gudhi Library - https://gudhi.inria.fr/ - which is released under MIT.
     See file LICENSE or go to https://gudhi.inria.fr/licensing/ for full license details.
@@ -56,8 +57,8 @@ with open(args.file, "r") as f:
         print(st.betti_numbers())
 
         if args.no_diagram == False:
-            pplot = gudhi.plot_persistence_diagram(diag, band=args.band)
-            pplot.show()
+            gudhi.plot_persistence_diagram(diag, band=args.band)
+            plot.show()
     else:
         print(args.file, "is not a valid OFF file")
 
diff --git a/src/python/gudhi/persistence_graphical_tools.py b/src/python/gudhi/persistence_graphical_tools.py
index 43a00459..c9dab323 100644
--- a/src/python/gudhi/persistence_graphical_tools.py
+++ b/src/python/gudhi/persistence_graphical_tools.py
@@ -43,27 +43,6 @@ def __min_birth_max_death(persistence, band=0.0):
         max_death += band
     return (min_birth, max_death)
 
-
-"""
-Only 13 colors for the palette
-"""
-palette = [
-    "#ff0000",
-    "#00ff00",
-    "#0000ff",
-    "#00ffff",
-    "#ff00ff",
-    "#ffff00",
-    "#000000",
-    "#880000",
-    "#008800",
-    "#000088",
-    "#888800",
-    "#880088",
-    "#008888",
-]
-
-
 def plot_persistence_barcode(
     persistence=[],
     persistence_file="",
@@ -72,6 +51,8 @@ def plot_persistence_barcode(
     max_barcodes=1000,
     inf_delta=0.1,
     legend=False,
+    colormap=None,
+    axes=None
 ):
     """This function plots the persistence bar code from persistence values list
     or from a :doc:`persistence file <fileformats>`.
@@ -94,14 +75,19 @@ def plot_persistence_barcode(
     :type inf_delta: float.
     :param legend: Display the dimension color legend (default is False).
     :type legend: boolean.
-    :returns: A matplotlib object containing horizontal bar plot of persistence
-        (launch `show()` method on it to display it).
+    :param colormap: A matplotlib-like qualitative colormaps. Default is None
+        which means :code:`matplotlib.cm.Set1.colors`.
+    :type colormap: tuple of colors (3-tuple of float between 0. and 1.).
+    :param axes: A matplotlib-like subplot axes. If None, the plot is drawn on
+        a new set of axes.
+    :type axes: `matplotlib.axes.Axes`
+    :returns: (`matplotlib.axes.Axes`): The axes on which the plot was drawn.
     """
     try:
         import matplotlib.pyplot as plt
         import matplotlib.patches as mpatches
 
-        if persistence_file is not "":
+        if persistence_file != "":
             if path.isfile(persistence_file):
                 # Reset persistence
                 persistence = []
@@ -115,7 +101,7 @@ def plot_persistence_barcode(
                 print("file " + persistence_file + " not found.")
                 return None
 
-        if max_barcodes is not 1000:
+        if max_barcodes != 1000:
             print("Deprecated parameter. It has been replaced by max_intervals")
             max_intervals = max_barcodes
 
@@ -126,6 +112,11 @@ def plot_persistence_barcode(
                 key=lambda life_time: life_time[1][1] - life_time[1][0],
                 reverse=True,
             )[:max_intervals]
+        
+        if colormap == None:
+            colormap = plt.cm.Set1.colors
+        if axes == None:
+            fig, axes = plt.subplots(1, 1)
 
         persistence = sorted(persistence, key=lambda birth: birth[1][0])
 
@@ -140,41 +131,43 @@ def plot_persistence_barcode(
         for interval in reversed(persistence):
             if float(interval[1][1]) != float("inf"):
                 # Finite death case
-                plt.barh(
+                axes.barh(
                     ind,
                     (interval[1][1] - interval[1][0]),
                     height=0.8,
                     left=interval[1][0],
                     alpha=alpha,
-                    color=palette[interval[0]],
+                    color=colormap[interval[0]],
                     linewidth=0,
                 )
             else:
                 # Infinite death case for diagram to be nicer
-                plt.barh(
+                axes.barh(
                     ind,
                     (infinity - interval[1][0]),
                     height=0.8,
                     left=interval[1][0],
                     alpha=alpha,
-                    color=palette[interval[0]],
+                    color=colormap[interval[0]],
                     linewidth=0,
                 )
             ind = ind + 1
 
         if legend:
             dimensions = list(set(item[0] for item in persistence))
-            plt.legend(
+            axes.legend(
                 handles=[
-                    mpatches.Patch(color=palette[dim], label=str(dim))
+                    mpatches.Patch(color=colormap[dim], label=str(dim))
                     for dim in dimensions
                 ],
                 loc="lower right",
             )
-        plt.title("Persistence barcode")
+
+        axes.set_title("Persistence barcode")
+
         # Ends plot on infinity value and starts a little bit before min_birth
-        plt.axis([axis_start, infinity, 0, ind])
-        return plt
+        axes.axis([axis_start, infinity, 0, ind])
+        return axes
 
     except ImportError:
         print("This function is not available, you may be missing matplotlib.")
@@ -189,6 +182,8 @@ def plot_persistence_diagram(
     max_plots=1000,
     inf_delta=0.1,
     legend=False,
+    colormap=None,
+    axes=None
 ):
     """This function plots the persistence diagram from persistence values
     list or from a :doc:`persistence file <fileformats>`.
@@ -213,14 +208,19 @@ def plot_persistence_diagram(
     :type inf_delta: float.
     :param legend: Display the dimension color legend (default is False).
     :type legend: boolean.
-    :returns: A matplotlib object containing diagram plot of persistence
-        (launch `show()` method on it to display it).
+    :param colormap: A matplotlib-like qualitative colormaps. Default is None
+        which means :code:`matplotlib.cm.Set1.colors`.
+    :type colormap: tuple of colors (3-tuple of float between 0. and 1.).
+    :param axes: A matplotlib-like subplot axes. If None, the plot is drawn on
+        a new set of axes.
+    :type axes: `matplotlib.axes.Axes`
+    :returns: (`matplotlib.axes.Axes`): The axes on which the plot was drawn.
     """
     try:
         import matplotlib.pyplot as plt
         import matplotlib.patches as mpatches
 
-        if persistence_file is not "":
+        if persistence_file != "":
             if path.isfile(persistence_file):
                 # Reset persistence
                 persistence = []
@@ -234,7 +234,7 @@ def plot_persistence_diagram(
                 print("file " + persistence_file + " not found.")
                 return None
 
-        if max_plots is not 1000:
+        if max_plots != 1000:
             print("Deprecated parameter. It has been replaced by max_intervals")
             max_intervals = max_plots
 
@@ -246,6 +246,11 @@ def plot_persistence_diagram(
                 reverse=True,
             )[:max_intervals]
 
+        if colormap == None:
+            colormap = plt.cm.Set1.colors
+        if axes == None:
+            fig, axes = plt.subplots(1, 1)
+
         (min_birth, max_death) = __min_birth_max_death(persistence, band)
         delta = (max_death - min_birth) * inf_delta
         # Replace infinity values with max_death + delta for diagram to be more
@@ -256,44 +261,44 @@ def plot_persistence_diagram(
         # line display of equation : birth = death
         x = np.linspace(axis_start, infinity, 1000)
         # infinity line and text
-        plt.plot(x, x, color="k", linewidth=1.0)
-        plt.plot(x, [infinity] * len(x), linewidth=1.0, color="k", alpha=alpha)
-        plt.text(axis_start, infinity, r"$\infty$", color="k", alpha=alpha)
+        axes.plot(x, x, color="k", linewidth=1.0)
+        axes.plot(x, [infinity] * len(x), linewidth=1.0, color="k", alpha=alpha)
+        axes.text(axis_start, infinity, r"$\infty$", color="k", alpha=alpha)
         # bootstrap band
         if band > 0.0:
-            plt.fill_between(x, x, x + band, alpha=alpha, facecolor="red")
+            axes.fill_between(x, x, x + band, alpha=alpha, facecolor="red")
 
         # Draw points in loop
         for interval in reversed(persistence):
             if float(interval[1][1]) != float("inf"):
                 # Finite death case
-                plt.scatter(
+                axes.scatter(
                     interval[1][0],
                     interval[1][1],
                     alpha=alpha,
-                    color=palette[interval[0]],
+                    color=colormap[interval[0]],
                 )
             else:
                 # Infinite death case for diagram to be nicer
-                plt.scatter(
-                    interval[1][0], infinity, alpha=alpha, color=palette[interval[0]]
+                axes.scatter(
+                    interval[1][0], infinity, alpha=alpha, color=colormap[interval[0]]
                 )
 
         if legend:
             dimensions = list(set(item[0] for item in persistence))
-            plt.legend(
+            axes.legend(
                 handles=[
-                    mpatches.Patch(color=palette[dim], label=str(dim))
+                    mpatches.Patch(color=colormap[dim], label=str(dim))
                     for dim in dimensions
                 ]
             )
 
-        plt.title("Persistence diagram")
-        plt.xlabel("Birth")
-        plt.ylabel("Death")
+        axes.set_xlabel("Birth")
+        axes.set_ylabel("Death")
         # Ends plot on infinity value and starts a little bit before min_birth
-        plt.axis([axis_start, infinity, axis_start, infinity + delta])
-        return plt
+        axes.axis([axis_start, infinity, axis_start, infinity + delta])
+        axes.set_title("Persistence diagram")
+        return axes
 
     except ImportError:
         print("This function is not available, you may be missing matplotlib.")
@@ -308,6 +313,7 @@ def plot_persistence_density(
     dimension=None,
     cmap=None,
     legend=False,
+    axes=None
 ):
     """This function plots the persistence density from persistence
     values list or from a :doc:`persistence file <fileformats>`. Be
@@ -346,14 +352,16 @@ def plot_persistence_density(
     :type cmap: cf. matplotlib colormap.
     :param legend: Display the color bar values (default is False).
     :type legend: boolean.
-    :returns: A matplotlib object containing diagram plot of persistence
-        (launch `show()` method on it to display it).
+    :param axes: A matplotlib-like subplot axes. If None, the plot is drawn on
+        a new set of axes.
+    :type axes: `matplotlib.axes.Axes`
+    :returns: (`matplotlib.axes.Axes`): The axes on which the plot was drawn.
     """
     try:
         import matplotlib.pyplot as plt
         from scipy.stats import kde
 
-        if persistence_file is not "":
+        if persistence_file != "":
             if dimension is None:
                 # All dimension case
                 dimension = -1
@@ -390,9 +398,15 @@ def plot_persistence_density(
         birth = persistence_dim[:, 0]
         death = persistence_dim[:, 1]
 
+        # default cmap value cannot be done at argument definition level as matplotlib is not yet defined.
+        if cmap is None:
+            cmap = plt.cm.hot_r
+        if axes == None:
+            fig, axes = plt.subplots(1, 1)
+
         # line display of equation : birth = death
         x = np.linspace(death.min(), birth.max(), 1000)
-        plt.plot(x, x, color="k", linewidth=1.0)
+        axes.plot(x, x, color="k", linewidth=1.0)
 
         # Evaluate a gaussian kde on a regular grid of nbins x nbins over data extents
         k = kde.gaussian_kde([birth, death], bw_method=bw_method)
@@ -402,19 +416,16 @@ def plot_persistence_density(
         ]
         zi = k(np.vstack([xi.flatten(), yi.flatten()]))
 
-        # default cmap value cannot be done at argument definition level as matplotlib is not yet defined.
-        if cmap is None:
-            cmap = plt.cm.hot_r
         # Make the plot
-        plt.pcolormesh(xi, yi, zi.reshape(xi.shape), cmap=cmap)
+        axes.pcolormesh(xi, yi, zi.reshape(xi.shape), cmap=cmap)
 
         if legend:
-            plt.colorbar()
+            axes.colorbar()
 
-        plt.title("Persistence density")
-        plt.xlabel("Birth")
-        plt.ylabel("Death")
-        return plt
+        axes.set_xlabel("Birth")
+        axes.set_ylabel("Death")
+        axes.set_title("Persistence density")
+        return axes
 
     except ImportError:
         print(