Make CubicalPersistence returns all dimensions. Post processing DimensionSelector can select the desired dimension

2 files changed, 100 insertions, 10 deletions

diff --git a/src/python/gudhi/sklearn/cubical_persistence.py b/src/python/gudhi/sklearn/cubical_persistence.py
index 9af683d7..7b77000d 100644
--- a/src/python/gudhi/sklearn/cubical_persistence.py
+++ b/src/python/gudhi/sklearn/cubical_persistence.py
@@ -13,27 +13,44 @@ from sklearn.base import BaseEstimator, TransformerMixin
 # joblib is required by scikit-learn
 from joblib import Parallel, delayed
 
+# Mermaid sequence diagram - https://mermaid-js.github.io/mermaid-live-editor/
+# sequenceDiagram
+#     USER->>CubicalPersistence: fit_transform(X)
+#     CubicalPersistence->>thread1: _tranform(X[0])
+#     CubicalPersistence->>thread2: _tranform(X[1])
+#     Note right of CubicalPersistence: ...
+#     thread1->>CubicalPersistence: [array( H0(X[0]) ), array( H1(X[0]) )]
+#     thread2->>CubicalPersistence: [array( H0(X[1]) ), array( H1(X[1]) )]
+#     Note right of CubicalPersistence: ...
+#     CubicalPersistence->>USER: [[array( H0(X[0]) ), array( H1(X[0]) )],<br/> [array( H0(X[1]) ), array( H1(X[1]) )],<br/> ...]
+
 
 class CubicalPersistence(BaseEstimator, TransformerMixin):
     """
     This is a class for computing the persistence diagrams from a cubical complex.
     """
 
-    def __init__(self, dimensions=None, persistence_dim=0, homology_coeff_field=11, min_persistence=0., n_jobs=None):
+    def __init__(self, dimensions=None, max_persistence_dimension=0, only_this_dim=-1, homology_coeff_field=11, min_persistence=0., n_jobs=None):
         """
         Constructor for the CubicalPersistence class.
 
         Parameters:
             dimensions (list of int): A list of number of top dimensional cells if cells filtration values will require
                 to be reshaped (cf. :func:`~gudhi.sklearn.cubical_persistence.CubicalPersistence.transform`)
-            persistence_dim (int): The returned persistence diagrams dimension. Default value is `0`.
+            max_persistence_dimension (int): The returned persistence diagrams maximal dimension. Default value is `0`.
+                Ignored if `only_this_dim` is set.
+            only_this_dim (int): The returned persistence diagrams dimension. If `only_this_dim` is set,
+                `max_persistence_dimension` will be ignored. 
+                Short circuit the use of :class:`~gudhi.sklearn.post_processing.DimensionSelector` when only one
+                dimension matters.
             homology_coeff_field (int): The homology coefficient field. Must be a prime number. Default value is 11.
             min_persistence (float): The minimum persistence value to take into account (strictly greater than
                 `min_persistence`). Default value is `0.0`. Sets `min_persistence` to `-1.0` to see all values.
             n_jobs (int): cf. https://joblib.readthedocs.io/en/latest/generated/joblib.Parallel.html
         """
         self.dimensions = dimensions
-        self.persistence_dim = persistence_dim
+        self.max_persistence_dimension = max_persistence_dimension
+        self.only_this_dim = only_this_dim
         self.homology_coeff_field = homology_coeff_field
         self.min_persistence = min_persistence
         self.n_jobs = n_jobs
@@ -49,8 +66,14 @@ class CubicalPersistence(BaseEstimator, TransformerMixin):
         cubical_complex.compute_persistence(
             homology_coeff_field=self.homology_coeff_field, min_persistence=self.min_persistence
         )
-        diagrams = cubical_complex.persistence_intervals_in_dimension(self.persistence_dim)
-        return diagrams
+        return [cubical_complex.persistence_intervals_in_dimension(dim) for dim in range(self.max_persistence_dimension + 1)]
+
+    def __transform_only_this_dim(self, cells):
+        cubical_complex = CubicalComplex(top_dimensional_cells=cells, dimensions=self.dimensions)
+        cubical_complex.compute_persistence(
+            homology_coeff_field=self.homology_coeff_field, min_persistence=self.min_persistence
+        )
+        return cubical_complex.persistence_intervals_in_dimension(self.only_this_dim)
 
     def transform(self, X, Y=None):
         """
@@ -58,12 +81,18 @@ class CubicalPersistence(BaseEstimator, TransformerMixin):
 
         Parameters:
             X (list of list of double OR list of numpy.ndarray): List of cells filtration values that can be flatten if
-                dimensions is set in the constructor, or already with the correct shape in a numpy.ndarray (and
-                dimensions must not be set).
+                `dimensions` is set in the constructor, or already with the correct shape in a numpy.ndarray (and
+                `dimensions` must not be set).
 
         Returns:
-            Persistence diagrams
+            Persistence diagrams in the format:
+            - If `only_this_dim` was set to `n`: `[array( Hn(X[0]) ), array( Hn(X[1]) ), ...]` 
+            - else: `[[array( H0(X[0]) ), array( H1(X[0]) ), ...], [array( H0(X[1]) ), array( H1(X[1]) ), ...], ...]` 
         """
 
-        # threads is preferred as cubical construction and persistence computation releases the GIL
-        return Parallel(n_jobs=self.n_jobs, prefer="threads")(delayed(self.__transform)(cells) for cells in X)
+        if self.only_this_dim == -1:
+            # threads is preferred as cubical construction and persistence computation releases the GIL
+            return Parallel(n_jobs=self.n_jobs, prefer="threads")(delayed(self.__transform)(cells) for cells in X)
+        else:
+            # threads is preferred as cubical construction and persistence computation releases the GIL
+            return Parallel(n_jobs=self.n_jobs, prefer="threads")(delayed(self.__transform_only_this_dim)(cells) for cells in X)
diff --git a/src/python/gudhi/sklearn/post_processing.py b/src/python/gudhi/sklearn/post_processing.py
new file mode 100644
index 00000000..79276e1e
--- /dev/null
+++ b/src/python/gudhi/sklearn/post_processing.py
@@ -0,0 +1,61 @@
+# 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.
+# Author(s):       Vincent Rouvreau
+#
+# Copyright (C) 2021 Inria
+#
+# Modification(s):
+#   - YYYY/MM Author: Description of the modification
+
+from sklearn.base import BaseEstimator, TransformerMixin
+
+# joblib is required by scikit-learn
+from joblib import Parallel, delayed
+
+# Mermaid sequence diagram - https://mermaid-js.github.io/mermaid-live-editor/
+# sequenceDiagram
+#     USER->>DimensionSelector: fit_transform(<br/>[[array( H0(X0) ), array( H1(X0) ), ...],<br/> [array( H0(X1) ), array( H1(X1) ), ...],<br/> ...])
+#     DimensionSelector->>thread1: _transform([array( H0(X0) ), array( H1(X0) )], ...)
+#     DimensionSelector->>thread2: _transform([array( H0(X1) ), array( H1(X1) )], ...)
+#     Note right of DimensionSelector: ...
+#     thread1->>DimensionSelector: array( Hn(X0) )
+#     thread2->>DimensionSelector: array( Hn(X1) )
+#     Note right of DimensionSelector: ...
+#     DimensionSelector->>USER: [array( Hn(X0) ), <br/> array( Hn(X1) ), <br/> ...]
+
+
+class DimensionSelector(BaseEstimator, TransformerMixin):
+    """
+    This is a class to select persistence diagrams in a specific dimension.
+    """
+
+    def __init__(self, persistence_dimension=0, n_jobs=None):
+        """
+        Constructor for the DimensionSelector class.
+
+        Parameters:
+            persistence_dimension (int): The returned persistence diagrams dimension. Default value is `0`.
+        """
+        self.persistence_dimension = persistence_dimension
+        self.n_jobs = n_jobs
+
+    def fit(self, X, Y=None):
+        """
+        Nothing to be done, but useful when included in a scikit-learn Pipeline.
+        """
+        return self
+
+    def transform(self, X, Y=None):
+        """
+        Select persistence diagrams from its dimension.
+
+        Parameters:
+            X (list of list of pairs): List of list of persistence pairs, i.e.
+            `[[array( H0(X0) ), array( H1(X0) ), ...], [array( H0(X1) ), array( H1(X1) ), ...], ...]` 
+
+        Returns:
+            Persistence diagrams in a specific dimension, i.e.
+            `[array( Hn(X0) ), array( Hn(X1), ...]`
+        """
+
+        return [persistence[self.persistence_dimension] for persistence in X]