[WIP] Clarify transported mass in partial OT

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 overflow to NaN in `ot.stochastic.coordinate_grad_semi_dual` by shifting before exponentiating, as `ot.stochastic.c_transform_entropic` already does. The factor cancels in the normalisation, so results are unchanged where the old code did not overflow (PR #866, Issue #264)

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
+<sphx_glr_auto_examples_unbalanced-partial_plot_partial_wass_and_gromov.py>`
+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

examples/unbalanced-partial/plot_partial_wass_and_gromov.py

@@ -15,6 +15,7 @@
 """
 
 # Author: Laetitia Chapel <laetitia.chapel@irisa.fr>
+# Mass-interpretation extension: Jialin Tang <jtang3898@sdsu.edu>
 # 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)