From 3378251d0fcf84713950164b8f8cecfe0d0da45a Mon Sep 17 00:00:00 2001 From: Jialin Tang Date: Sun, 25 Oct 2026 19:42:22 -0700 Subject: [PATCH] Clarify transported mass in partial OT documentation --- RELEASES.md | 4 + docs/source/user_guide.rst | 79 ++++++++- .../plot_partial_wass_and_gromov.py | 158 ++++++++++++++++-- 3 files changed, 222 insertions(+), 19 deletions(-) diff --git a/RELEASES.md b/RELEASES.md index d01ff927f..02458b3ca 100644 --- a/RELEASES.md +++ b/RELEASES.md @@ -16,6 +16,10 @@ blocks, selectable with `sampling_slices` in `sliced_wasserstein_distance`, as described in [97] (PR #853) +#### Documentation + +- Clarify how transported mass is interpreted in partial optimal transport and extend the existing 2D example with a deterministic mass sweep and matched-mass summaries (Issue #890). + #### Closed issues - Fix duplicate function and class descriptions in generated module API pages while preserving top-level aliases and example backreferences (Issue #527) diff --git a/docs/source/user_guide.rst b/docs/source/user_guide.rst index ecb314a1a..6e3eb75af 100644 --- a/docs/source/user_guide.rst +++ b/docs/source/user_guide.rst @@ -774,8 +774,10 @@ implemented the main function :any:`ot.barycenter_unbalanced`. Partial optimal transport ^^^^^^^^^^^^^^^^^^^^^^^^^^^^ -Partial OT is a variant of the optimal transport problem when only a fixed amount of mass m -is to be transported. The partial OT metric between two histograms a and b is defined as [28]_: +Partial OT is a variant of the optimal transport problem in which only a fixed +amount of mass :math:`m` is transported. Each source and target weight is a +capacity: some of that weight can remain unused. For nonnegative histograms +:math:`a` and :math:`b`, the partial OT problem is defined as [28]_: .. math:: \gamma = \arg\min_\gamma <\gamma,M>_F @@ -785,13 +787,74 @@ is to be transported. The partial OT metric between two histograms a and b is de \gamma 1 \leq a\\ \gamma^T 1 \leq b\\ 1^T \gamma^T 1 = m \leq \min\{\|a\|_1, \|b\|_1\} - -Interestingly the problem can be casted into a regular OT problem by adding reservoir points -in which the surplus mass is sent [29]_. We provide a solver for partial OT -in :any:`ot.partial`. The exact resolution of the problem is computed in :any:`ot.partial.partial_wasserstein` -and :any:`ot.partial.partial_wasserstein2` that return respectively the OT matrix and the value of the -linear term. The entropic solution of the problem is computed in :any:`ot.partial.entropic_partial_wasserstein` + +Interpreting the transported mass +""""""""""""""""""""""""""""""""" + +The parameter ``m`` is an **absolute amount of mass**, in the same units as the +weights in ``a`` and ``b``. A feasible value cannot exceed +:math:`\min\{\sum_i a_i, \sum_j b_j\}`. When both histograms are normalized to +sum to one, ``m=0.5`` also means transporting half of each distribution's total +mass. Without that normalization, ``m`` is not a fraction: if the source and +target totals are two and three, ``m=0.5`` transports one quarter of the source +mass and one sixth of the target mass. To request a fraction :math:`\rho` of +the smaller total, set :math:`m=\rho\min\{\sum_i a_i,\sum_j b_j\}` explicitly. + +Mass is also different from a number of points. With 40 uniformly weighted +points in a normalized histogram, each point has weight :math:`1/40`, so +``m=0.5`` transports the equivalent of 20 points' total weight. It does not +require selecting exactly 20 points or returning 20 one-to-one matches. A +point's weight can be used partially or split across several target points; +the row and column sums of the returned plan describe how much weight each +point contributes or receives. For nonuniform weights, counting matched +points is an especially poor substitute for measuring transported mass. + +If both histograms have the same total mass :math:`s` and ``m=s``, all row and +column capacities must be filled, recovering the balanced OT constraints. +Thus ``m=1`` is the balanced endpoint only when **both totals are one**. If +the totals differ, transporting the smaller total still leaves unused mass +in the larger distribution. + +Choosing and inspecting the mass +"""""""""""""""""""""""""""""""" + +The mass controls a tradeoff. A small ``m`` can restrict transport to nearby +samples while omitting some genuine overlap. Increasing ``m`` requires more +mass to participate and can force matches involving remote samples. Partial +OT chooses a plan from the supplied weights and costs; it does not guarantee +outlier removal or choose ``m`` automatically. Expected overlap from the +application, or validation suited to that application, can inform the choice. +Inspecting several feasible masses helps show the consequences of that choice. + +The :ref:`2D partial Wasserstein example +` +keeps the same two distributions of 20 Gaussian samples and 20 remote samples +and compares exact plans for ``m=0.25``, ``0.5``, ``0.75``, and ``1``. It shows +the plans and the matched mass from the Gaussian and remote groups, separately +for the source and target. The synthetic inlier/outlier labels are used only +to summarize the results and are not passed to the solver. These categories +describe which groups participate, not point-correspondence accuracy. When +summarizing mass involving *either* outlier group, count an outlier-outlier +pair once: subtract the inlier-inlier block from the total transported mass, +instead of adding the source and target outlier masses. + +The raw transport cost :math:`\langle\gamma,M\rangle_F` is computed over the +mass actually transported. Costs at different values of ``m`` therefore do +not directly compare matching quality: a smaller cost can simply reflect +transporting less mass. Dividing by ``m`` gives a cost per unit of mass, but +the selected mass and pairs can still differ. Compare the plans and their +participating mass alongside the cost, using criteria appropriate to the task. + +Solvers +""""""" + +The problem can be cast into a regular OT problem by adding reservoir points +to receive the surplus mass [29]_. We provide solvers in :any:`ot.partial`. +The exact problem is solved by :any:`ot.partial.partial_wasserstein` and +:any:`ot.partial.partial_wasserstein2`, which return the OT matrix and the +value of the linear term, respectively. The entropic solution is computed +by :any:`ot.partial.entropic_partial_wasserstein` (see [3]_). Following the convention of :any:`ot.sinkhorn`, this solver takes a ``method`` parameter: ``method='sinkhorn'`` (default) runs the classical multiplicative-domain iterations, while ``method='sinkhorn_log'`` switches to a numerically stable log-domain solver diff --git a/examples/unbalanced-partial/plot_partial_wass_and_gromov.py b/examples/unbalanced-partial/plot_partial_wass_and_gromov.py index a42836f18..ba87aa6e3 100755 --- a/examples/unbalanced-partial/plot_partial_wass_and_gromov.py +++ b/examples/unbalanced-partial/plot_partial_wass_and_gromov.py @@ -15,6 +15,7 @@ """ # Author: Laetitia Chapel +# Mass-interpretation extension: Jialin Tang # License: MIT License # sphinx_gallery_thumbnail_number = 2 @@ -38,14 +39,15 @@ n_samples = 20 # nb samples (gaussian) n_noise = 20 # nb of samples (noise) +random_state = np.random.RandomState(42) mu = np.array([0, 0]) cov = np.array([[1, 0], [0, 2]]) -xs = ot.datasets.make_2D_samples_gauss(n_samples, mu, cov) -xs = np.append(xs, (np.random.rand(n_noise, 2) + 1) * 4).reshape((-1, 2)) -xt = ot.datasets.make_2D_samples_gauss(n_samples, mu, cov) -xt = np.append(xt, (np.random.rand(n_noise, 2) + 1) * -3).reshape((-1, 2)) +xs = ot.datasets.make_2D_samples_gauss(n_samples, mu, cov, random_state=random_state) +xs = np.append(xs, (random_state.rand(n_noise, 2) + 1) * 4).reshape((-1, 2)) +xt = ot.datasets.make_2D_samples_gauss(n_samples, mu, cov, random_state=random_state) +xt = np.append(xt, (random_state.rand(n_noise, 2) + 1) * -3).reshape((-1, 2)) M = sp.spatial.distance.cdist(xs, xt) @@ -72,7 +74,7 @@ print("Partial Wasserstein distance (m = 0.5): " + str(log0["partial_w_dist"])) print("Entropic partial Wasserstein distance (m = 0.5): " + str(log["partial_w_dist"])) -pl.figure(1, (10, 5)) +pl.figure(figsize=(10, 5)) pl.subplot(1, 2, 1) pl.imshow(w0, cmap="gray_r") pl.title("Partial Wasserstein") @@ -82,6 +84,138 @@ pl.show() +############################################################################# +# +# Interpret the transported mass on the same 2D samples +# ----------------------------------------------------- +# +# Both histograms have total mass one, so ``m`` is also the transported +# fraction here. With 40 uniformly weighted samples, each sample can supply +# or receive at most 1/40 of the mass. The parameter constrains total mass, +# rather than selecting a fixed number of points or one-to-one matches. +# +# We vary ``m`` using the exact partial Wasserstein solver. The first 20 +# samples in each distribution are Gaussian samples (called inliers below), +# and the last 20 are remote samples (called outliers). These synthetic +# labels are used only to summarize the plans; they are not solver inputs. +# A small mass can omit genuine overlap, while a large mass can require +# transporting remote samples. Partial OT does not guarantee outlier removal. + +mass_study = { + "masses": np.array([0.25, 0.5, 0.75, 1.0]), + "plans": [], + "source_weights": p.copy(), + "target_weights": q.copy(), + "cost_matrix": M.copy(), + "inlier_count": n_samples, + "balanced_plan": ot.emd(p, q, M), + "records": [], +} +tolerance = 1e-8 + +for mass in mass_study["masses"]: + plan = ot.partial.partial_wasserstein(p, q, M, m=mass) + assert np.isfinite(plan).all() and (plan >= 0).all() + assert (plan.sum(axis=1) <= p + tolerance).all() + assert (plan.sum(axis=0) <= q + tolerance).all() + assert np.isclose(plan.sum(), mass, atol=tolerance, rtol=0) + mass_study["plans"].append(plan) + mass_study["records"].append( + { + "mass": float(mass), + "cost": float(np.sum(plan * M)), + "source_inlier_mass": float(plan[:n_samples, :].sum()), + "source_outlier_mass": float(plan[n_samples:, :].sum()), + "target_inlier_mass": float(plan[:, :n_samples].sum()), + "target_outlier_mass": float(plan[:, n_samples:].sum()), + # Count a pair involving either outlier group only once. + "any_outlier_mass": float(plan.sum() - plan[:n_samples, :n_samples].sum()), + } + ) + +# At full mass, these normalized histograms recover balanced OT. Optimal +# plans need not be unique, so compare the marginals and objective values. +full_mass_plan = mass_study["plans"][-1] +assert np.allclose(full_mass_plan.sum(axis=1), p, atol=tolerance, rtol=0) +assert np.allclose(full_mass_plan.sum(axis=0), q, atol=tolerance, rtol=0) +assert np.isclose( + np.sum(full_mass_plan * M), + np.sum(mass_study["balanced_plan"] * M), + atol=tolerance, + rtol=0, +) + +print("\nExact partial Wasserstein mass study (raw transport costs):") +print(" m cost source in source out target in target out any out") +for record in mass_study["records"]: + print( + f'{record["mass"]:4.2f} {record["cost"]:9.4f}' + f' {record["source_inlier_mass"]:10.4f}' + f' {record["source_outlier_mass"]:10.4f}' + f' {record["target_inlier_mass"]:10.4f}' + f' {record["target_outlier_mass"]:10.4f}' + f' {record["any_outlier_mass"]:10.4f}' + ) + +############################################################################# +# +# The plan matrices show which sample groups are coupled. Dashed lines +# separate the Gaussian and remote samples on each axis. All four panels +# use the same color scale, in units of mass per source-target pair. + +fig_plans, plan_axes = pl.subplots( + 1, 4, figsize=(12, 3.6), sharex=True, sharey=True, constrained_layout=True +) +for ax, mass, plan in zip(plan_axes, mass_study["masses"], mass_study["plans"]): + image = ax.imshow(plan, cmap="Blues", vmin=0, vmax=p.max()) + ax.axhline(n_samples - 0.5, color="gray", linestyle="--", linewidth=1) + ax.axvline(n_samples - 0.5, color="gray", linestyle="--", linewidth=1) + ax.set_title(f"m = {mass:g}") + ax.set_xlabel("Target sample index") + ax.set_xticks([0, n_samples, n_samples + n_noise - 1]) +plan_axes[0].set_ylabel("Source sample index") +fig_plans.suptitle("Exact partial OT: indices 0-19 are inliers, 20-39 are outliers") +fig_plans.colorbar(image, ax=plan_axes, label="Mass per pair", shrink=0.75) +pl.show() + +############################################################################# +# +# Summarize the mass used from each distribution separately. Each stacked +# bar has height ``m``. Mass involving any outlier is reported in the table +# as total mass minus the inlier-inlier block, since adding the source and +# target outlier masses would count outlier-outlier pairs twice. +# +# These summaries describe participation of the synthetic groups, not +# point-correspondence accuracy. Raw costs at different masses are not a +# direct comparison of matching quality, because the amount transported +# changes. The solver also does not choose ``m`` automatically. + +fig_mass, mass_axes = pl.subplots( + 1, 2, figsize=(9, 4), sharey=True, constrained_layout=True +) +for ax, side in zip(mass_axes, ("source", "target")): + inlier_masses = [record[f"{side}_inlier_mass"] for record in mass_study["records"]] + outlier_masses = [ + record[f"{side}_outlier_mass"] for record in mass_study["records"] + ] + ax.bar(mass_study["masses"], inlier_masses, width=0.15, label="Gaussian (inlier)") + ax.bar( + mass_study["masses"], + outlier_masses, + bottom=inlier_masses, + width=0.15, + label="Remote (outlier)", + ) + ax.set_title(f"{side.capitalize()} matched mass") + ax.set_xlabel("Requested transported mass m") + ax.set_xticks(mass_study["masses"]) + ax.set_ylim(0, 1.05) +mass_axes[0].set_ylabel("Matched mass") +mass_axes[1].legend(loc="upper left") +mass_study["figures"] = {"plans": fig_plans, "mass_accounting": fig_mass} +pl.show() + + ############################################################################# # # Sample one 2D and 3D Gaussian distributions and plot them @@ -104,11 +238,13 @@ cov_t = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]]) -xs = ot.datasets.make_2D_samples_gauss(n_samples, mu_s, cov_s) -xs = np.concatenate((xs, ((np.random.rand(n_noise, 2) + 1) * 4)), axis=0) +xs = ot.datasets.make_2D_samples_gauss( + n_samples, mu_s, cov_s, random_state=random_state +) +xs = np.concatenate((xs, ((random_state.rand(n_noise, 2) + 1) * 4)), axis=0) P = sp.linalg.sqrtm(cov_t) -xt = np.random.randn(n_samples, 3).dot(P) + mu_t -xt = np.concatenate((xt, ((np.random.rand(n_noise, 3) + 1) * 10)), axis=0) +xt = random_state.randn(n_samples, 3).dot(P) + mu_t +xt = np.concatenate((xt, ((random_state.rand(n_noise, 3) + 1) * 10)), axis=0) fig = pl.figure() ax1 = fig.add_subplot(121) @@ -137,7 +273,7 @@ print("Wasserstein distance (m = 1): " + str(log0["partial_gw_dist"])) print("Entropic Wasserstein distance (m = 1): " + str(log["partial_gw_dist"])) -pl.figure(1, (10, 5)) +pl.figure(figsize=(10, 5)) pl.title("mass to be transported m = 1") pl.axis("off") pl.subplot(1, 2, 1) @@ -161,7 +297,7 @@ print("Partial Wasserstein distance (m = 2/3): " + str(log0["partial_gw_dist"])) print("Entropic partial Wasserstein distance (m = 2/3): " + str(log["partial_gw_dist"])) -pl.figure(2, (10, 5)) +pl.figure(figsize=(10, 5)) pl.title("mass to be transported m = 2/3") pl.axis("off") pl.subplot(1, 2, 1)