Overview of the stratified Wasserstein framework

We propose stratified Wasserstein, a framework that embeds into Euclidean space each shape, represented as an unstructured point cloud, and facilitates kernel methods in that space for population-level quantification tasks. Compared to existing shape distances such as Gromov–Wasserstein and its lower bounds [12] as defined in Eqs M3, M4, and M5 (Methods), the proposed framework achieves similar discriminative power while being computationally more efficient, with complexity nearly quadratic in the number of points and empirical runtimes typically below 1% of those required by Gromov–Wasserstein methods. The induced kernel is characteristic, so population-level statistics via kernel methods retain their standard statistical guarantees, including consistency and power against alternatives [18,19,25]‌‌. Fig 1 provides a schematic overview of our procedure, that takes local distance distribution to produce shape embeddings and derive population-level statistics. Here we describe the main construction and summarize its key properties; detailed statements and proofs are deferred to section Methods and the Supplementary Information.

Download:

• PNG
  larger image
• TIFF
  original image

Fig 1. A summary of our procedure.

Given unstructured point cloud in 2D or 3D representing a shape of a wide range of types, we first compute the intrinsic distances on each shape and localized for each point. Using these local distances, a wide range of downstream tasks can be performed given a large population of shapes, such as dimension reduction, clustering, hypothesis testing, and feature selection. The detailed methodology is presented in subsequent sections, with rigorous theoretical foundation discussed.

https://doi.org/10.1371/journal.pcbi.1014254.g001

First, we represent each shape by a point cloud together with an intrinsic distance, such as the geodesic distance on the surface or the Euclidean distance in ambient space. For every point on the shape, we compute the distribution of distances to all other points, which captures its local geometric structure. We then rank the points according to a continuous functional of these local distance distributions, which induces a one-dimensional stratification of the shape. Within each stratum, we summarize the local geometry by recording a set of quantiles of the distance distributions. This defines an embedding of the shape into a two-dimensional function space indexed by the stratification variable and the quantile level. The stratified Wasserstein distance between two shapes is simply the distance between their embeddings, computed over this two-dimensional domain.

In this work, we show that the stratified Wasserstein distance inherits many of the desirable properties of Gromov–Wasserstein while significantly improving computational efficiency. It is naturally invariant to isometric transformations because of the invariance property of local distance distributions. Under mild regularity conditions on the ranking functional used for disintegrating the shape, the distance is injective up to isometry: two shapes have zero distance if and only if they are isometric. Because the distance is induced by an norm in the embedded space, standard kernels such as the Gaussian and Laplace built from it are positive definite. When the distance is injective, these kernels are also characteristic, meaning they can distinguish any two distributions of shapes. The estimator of the distance based on finite point clouds is statistically consistent: as the number of points per shape grows, with discretization refined accordingly, the estimated distance converges to the population distance. The computational cost scales nearly quadratically with the number of points per shape, which is far more efficient than Gromov–Wasserstein methods.

Finally, because the kernel is characteristic, standard population-level statistics built from it inherit strong theoretical guarantees. For example, maximum mean discrepancy (MMD) [18] measures any form of discrepancy between two populations, while the Hilbert–Schmidt Independence Criterion (HSIC) [19] measures the strength of dependence between variables. Both statistics are consistent and have power against all fixed alternatives, which enables rigorous nonparametric two-sample testing and dependence detection between shape distributions and external covariates.

We provide Table 1 to compare our proposed stratified distance between shapes against other relative alternatives. Compared to the other distances, stratified Wasserstein is Hilbertian, can be consistently estimated from finite samples, and achieves the lowest asymptotic computational complexity. It is conditionally injective and yields a characteristic kernel if the oracle functional for sorting and binning is chosen carefully. The balance between discriminative power and computational complexity is further demonstrated in numerical examples on real world biological shape datasets as discussed in the next section. Empirical comparison on these distances on simple synthetic shapes can be found in S3 and S4 Figs.

Download:

• PNG
  larger image
• TIFF
  original image

Table 1. Comparison between distances with shapes represented as point clouds in 2D/3D. For Gromov–Wasserstein and Wasserstein between local distances, computational cost for both exact computation (left) and entropic approximation (right) are both shown. N = points per shape. ^*Exact GW is widely believed to be NP-hard [16]; ^†Entropic GW complexity depends on (regularization) and (tolerance) [26], both are small, positive values for accurate approximation; ^‡Systematic studies on computational costs can be found in [27–29]. ^§ Our framework has a conditionally injective and characteristic kernel under Hypotheses (H1) and (H2) with details in Supplementary Information.

https://doi.org/10.1371/journal.pcbi.1014254.t001

Breast cancer cell contour shapes from fluorescence microscopy

We apply our framework to 2D cancer cell shapes obtained from fluorescence microscopy [32,33]. Cell images are binarized, and their boundaries are extracted to form discrete curves given by the 2D coordinates of the cell contours. The dataset includes cancer cell shapes from three different cell lines for breast cancer, corresponding to: (1) MCF10A (228 cells): non-tumorigenic human breast mammary gland epithelial cell line, which is a classic model for normal cells, (2) MCF7 (225 cells): breast cancer line with relatively low metastatic potential, and (3) MDA-MB-231 (abbreviated MDA, 224 cells): a triple-negative breast cancer cell line that is highly invasive and commonly used as a model for metastatic progression. Given populations of these cell contours, the goal is to test whether the underlying shape distributions differ significantly across groups defined by cell lines. Understanding these shape differences may help reveal whether and how cell morphology encodes functional behaviors for different cancer types. An example of the raw images, as well as random samples of cells from each cell line are visualized in Fig 2A and 2B, and more images can be found in S1 Fig. Implementation details can be found in Methods and Table 2.

Download:

• PNG
  larger image
• TIFF
  original image

Table 2. Summary of experimental parameters and implementation details for 2D and 3D shape analyses.

https://doi.org/10.1371/journal.pcbi.1014254.t002

Download:

• PNG
  larger image
• TIFF
  original image

Fig 2. Case study on 2D cancer cell shape populations represented by contours.

A: Example fluorescence microscopy images of cell lines curated by [32]. B: Cell contours after segmentation, with a random sample of 8 shapes from each population visualized. C: Inspection of first and second moments suggests similarity between the MCF10A and MDA populations, as shown via Fréchet mean shape and distances to the mean shape within each population under the SRV metric, computed using Geomstats [10]. D: To further probe differences between MCF10A and MDA, we compute the MMD-based test statistic and its p-value using our proposed stratified kernel, leveraging all available cell shapes. E: Cells with the top 8 witness scores in each population highlight those whose shapes differ most from the respective population majority. F: Empirical probability of Type II error and computational time under small (15) to moderate (60) sample sizes, benchmarked across various kernel-based tests, with results averaged over five independent trials (error bars: ± 1 SD).

https://doi.org/10.1371/journal.pcbi.1014254.g002

Relative MMD test reveals similarity between triple negative cancer cells and non-cancerous cells: Upon testing relative MMD between three populations, we observed that: Despite being a cancer line, MDA resembles MCF10A (normal cells) more than MCF7 (low-metastasis cancer). The empirical estimate of squared MMD reveals that MMD²(MCF10A,MDA)=0.0189, much smaller than the estimate of MMD²(MCF7,MDA)=0.1409. To test whether the difference is statistically significant, we perform a relative MMD test for the following:

The p-value is computed using a permutation test that pools MCF7 and MCF10A, under the null hypothesis that the two populations are equally distant from MDA. Specifically, we permute the MCF7/MCF10A labels and recompute the difference in squared MMD to MDA for each shuffle. With 1000 permutations, this test results in a p-value of 0.001 (S1C Fig), suggesting that the MDA line (high metastatic cancer) is morphologically closer to MCF10A (normal cell line) than to MCF7 (low metastatic cancer). This observation is consistent to findings from gene expression studies [34], which reported that both MCF10A and MDA share a basal-like subtype, while MCF7 belongs to the luminal class. Our shape-based results suggest that cell morphology is more closely aligned with molecular subtype (basal versus luminal) than with cancer status.

The result is also consistent with the prior approach [35] that computes Fréchet mean shape under the square root velocity (SRV) metric [6] through geomstats [10] package. The mean shapes of processed contours from all 3 populations are shown in S1B Fig, where MCF10A and MDA mean shapes are both elongated and almost coincide.

Absolute MMD test indicates discrepancy between populations on the tails: Given almost identical mean shapes and distribution of distance to mean shown in Fig 2C, it seems that the two shape populations of MDA and MCF10A agree up to mean and variance. A natural follow-up question is whether the generating distributions of the two shape populations are truly identical. To this end, we zoom in on the unexpected pair and perform an MMD test on those two populations, which detects arbitrary order of discrepancy in distribution that extends beyond second order:

With 1,000 permutations, it shows a p-value of 0.002 (Fig 2D), suggesting a significant difference in the two populations. To localize the morphological differences, we examine extreme shapes identified by the MMD witness function, which highlights regions where the shape distributions differ most. Top 8 cells in each population with the highest values of the MMD witness function are visualized in Fig 2E. The extremal subset of MDA cells exhibits highly irregular and protrusive morphologies, which are key contributors to the shape-based distinction from MCF10A. We hypothesize that such enriched shapes could correspond to invasive behaviors as known for MDA [36]. Notably, the statistical difference in MMD is not driven by the mean shape or the distances to mean shape (Fig 2C), but by this tail population in MDA, which can only be captured using nonparametric test like the kernel MMD.

The tests are computationally efficient and statistically well-powered: To assess the power of the absolute and relative tests, we perform randomized experiments to quantify the empirical Type I and II error rates. We compare our kernels to the framework suggested by Zhang et al [37] using square-root velocity (SRV) metric and its variant (Elastic) for MMD. To ensure a fair comparison with methods that require smoother, parameterized curves, the benchmark approaches involving SRV and elastic metrics are applied to preprocessed data, where each shape has been interpolated to have 2,000 points and aligned to the same reference shape beforehand. In contrast, in our approaches, we perform the test directly on raw data without any processing.

We perform the tests for sample sizes varying between 15 and 60 via MMD. To systematically quantify the error, for each sample size, we resample and repeat the experiments 1,000 times, compute the p-value via permutation test, and report the relative frequency of making an error in Fig 2F. For Type II error, we draw one sample from each group, perform the test, and compute the relative frequency that the equality of shape distributions cannot be rejected. Among our proposed distances, Gromov–Wasserstein consistently achieves the lowest Type II error, followed by stratified Wasserstein distances and the second lower bound (SLB) of Gromov–Wasserstein. All Wasserstein-type distances outperform SRV-based metrics at sample sizes under 40, which is consistent with the fact that Wasserstein-type distances operate on distributions of intrinsic distances and are less sensitive to curve parameterization and discretization effects. Similar experiments and observation on error rates for both Type I and II are done for relative test and shown in S1 Fig. We also note that the error rate is robust to the choices of hyperparameters, and sensitivity results are reported in S2 Fig.

To assess computational efficiency, we record the wall time required to compute all pairwise distances between shapes in each subsample, averaged over five independent runs. These results are shown in Fig 2. While GW offers the best statistical power, its computational cost is substantial. Our stratified Wasserstein achieves a favorable balance: while offering second best Type II error control, slightly worse than Gromov–Wasserstein, it is significantly more efficient than Gromov–Wasserstein and its second lower bound, and is second only to the pooled Wasserstein method, which simply aggregates and sorts all pairwise distances.

In sum, our nonparametric, distribution-level approach operates directly on raw contour shapes, detects tail effects, and localizes the drivers of discrepancy while remaining computationally practical.

Allen institute 3D cell and nucleus shapes

In this example, we demonstrate that quantiles of distance matrices encode meaningful morphological information, using 3D cell imaging data from Viana et al. [1]. We analyze a subset of 5,764 cells used by the original authors to train a classifier, including annotations for cell types reflecting six mitotic stages (M0, M1M2, M3, M4M5, M6M7 early, and M6M7 half) as well as three outlier types (blob, dead, wrong). Implementation details on our method can be found in Methods and Table 2. Existing dimension reduction methods rely either on hand-crafted features (e.g., volume and height of the cell and nucleus) [1] or on learned latent representations from deep models such as variational autoencoders [38], followed by UMAP. However, these methods are either highly specialized for the application or require significant training effort.

Dimension reduction for shapes reveals relevant morphological variations.

We pre-process each binary image by eroding it to extract surface points, followed by down-sampling to retain 200 points per cell. For each cell, we divide it into 100 bins, compute 100 quantiles for both the cell shape and the nucleus shape in each bin, concatenate them, and apply UMAP to the resulting 20,000-dimensional vector. Fig 3A shows 2D UMAP embeddings obtained using stratified Wasserstein distance. To benchmark our result, we performed dimension reduction using 4 additional methods and visualized them in Fig 3B that are based on (1) Euclidean distance in raw binary images, (2) dominant features through volumes of cell and nucleus, (3) the intermediate latent layer using a pre-trained PointNet model [39], and (4) Gromov-Wasserstein distance with each cell downsampled to 100 points due to prohibitive computational cost. For a fair comparison, we also applied the proposed stratified Wasserstein embedding to shapes downsampled to 100 points. As shown in S4 Fig, the resulting embedding is qualitatively similar to that obtained with 200 points. Our dimension reduction best preserves the cyclic and continuous nature of the shape dynamics, consistent with the feature-based and PointNet embeddings, but with a more smooth trajectory that reflects the remark on the ambiguity of manual annotation by Viana et al. [1].

Download:

• PNG
  larger image
• TIFF
  original image

Fig 3. Case study on 3D cell and nucleus shape populations.

A: UMAP embedding of the annotated mitotic dataset from Viana et al. [1] through stratified local distance distributions, and representative cell (blue) and nucleus (red) shapes along the cycle. B: Cell-level embeddings from standard methods fail to reveal dynamics. Euclidean distances between raw binary images and Gromov–Wasserstein distances on downsampled point clouds show no meaningful pattern, while feature-based embeddings (volumes or neural networks) capture heterogeneity but not the cyclic continuum of mitosis. C: Cell features displayed against UMAP allows for interpretation of shape changes within each stage. D: Population-level embedding using proposed MMD reveals cyclic process consistent to A. E: HSIC witness scores for individual cells, revealing regions with mixed positive and negative values near late mitotic stages. We selected 3 cells with high positive witness scores (yellow crosses) and 3 with high negative witness scores (magenta stars), and display the nucleus shapes on the sides.

https://doi.org/10.1371/journal.pcbi.1014254.g003

This embedding provides biological insight into morphological changes during the mitotic cycle. In Fig 3C, we color cells by morphological features to show how these evolve along the cycle. Starting from the M1M2 phase, both cell and nucleus volumes increase significantly, followed by a marked rise in nucleus height during M3 and M4M5. These trends align with the known stages of mitosis, where volume increase reflects DNA and organelle replication, and structural transitions (e.g., chromatin condensation and nuclear envelope breakdown) are hallmarks of the later phases [40]. The outlier cells have particularly smaller volumes and heights in both cell and nucleus shapes, likely reflecting incomplete mitosis, abnormal states, or segmentation artifacts.

Shape representation and intrinsic distances

We represent each shape as an unstructured point cloud with . Each shape is equipped with an intrinsic distance and a uniform measure . Distances can be chosen flexibly, depending on the application, including Euclidean distance, geodesic distance on a k-nearest neighbor graph [41], or diffusion distance [42]. The intrinsic distance matrix is .

A shape statistic central to this work is the local distance distribution, defined as follows. For each point , the local distance distribution is the pushforward measure

where denotes the uniform probability measure on and # denotes the pushforward operator. In the discrete setting of a point cloud , this reduces to the empirical measure

(M1)

For comparison, one may also define the global distance distribution of a shape as the pushforward

(M2)

which captures the overall distribution of pairwise distances on the shape. In the discrete setting, this corresponds to the empirical distribution of all pairwise distances . The local distance distributions are known to have better discriminative power [12,14] compared to global distance distributions for their ability to retain pointwise geometric information and are not recoverable from alone.

For scale invariance, shapes are rescaled so that the median pairwise distance equals one.

Gromov–Wasserstein and related distances

Gromov–Wasserstein distance.

For two metric measure spaces , i = 1,2, the p–Gromov–Wasserstein (GW) distance [12,43] is defined by

(M3)

In the discrete setting, this corresponds to minimizing over coupling matrices between points. Exact computation is widely believed to be NP-hard [16], and entropic regularization is typically used to obtain approximate solutions with complexity [26]. In all numerical experiments, we used the implementation provided by Python Optimal Transport [44,45].

Global Wasserstein distance.

A simpler lower bound compares the global distributions of pairwise distances [13,15,46]. Let and be the quantile functions of the empirical global distance distributions defined as in Eq (M2). The p–Wasserstein distance between these distributions is

(M4)

This distance can be computed efficiently in , but it is not injective: different shapes may yield the same global distance distribution. An illustrating example [12] can be found in Fig A in S3 Text.

Second lower bound (SLB).

A stronger metric compares local distance distributions [12,14]:

(M5)

where is the local distance distribution at point x and is a coupling between points. SLB is strictly stronger than the global Wasserstein distance and is injective for smooth closed shapes under regularity assumptions [14], but remains computationally demanding and is not Hermitian.

Stratified Wasserstein distance

The proposed stratified Wasserstein distance combines the geometric discriminative power of GW-type distances with the efficiency of quantile-based embeddings. For each point , we compute its local distance distribution , and further summarize it by a scalar ranking functional , which represents the 1D distribution in a way that reflects the point’s geometric context within the shape. In practice, we use the mean of local distance distributions, which tends to separate points near the shape barycenter from those near peripheral or protrusive regions.

To make this descriptor comparable across shapes and invariant to monotone transformations, we convert into a normalized rank

where denotes the cumulative distribution function of for . This transformation maps points on the shape to the unit interval [0,1], inducing a one-dimensional stratification of the shape. The stratification preserves the ordering of points by geometric context while discarding absolute scale information, enabling consistent aggregation of local geometry across shapes.

The shape measure is then disintegrated along ,

where denotes the conditional law at stratum u. We define the function Q by evaluating the quantile function of at level :

where denotes the (left-continuous) generalized inverse of . For each stratum and quantile level , we compute

Overall, we have defined an embedding . The stratified Wasserstein distance between two shapes is the distance between their embeddings:

Under mild regularity conditions on the ranking functional, the distance is injective up to isometry (Theorem 2 in S1 Text). Since it is induced by an L² norm, Gaussian and Laplace kernels constructed from D₂ are positive definite, and if the embedding is injective, these kernels are characteristic. The empirical estimator of the distance, based on binning and quantiles, is statistically consistent in the joint regime where the number of sampled points per shape grows while the discretization is refined. Under regularity of the ranking functional, which ensures stability of ranks, and a bin-growth condition, which ensures sufficient samples per stratum, the estimated distance converges to the population distance (see S2 Text). Its computational complexity is nearly quadratic in the number of points, offering substantial savings over GW and SLB.

Benchmarks used for comparison

Elastic metrics between 2D planar curves.

In the task on 2D cancer contour shapes, for comparison, we used the elastic metric with two choices of parameters [47–49]. Let be a smooth parameterized curve, and be two tangent vector fields along . The elastic metric with parameters a, b > 0 is defined by

where s denotes the arc-length parameter of , and

is the derivative of h with respect to arc-length. The square root velocity has a = 1/2 and b = 1, while the elastic metric implemented by Geomstats [10] has default a = 1 and b = 1/2. For the 2D cell line dataset, we perform benchmarking via elastic and square root velocity metrics on curves post alignment and interpolated to consist 2,000 points each, which can be accessed from https://github.com/wxli0/dyn/tree/main%4092c7a58/dyn/datasets/breast_cancer/aligned/projection_rescale_rotation_reparameterization.

Kernel methods on shape populations

Positive-definite kernels constructed from the stratified Wasserstein distance enable population-level statistical analysis using kernel methods. We focus on maximum mean discrepancy (MMD) and Hilbert–Schmidt independence criterion (HSIC), which are widely used nonparametric statistics with strong theoretical guarantees.

Maximum mean discrepancy (MMD).

MMD [18] is the squared distance between kernel mean embeddings of two distributions and in the reproducing kernel Hilbert space (RKHS) associated with kernel k:

Given samples and , the (biased) empirical estimator is

Significance is assessed by permutation testing of group labels. When k is characteristic (as is the case for Gaussian or Laplace kernels on the stratified Wasserstein metric), MMD equals zero if and only if . The test is consistent against all fixed alternatives, meaning its power converges to one as sample sizes grow [18,51].

MMD witness function.

The MMD witness function [18]

identifies regions of shape space where the two distributions differ most strongly. Evaluating at sample shapes highlights representative shapes that contribute to the observed differences.

Hilbert–Schmidt independence criterion (HSIC).

HSIC [19] quantifies dependence between random variables using kernels on each domain, and can be interpreted as the maximum mean discrepancy (MMD) between the joint distribution and the product of its marginals . For paired data , with characteristic kernels k on shapes and ℓ on covariates, HSIC is defined as the squared Hilbert–Schmidt norm of the cross-covariance operator in the associated RKHSs:

With Gram matrices , , and centering matrix , the (biased) estimator is

HSIC is zero if and only if and Y are independent, provided both kernels are characteristic. It is consistent against all alternatives and detects arbitrary nonlinear dependencies [19,21].

HSIC witness function.

Similar to MMD, HSIC admits a witness function that localizes the contribution of individual shape–covariate pairs to the overall dependence. Given kernels k and ℓ centered in RKHS, the empirical HSIC witness is:

where and are the centered kernels. Positive values indicate that the pair (s,y) supports the observed dependence, while negative values indicate opposition. It is precisely the MMD witness function applied to the joint distribution and the product of marginals. This function can be used to interpret the contributions of specific shapes or covariate values.

Statistical analyses

All population-level comparisons were performed using kernel methods derived from the stratified Wasserstein distance and other benchmarking distances.

Two-sample tests between shape populations were carried out using the maximum mean discrepancy (MMD) statistic [18]. For each test, we used a permutation test with 1,000 permutations to estimate the null distribution, and reported p-values based on the proportion of permuted statistics exceeding the observed value. The MMD witness function was used to localize representative shapes contributing to significant differences between groups.

To assess dependence between shape distributions and external covariates (e.g., developmental stage), we used the Hilbert–Schmidt independence criterion (HSIC) [19] with Laplace kernels on both shape and covariate domains, with bandwidth chosen via the median distance.

For clarity, all the tests in this section are performed using Laplace kernels between shapes with median bandwidth heuristic:

where represents a distance in shape space, such as Gromov–Wasserstein and its alternatives.

For the 3D mitotic cell shape data, the mitotic stage variable consists of six ordered categories (M0, M1M2, M3, M4M5, M6M7_early, M6M7_half), which were encoded as integers . The mitotic stage label was encoded as an angular variable , and mapped to circular coordinates on the unit circle. A Gaussian kernel applied to these circular embeddings,

is equivalent (up to a constant factor) to a von Mises kernel (under a rescaling of )

The bandwidth parameter was selected using the median heuristic, that is, was set to the median of squared pairwise distances between stage embeddings on the unit circle. This encoding preserves the cyclic topology of the mitotic progression while allowing the stage variable to interact naturally with Euclidean distances between shapes.

Permutation testing over 1,000 trials of the covariate labels was used to assess significance. The HSIC witness function was used to visualize shape–covariate pairs with the strongest contribution to the dependence.

Computational details

All experiments were implemented in Python (v3.13.7). Numerical computations used NumPy (v2.3.4). Optimal transport computations were performed using the POT library (v0.9.6). Kernel-based statistical tests relied on scikit-learn (v1.7.2), dimensionality reduction used UMAP (v0.5.9), and elastic shape analysis used in benchmarks employed Geomstats (v2.8.0).

Timing experiments were conducted on a MacBook Pro (14-inch, 2024) equipped with an Apple M4 Pro chip and 48 GB unified memory, running macOS Sequoia (v15.7.3). All methods were evaluated under identical hardware and software conditions; reported runtimes are therefore intended for relative comparison rather than absolute benchmarking.