multipers 2.3.3b6__cp311-cp311-manylinux_2_39_x86_64.whl

multipers/ml/sliced_wasserstein.py

@@ -0,0 +1,461 @@
+# This code was written by Mathieu Carrière.
+
+import numpy as np
+from sklearn.base import BaseEstimator, TransformerMixin
+from sklearn.metrics import pairwise_distances, pairwise_kernels
+from joblib import Parallel, delayed
+
+
+def _pairwise(fallback, skipdiag, X, Y, metric, n_jobs):
+    if Y is not None:
+        return fallback(X, Y, metric=metric, n_jobs=n_jobs)
+    triu = np.triu_indices(len(X), k=skipdiag)
+    tril = (triu[1], triu[0])
+    par = Parallel(n_jobs=n_jobs, prefer="threads")
+    d = par(delayed(metric)([triu[0][i]], [triu[1][i]])
+            for i in range(len(triu[0])))
+    m = np.empty((len(X), len(X)))
+    m[triu] = d
+    m[tril] = d
+    if skipdiag:
+        np.fill_diagonal(m, 0)
+    return m
+
+
+def _sklearn_wrapper(metric, X, Y, **kwargs):
+    """
+    This function is a wrapper for any metric between two signed measures that takes two numpy arrays of shapes (nxD) and (mxD) as arguments.
+    """
+    if Y is None:
+
+        def flat_metric(a, b):
+            return metric(X[int(a[0])], X[int(b[0])], **kwargs)
+    else:
+
+        def flat_metric(a, b):
+            return metric(X[int(a[0])], Y[int(b[0])], **kwargs)
+
+    return flat_metric
+
+
+def _compute_signed_measure_parts(X):
+    """
+    This is a function for separating the positive and negative points of a list of signed measures. This function can be used as a preprocessing step in order to speed up the running time for computing all pairwise (sliced) Wasserstein distances on a list of signed measures.
+
+    Parameters:
+        X (list of n tuples): list of signed measures.
+
+    Returns:
+        list of n pairs of numpy arrays of shape (num x dimension): list of positive and negative signed measures.
+    """
+    XX = []
+    for C, M in X:
+        pos_idxs = np.argwhere(M > 0).ravel()
+        neg_idxs = np.setdiff1d(np.arange(len(M)), pos_idxs)
+        XX.append(
+            [
+                np.repeat(C[pos_idxs], M[pos_idxs], axis=0),
+                np.repeat(C[neg_idxs], -M[neg_idxs], axis=0),
+            ]
+        )
+    return XX
+
+
+def _compute_signed_measure_projections(X, num_directions, scales):
+    """
+    This is a function for projecting the points of a list of signed measures onto a fixed number of lines sampled uniformly. This function can be used as a preprocessing step in order to speed up the running time for computing all pairwise sliced Wasserstein distances on a list of signed measures.
+
+    Parameters:
+        X (list of n tuples): list of signed measures.
+        num_directions (int): number of lines evenly sampled from [-pi/2,pi/2] in order to approximate and speed up the distance computation.
+        scales (array of shape D): scales associated to the dimensions.
+
+    Returns:
+        list of n pairs of numpy arrays of shape (num x num_directions): list of positive and negative projected signed measures.
+    """
+    dimension = X[0][0].shape[1]
+    np.random.seed(42)
+    thetas = np.random.normal(0, 1, [num_directions, dimension])
+    lines = (thetas / np.linalg.norm(thetas, axis=1)[:, None]).T
+    weights = (
+        np.linalg.norm(np.multiply(scales[:, None], lines), axis=0)
+        if scales is not None
+        else np.ones(num_directions)
+    )
+    XX = []
+    for C, M in X:
+        pos_idxs = np.argwhere(M > 0).ravel()
+        neg_idxs = np.setdiff1d(np.arange(len(M)), pos_idxs)
+        XX.append(
+            [
+                np.matmul(np.repeat(C[pos_idxs], M[pos_idxs], axis=0), lines),
+                np.matmul(np.repeat(C[neg_idxs], -M[neg_idxs], axis=0), lines),
+                weights,
+            ]
+        )
+    return XX
+
+
+def pairwise_signed_measure_distances(
+    X, Y=None, metric="sliced_wasserstein", n_jobs=None, **kwargs
+):
+    """
+    This function computes the distance matrix between two lists of signed measures given as numpy arrays of shape (nxD).
+
+    Parameters:
+        X (list of n tuples): first list of signed measures.
+        Y (list of m tuples): second list of signed measures (optional). If None, pairwise distances are computed from the first list only.
+        metric: distance to use. It can be either a string ("sliced_wasserstein", "wasserstein") or a function taking two tuples as inputs. If it is a function, make sure that it is symmetric and that it outputs 0 if called on the same two tuples.
+        n_jobs (int): number of jobs to use for the computation. This uses joblib.Parallel(prefer="threads"), so metrics that do not release the GIL may not scale unless run inside a `joblib.parallel_backend <https://joblib.readthedocs.io/en/latest/parallel.html#joblib.parallel_backend>`_ block.
+        **kwargs: optional keyword parameters. Any further parameters are passed directly to the distance function. See the docs of the various distance classes in this module.
+
+    Returns:
+        numpy array of shape (nxm): distance matrix
+    """
+    XX = np.reshape(np.arange(len(X)), [-1, 1])
+    YY = None if Y is None or Y is X else np.reshape(
+        np.arange(len(Y)), [-1, 1])
+    if metric == "sliced_wasserstein":
+        Xproj = _compute_signed_measure_projections(X, **kwargs)
+        Yproj = None if Y is None else _compute_signed_measure_projections(
+            Y, **kwargs)
+        return _pairwise(
+            pairwise_distances,
+            True,
+            XX,
+            YY,
+            metric=_sklearn_wrapper(
+                _sliced_wasserstein_distance_on_projections, Xproj, Yproj
+            ),
+            n_jobs=n_jobs,
+        )
+    elif metric == "wasserstein":
+        Xproj = _compute_signed_measure_parts(X)
+        Yproj = None if Y is None else _compute_signed_measure_parts(Y)
+        return _pairwise(
+            pairwise_distances,
+            True,
+            XX,
+            YY,
+            metric=_sklearn_wrapper(
+                _wasserstein_distance_on_parts(**kwargs), Xproj, Yproj
+            ),
+            n_jobs=n_jobs,
+        )
+    else:
+        return _pairwise(
+            pairwise_distances,
+            True,
+            XX,
+            YY,
+            metric=_sklearn_wrapper(metric, X, Y, **kwargs),
+            n_jobs=n_jobs,
+        )
+
+
+def _wasserstein_distance_on_parts(ground_norm=1, epsilon=1.0):
+    """
+    This is a function for computing the Wasserstein distance between two signed measures that have already been separated into their positive and negative parts.
+
+    Parameters:
+        meas1: pair of (n x dimension) numpy.arrays containing the points of the positive and negative parts of the first measure.
+        meas2: pair of (m x dimension) numpy.arrays containing the points of the positive and negative parts of the second measure.
+
+    Returns:
+        float: the sliced Wasserstein distance between the projected signed measures.
+    """
+
+    def metric(meas1, meas2):
+        meas1_plus, meas1_minus = meas1[0], meas1[1]
+        meas2_plus, meas2_minus = meas2[0], meas2[1]
+        num_pts = len(meas1_plus) + len(meas2_minus)
+        meas_t1 = np.vstack([meas1_plus, meas2_minus])
+        meas_t2 = np.vstack([meas2_plus, meas1_minus])
+        import ot
+
+        if epsilon > 0:
+            wass = ot.sinkhorn2(
+                1 / num_pts * np.ones(num_pts),
+                1 / num_pts * np.ones(num_pts),
+                pairwise_distances(
+                    meas_t1, meas_t2, metric="minkowski", p=ground_norm),
+                epsilon,
+            )
+            return wass[0]
+        else:
+            wass = ot.lp.emd2(
+                [],
+                [],
+                np.ascontiguousarray(
+                    pairwise_distances(
+                        meas_t1, meas_t2, metric="minkowski", p=ground_norm
+                    ),
+                    dtype=np.float64,
+                ),
+            )
+            return wass
+
+    return metric
+
+
+def _sliced_wasserstein_distance_on_projections(meas1, meas2, scales=None):
+    """
+    This is a function for computing the sliced Wasserstein distance between two signed measures that have already been projected onto some lines. It simply amounts to comparing the sorted projections with the 1-norm, and averaging over the lines. See http://proceedings.mlr.press/v70/carriere17a.html for more details.
+
+    Parameters:
+        meas1: pair of (n x number_of_lines) numpy.arrays containing the projected points of the positive and negative parts of the first measure.
+        meas2: pair of (m x number_of_lines) numpy.arrays containing the projected points of the positive and negative parts of the second measure.
+        scales (array of shape D): scales associated to the dimensions.
+
+    Returns:
+        float: the sliced Wasserstein distance between the projected signed measures.
+    """
+    # assert np.array_equal(  meas1[2], meas2[2]  )
+    weights = meas1[2]
+    meas1_plus, meas1_minus = meas1[0], meas1[1]
+    meas2_plus, meas2_minus = meas2[0], meas2[1]
+    A = np.sort(np.vstack([meas1_plus, meas2_minus]), axis=0)
+    B = np.sort(np.vstack([meas2_plus, meas1_minus]), axis=0)
+    L1 = np.sum(np.abs(A - B), axis=0)
+    return np.mean(np.multiply(L1, weights))
+
+
+def _sliced_wasserstein_distance(meas1, meas2, num_directions, scales=None):
+    """
+    This is a function for computing the sliced Wasserstein distance from two signed measures. The Sliced Wasserstein distance is computed by projecting the signed measures onto lines, comparing the projections with the 1-norm, and finally averaging over the lines. See http://proceedings.mlr.press/v70/carriere17a.html for more details.
+
+    Parameters:
+        meas1: ((n x D), (n)) tuple with numpy.array encoding the (finite points of the) first measure and their multiplicities. Must not contain essential points (i.e. with infinite coordinate).
+        meas2: ((m x D), (m)) tuple encoding the second measure.
+        num_directions (int): number of lines evenly sampled from [-pi/2,pi/2] in order to approximate and speed up the distance computation.
+        scales (array of shape D): scales associated to the dimensions.
+
+    Returns:
+        float: the sliced Wasserstein distance between signed measures.
+    """
+    C1, M1 = meas1[0], meas1[1]
+    C2, M2 = meas2[0], meas2[1]
+    dimension = C1.shape[1]
+    C1_plus_idxs, C2_plus_idxs = (
+        np.argwhere(M1 > 0).ravel(),
+        np.argwhere(M2 > 0).ravel(),
+    )
+    C1_minus_idxs, C2_minus_idxs = (
+        np.setdiff1d(np.arange(len(M1)), C1_plus_idxs),
+        np.setdiff1d(np.arange(len(M2)), C2_plus_idxs),
+    )
+    np.random.seed(42)
+    thetas = np.random.normal(0, 1, [num_directions, dimension])
+    lines = (thetas / np.linalg.norm(thetas, axis=1)[:, None]).T
+    weights = (
+        np.linalg.norm(np.multiply(scales[:, None], lines), axis=0)
+        if scales is not None
+        else np.ones(num_directions)
+    )
+    approx1 = np.matmul(
+        np.vstack(
+            [
+                np.repeat(C1[C1_plus_idxs], M1[C1_plus_idxs], axis=0),
+                np.repeat(C2[C2_minus_idxs], -M2[C2_minus_idxs], axis=0),
+            ]
+        ),
+        lines,
+    )
+    approx2 = np.matmul(
+        np.vstack(
+            [
+                np.repeat(C2[C2_plus_idxs], M2[C2_plus_idxs], axis=0),
+                np.repeat(C1[C1_minus_idxs], -M1[C1_minus_idxs], axis=0),
+            ]
+        ),
+        lines,
+    )
+    A = np.sort(approx1, axis=0)
+    B = np.sort(approx2, axis=0)
+    L1 = np.sum(np.abs(A - B), axis=0)
+    return np.mean(np.multiply(L1, weights))
+
+
+def _wasserstein_distance(meas1, meas2, epsilon, ground_norm):
+    """
+    This is a function for computing the Wasserstein distance from two signed measures.
+
+    Parameters:
+        meas1: ((n x D), (n)) tuple with numpy.array encoding the (finite points of the) first measure and their multiplicities. Must not contain essential points (i.e. with infinite coordinate).
+        meas2: ((m x D), (m)) tuple encoding the second measure.
+        epsilon (float): entropy regularization parameter.
+        ground_norm (int): norm to use for ground metric cost.
+
+    Returns:
+        float: the Wasserstein distance between signed measures.
+    """
+    C1, M1 = meas1[0], meas1[1]
+    C2, M2 = meas2[0], meas2[1]
+    C1_plus_idxs, C2_plus_idxs = (
+        np.argwhere(M1 > 0).ravel(),
+        np.argwhere(M2 > 0).ravel(),
+    )
+    C1_minus_idxs, C2_minus_idxs = (
+        np.setdiff1d(np.arange(len(M1)), C1_plus_idxs),
+        np.setdiff1d(np.arange(len(M2)), C2_plus_idxs),
+    )
+    approx1 = np.vstack(
+        [
+            np.repeat(C1[C1_plus_idxs], M1[C1_plus_idxs], axis=0),
+            np.repeat(C2[C2_minus_idxs], -M2[C2_minus_idxs], axis=0),
+        ]
+    )
+    approx2 = np.vstack(
+        [
+            np.repeat(C2[C2_plus_idxs], M2[C2_plus_idxs], axis=0),
+            np.repeat(C1[C1_minus_idxs], -M1[C1_minus_idxs], axis=0),
+        ]
+    )
+    num_pts = len(approx1)
+    import ot
+
+    if epsilon > 0:
+        wass = ot.sinkhorn2(
+            1 / num_pts * np.ones(num_pts),
+            1 / num_pts * np.ones(num_pts),
+            pairwise_distances(
+                approx1, approx2, metric="minkowski", p=ground_norm),
+            epsilon,
+        )
+        return wass[0]
+    else:
+        wass = ot.lp.emd2(
+            1 / num_pts * np.ones(num_pts),
+            1 / num_pts * np.ones(num_pts),
+            pairwise_distances(
+                approx1, approx2, metric="minkowski", p=ground_norm),
+        )
+        return wass
+
+
+class SlicedWassersteinDistance(BaseEstimator, TransformerMixin):
+    """
+    This is a class for computing the sliced Wasserstein distance matrix from a list of signed measures. The Sliced Wasserstein distance is computed by projecting the signed measures onto lines, comparing the projections with the 1-norm, and finally integrating over all possible lines. See http://proceedings.mlr.press/v70/carriere17a.html for more details.
+    """
+
+    def __init__(self, num_directions=10, scales=None, n_jobs=None):
+        """
+        Constructor for the SlicedWassersteinDistance class.
+
+        Parameters:
+            num_directions (int): number of lines evenly sampled in order to approximate and speed up the distance computation (default 10).
+            scales (array of shape D): scales associated to the dimensions.
+            n_jobs (int): number of jobs to use for the computation. See :func:`pairwise_signed_measure_distances` for details.
+        """
+        self.num_directions = num_directions
+        self.scales = scales
+        self.n_jobs = n_jobs
+
+    def fit(self, X, y=None):
+        """
+        Fit the SlicedWassersteinDistance class on a list of signed measures: signed measures are projected onto the different lines. The measures themselves are then stored in numpy arrays, called **measures_**.
+
+        Parameters:
+            X (list of tuples): input signed measures.
+            y (n x 1 array): signed measure labels (unused).
+        """
+        self.measures_ = X
+        return self
+
+    def transform(self, X):
+        """
+        Compute all sliced Wasserstein distances between the signed measures that were stored after calling the fit() method, and a given list of (possibly different) signed measures.
+
+        Parameters:
+            X (list of tuples): input signed measures.
+
+        Returns:
+            numpy array of shape (number of measures in **measures**) x (number of measures in X): matrix of pairwise sliced Wasserstein distances.
+        """
+        return pairwise_signed_measure_distances(
+            X,
+            self.measures_,
+            metric="sliced_wasserstein",
+            num_directions=self.num_directions,
+            scales=self.scales,
+            n_jobs=self.n_jobs,
+        )
+
+    def __call__(self, meas1, meas2):
+        """
+        Apply SlicedWassersteinDistance on a single pair of signed measures and outputs the result.
+
+        Parameters:
+            meas1: ((n x D), (n)) tuple with numpy.array encoding the (finite points of the) first measure and their multiplicities. Must not contain essential points (i.e. with infinite coordinate).
+            meas2: ((m x D), (m)) tuple encoding the second measure.
+
+        Returns:
+            float: sliced Wasserstein distance.
+        """
+        return _sliced_wasserstein_distance(
+            meas1, meas2, num_directions=self.num_directions, scales=self.scales
+        )
+
+
+class WassersteinDistance(BaseEstimator, TransformerMixin):
+    """
+    This is a class for computing the Wasserstein distance matrix from a list of signed measures.
+    """
+
+    def __init__(self, epsilon=1.0, ground_norm=1, n_jobs=None):
+        """
+        Constructor for the WassersteinDistance class.
+
+        Parameters:
+            epsilon (float): entropy regularization parameter.
+            ground_norm (int): norm to use for ground metric cost.
+            n_jobs (int): number of jobs to use for the computation. See :func:`pairwise_signed_measure_distances` for details.
+        """
+        self.epsilon = epsilon
+        self.ground_norm = ground_norm
+        self.n_jobs = n_jobs
+
+    def fit(self, X, y=None):
+        """
+        Fit the WassersteinDistance class on a list of signed measures. The measures themselves are then stored in numpy arrays, called **measures_**.
+
+        Parameters:
+            X (list of tuples): input signed measures.
+            y (n x 1 array): signed measure labels (unused).
+        """
+        self.measures_ = X
+        return self
+
+    def transform(self, X):
+        """
+        Compute all Wasserstein distances between the signed measures that were stored after calling the fit() method, and a given list of (possibly different) signed measures.
+
+        Parameters:
+            X (list of tuples): input signed measures.
+
+        Returns:
+            numpy array of shape (number of measures in **measures**) x (number of measures in X): matrix of pairwise Wasserstein distances.
+        """
+        return pairwise_signed_measure_distances(
+            X,
+            self.measures_,
+            metric="wasserstein",
+            epsilon=self.epsilon,
+            ground_norm=self.ground_norm,
+            n_jobs=self.n_jobs,
+        )
+
+    def __call__(self, meas1, meas2):
+        """
+        Apply WassersteinDistance on a single pair of signed measures and outputs the result.
+
+        Parameters:
+            meas1: ((n x D), (n)) tuple with numpy.array encoding the (finite points of the) first measure and their multiplicities. Must not contain essential points (i.e. with infinite coordinate).
+            meas2: ((m x D), (m)) tuple encoding the second measure.
+
+        Returns:
+            float: Wasserstein distance.
+        """
+        return _wasserstein_distance(
+            meas1, meas2, epsilon=self.epsilon, ground_norm=self.ground_norm
+        )

multipers/ml/tools.py

@@ -0,0 +1,113 @@
+from types import FunctionType
+from typing import Iterable
+
+import numpy as np
+from joblib import Parallel, delayed
+from sklearn.base import BaseEstimator, TransformerMixin
+
+import multipers as mp
+from multipers.simplex_tree_multi import SimplexTreeMulti
+
+
+def get_simplex_tree_from_delayed(x) -> mp.SimplexTreeMulti:
+    f, args, kwargs = x
+    return f(*args, **kwargs)
+
+
+def get_simplextree(x) -> mp.SimplexTreeMulti:
+    if isinstance(x, mp.SimplexTreeMulti):
+        return x
+    if len(x) == 3 and isinstance(x[0], FunctionType):
+        return get_simplex_tree_from_delayed(x)
+    else:
+        raise TypeError("Not a valid SimplexTree !")
+
+
+def filtration_grid_to_coordinates(F, return_resolution):
+    # computes the mesh as a coordinate list
+    mesh = np.meshgrid(*F)
+    coordinates = np.concatenate([stuff.flatten()[:, None] for stuff in mesh], axis=1)
+    if return_resolution:
+        return coordinates, tuple(len(f) for f in F)
+    return coordinates
+
+
+def get_filtration_weights_grid(
+    num_parameters: int = 2,
+    resolution: int | Iterable[int] = 3,
+    *,
+    min: float = 0,
+    max: float = 20,
+    dtype=float,
+    remove_homothetie: bool = True,
+    weights=None,
+):
+    """
+    Provides a grid of weights, for filtration rescaling.
+     - num parameter : the dimension of the grid tensor
+     - resolution :  the size of each coordinate
+     - min : minimum weight
+     - max : maximum weight
+     - weights : custom weights (instead of linspace between min and max)
+     - dtype : the type of the grid values (useful for int weights)
+    """
+    from itertools import product
+
+    # if isinstance(resolution, int):
+    try:
+        float(resolution)
+        resolution = [resolution] * num_parameters
+    except:
+        pass
+    if weights is None:
+        weights = [
+            np.linspace(start=min, stop=max, num=r, dtype=dtype) for r in resolution
+        ]
+    try:
+        float(weights[0])  # same weights for each filtrations
+        weights = [weights] * num_parameters
+    except:
+        None
+    out = np.asarray(list(product(*weights)))
+    if remove_homothetie:
+        _, indices = np.unique(
+            [x / x.max() for x in out if x.max() != 0], axis=0, return_index=True
+        )
+        out = out[indices]
+    return list(out)
+
+
+class SimplexTreeEdgeCollapser(BaseEstimator, TransformerMixin):
+    def __init__(
+        self,
+        num_collapses: int = 0,
+        full: bool = False,
+        max_dimension: int | None = None,
+        n_jobs: int = 1,
+    ) -> None:
+        super().__init__()
+        self.full = full
+        self.num_collapses = num_collapses
+        self.max_dimension = max_dimension
+        self.n_jobs = n_jobs
+        return
+
+    def fit(self, X: np.ndarray | list, y=None):
+        return self
+
+    def transform(self, X):
+        edges_list = Parallel(n_jobs=-1, prefer="threads")(
+            delayed(mp.SimplextreeMulti.get_edge_list)(x) for x in X
+        )
+        collapsed_edge_lists = Parallel(n_jobs=self.n_jobs)(
+            delayed(mp._collapse_edge_list)(
+                edges, full=self.full, num=self.num_collapses
+            )
+            for edges in edges_list
+        )
+        collapsed_simplextrees = Parallel(n_jobs=-1, prefer="threads")(
+            delayed(mp.SimplexTreeMulti._reconstruct_from_edge_list)(
+                collapsed_edge_lists, swap=True, expand_dim=self.max_dimension
+            )
+        )
+        return collapsed_simplextrees