[MRG] entropic_partial_wasserstein: add method='sinkhorn_log' stable log-domain solver (rescue of #724)

[hinanohart]

Context

This is a rescue follow-up to PR #724 (and issue #723), both opened by @wzm2256 in March 2025. The original PR adds a numerically-stable entropic_partial_wasserstein_logscale to fix the NaN regime that #723 documents, but has been stuck CONFLICTING since 2025-09 — the author wrote in the PR body that they could not (a) build the docs locally (sphinx_rtd_theme missing) or (b) write pytest cases, and asked for help.

I hit the same NaN regime as a downstream user of entropic_partial_wasserstein, so this PR re-applies @wzm2256's function on top of current master and supplies the missing pieces.

Why a new PR rather than pushing to #724's branch

When PR #724 was opened, ot/partial.py was a single file. The maintainers later split it into the ot/partial/ package, so the diff between #724's branch and current master is now +2k / −13k lines of mostly unrelated movement. Re-applying just the new function on the new layout is cleaner than fighting that rebase.

If maintainers prefer to keep #724 as the canonical PR, I'm happy to close this and instead push these changes to that branch (write access permitting).

What's added (additive only; nothing removed)

• ot/partial/partial_solvers.py — entropic_partial_wasserstein_logscale function, identical algorithm to PR [WIP] Add a stablized function entropic_partial_wasserstein_logscale #724 re-applied on the current layout.
• ot/partial/__init__.py — export + __all__ entry.
• test/test_partial.py — 5 new test functions / 11 parametrised cases:
  • No NaN/Inf at small ε across [0.1, 0.05, 0.01, 5e-3, 1e-3, 5e-4] — reproduces issue entropic_partial_wasserstein not stable #723's failure mode and verifies the logscale solver stays finite.
  • Bit-for-bit agreement with the standard solver at large ε (atol=1e-10, rtol=1e-10) across [1.0, 10.0].
  • approaches_exact_at_small_reg — checks the plan-cost vs. exact partial-OT gap at reg=1e-3 (not just NaN-freeness; verifies mathematical correctness).
• examples/unbalanced-partial/plot_entropic_partial_wasserstein_logscale.py — Sphinx-Gallery example reproducing entropic_partial_wasserstein not stable #723's failure and showing the fix. Adapted from @wzm2256's compare_logscale_POT.py.
• docs/source/user_guide.rst — one-paragraph mention next to entropic_partial_wasserstein.
• RELEASES.md — one-line entry under 0.9.7.dev0, phrased as mitigation (the standard solver itself is unchanged; callers opt into the log-domain variant).

Local verification

• pytest test/test_partial.py -q → 19 passed (8 originals + 11 new parametrised cases).
• pytest test/ (full suite) → 1939 passed, 97 skipped, 6 xfailed — no regression outside partial.
• Example script runs end-to-end with MPLBACKEND=Agg. Standard solver returns NaN at reg ∈ {0.05, 0.01} on the 50×50 cost-scale-50 problem; logscale solver stays finite over the full sweep.
• At the time of the original commit, git diff master -- ot/partial/partial_solvers.py showed exactly one hunk: the new function block. (The follow-up below adds the method-dispatch refactor on top of that.)

Update (2026-06-01) — merged master + method parameter

Addressing @rflamary's review:

• Merged current master and resolved the conflict. The only conflict was in the RELEASES.md 0.9.7.dev0 "New features" list (the ot.sgot / ARM-wheel entries vs. the partial entry); the PR is now MERGEABLE.
• Added a method parameter following the ot.sinkhorn convention. entropic_partial_wasserstein(..., method=...) is now a thin wrapper: method='sinkhorn' (default) runs the classical multiplicative-domain solver and method='sinkhorn_log' the stable log-domain one. The classical iteration moved verbatim into a private _entropic_partial_wasserstein_sinkhorn, so the default path is byte-for-byte unchanged and the three internal ot.gromov / ot.partial callers (which pass a, b, M, reg, m positionally) are unaffected. Docs, the gallery example, RELEASES.md and tests were updated to cover the dispatch.

Co-authorship

Credit for the actual numerical work belongs to @wzm2256, and the commit carries a Co-authored-by: wzm2256 <wzm2256@qq.com> trailer to reflect that (matching this repo's convention). If you'd rather I drop the trailer or keep the acknowledgement to the PR description only, just say so.

cc @rflamary @cedricvincentcuaz — you both did master-merges on #724 earlier this year so it looked like a "we want this in" signal. Happy to do whatever's least friction here: this PR can be the canonical one, or I can close it and push to #724's branch.

The [WIP] prefix matches the original PR's convention and signals this is open for re-titling to [MRG] once you've had a look.

[rflamary]

Hello @hinanohart, thanks for this PR, we are OK with emrging this new rescue, we really appreciate your new contributions. I have fixed the failing tests and doc in master so I would appreciate if you merge master to this PR so that we can do a proper code review when all tests pass.

[rflamary]

Hello @hinanohart , could you merge and resolve conflicts please? I really don't want to rescue another rescue ;)

BTW I had alook at your code and I think it would be better to do like we did in ot.sinkhorn an add a method parameter that allows to choose between sinkhorn and sinkhorn_log to chose classical or log solver. That is making the classical partial solver a wrapper over the two solvers. Could you do that so that users can now use your new algorithm by just setting a parameter?

[codecov]

Codecov Report

❌ Patch coverage is 94.87179% with 8 lines in your changes missing coverage. Please review.
✅ Project coverage is 96.79%. Comparing base (fab3507) to head (f212ff4).

Additional details and impacted files

@@            Coverage Diff             @@
##           master     #811      +/-   ##
==========================================
- Coverage   96.80%   96.79%   -0.02%     
==========================================
  Files         118      118              
  Lines       23721    23877     +156     
==========================================
+ Hits        22963    23111     +148     
- Misses        758      766       +8     

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[hinanohart]

Done on both counts:

• Merged master and resolved the conflict. The only conflict was in the RELEASES.md 0.9.7.dev0 "New features" list (your ot.sgot / ARM-wheel entries vs. the partial entry); the PR is now MERGEABLE.
• Added a method parameter following the ot.sinkhorn convention. entropic_partial_wasserstein(..., method='sinkhorn'|'sinkhorn_log') is now a thin wrapper: 'sinkhorn' (default) runs the classical multiplicative-domain solver, 'sinkhorn_log' the stable log-domain one. The classical iteration moved verbatim into a private _entropic_partial_wasserstein_sinkhorn, so the default path is byte-for-byte unchanged and the three internal ot.gromov / ot.partial callers (which pass a, b, M, reg, m positionally) are unaffected.

Docs, the gallery example, RELEASES.md and tests were updated; the new tests cover the dispatch (default == classical, 'sinkhorn_log' == entropic_partial_wasserstein_logscale, case-insensitivity, log-dict forwarding, invalid-method ValueError). Happy to flip to [MRG] whenever you'd like a full review.

[hinanohart]

All CI is green on the merged branch — Lint, build/build_from_source, linux 3.11–3.14, macOS, Windows, the CircleCI docs build, and codecov/patch+codecov/project all pass. Flipping to [MRG]; ready for review/merge whenever convenient. Happy to address anything you spot.