Partial Optimal Transport

def ot.partial.partial_wasserstein(a, b, M, m=None, numItermax=1000000, log=False, **kwargs):
    """
    Solve partial optimal transport problem.
    
    Computes the optimal transport plan that transports at most mass m between
    source and target distributions, relaxing the total mass conservation constraint.
    This is useful for robust transport in the presence of outliers.
    
    Parameters:
    - a: array-like, shape (n_samples_source,)
         Source distribution. Must be non-negative.
    - b: array-like, shape (n_samples_target,)
         Target distribution. Must be non-negative.
    - M: array-like, shape (n_samples_source, n_samples_target)
         Ground cost matrix between source and target samples.
    - m: float, optional
         Fraction of mass to transport. Must be in (0, min(sum(a), sum(b))].
         If None, defaults to min(sum(a), sum(b)).
    - numItermax: int, default=1000000
         Maximum number of iterations for the underlying solver.
    - log: bool, default=False
         Return optimization log with convergence details.
    - kwargs: dict
         Additional arguments passed to the underlying linear programming solver.
    
    Returns:
    - transport_plan: ndarray, shape (n_samples_source, n_samples_target)
         Partial optimal transport plan. The sum of the plan equals the transported mass m.
    - log: dict (if log=True)
         Contains 'cost': optimal transport cost, 'u': source dual variables,
         'v': target dual variables, 'result_code': solver status.
    
    Example:
        a = np.array([0.5, 0.5])
        b = np.array([0.3, 0.4, 0.3])
        M = np.random.rand(2, 3)
        plan = ot.partial.partial_wasserstein(a, b, M, m=0.7)
        print(f"Transported mass: {np.sum(plan)}")
    """

def ot.partial.partial_wasserstein2(a, b, M, m=None, numItermax=1000000, log=False, **kwargs):
    """
    Solve partial optimal transport and return cost only.
    
    More efficient than partial_wasserstein() when only the optimal cost is needed.
    
    Parameters: Same as partial_wasserstein()
    
    Returns:
    - cost: float
         Partial optimal transport cost (optimal objective value).
    - log: dict (if log=True)
         Optimization information.
    
    Example:
        cost = ot.partial.partial_wasserstein2(a, b, M, m=0.7)
        print(f"Partial transport cost: {cost}")
    """

def ot.partial.partial_wasserstein_lagrange(a, b, M, reg_m=None, numItermax=1000000, log=False, **kwargs):
    """
    Solve partial optimal transport using Lagrangian relaxation.
    
    Alternative formulation where the mass constraint is regularized using
    a Lagrangian multiplier instead of being enforced as a hard constraint.
    
    Parameters:
    - a: array-like, shape (n_samples_source,)
         Source distribution.
    - b: array-like, shape (n_samples_target,)
         Target distribution.
    - M: array-like, shape (n_samples_source, n_samples_target)
         Cost matrix.
    - reg_m: float, optional
         Regularization parameter for mass constraint. Higher values enforce
         stronger mass conservation.
    - numItermax: int, default=1000000
         Maximum iterations.
    - log: bool, default=False
    - kwargs: dict
         Additional solver arguments.
    
    Returns:
    - transport_plan: ndarray
         Partial transport plan with regularized mass constraint.
    - log: dict (if log=True)
    """

def ot.partial.entropic_partial_wasserstein(a, b, M, reg, m=None, numItermax=1000, stopThr=1e-9, verbose=False, log=False):
    """
    Solve entropic regularized partial optimal transport.
    
    Combines partial transport with entropic regularization for better
    computational properties and differentiability.
    
    Parameters:
    - a: array-like, shape (n_samples_source,)
         Source distribution.
    - b: array-like, shape (n_samples_target,)
         Target distribution.
    - M: array-like, shape (n_samples_source, n_samples_target)
         Cost matrix.
    - reg: float
         Entropic regularization parameter. Higher values give smoother solutions.
    - m: float, optional
         Mass to transport. If None, uses min(sum(a), sum(b)).
    - numItermax: int, default=1000
         Maximum Sinkhorn-like iterations.
    - stopThr: float, default=1e-9
         Convergence threshold on marginal constraints.
    - verbose: bool, default=False
         Print iteration information.
    - log: bool, default=False
         Return optimization log.
    
    Returns:
    - transport_plan: ndarray, shape (n_samples_source, n_samples_target)
         Entropic partial transport plan.
    - log: dict (if log=True)
         Contains 'err': convergence errors, 'mass': transported mass,
         'u': source scaling factors, 'v': target scaling factors.
    
    Example:
        # Entropic partial transport with regularization
        reg = 0.1  # Regularization strength
        plan = ot.partial.entropic_partial_wasserstein(a, b, M, reg, m=0.8)
    """

def ot.partial.partial_gromov_wasserstein(C1, C2, p, q, m=None, loss_fun='square_loss', alpha=0.5, armijo=False, log=False, max_iter=1000, tol_rel=1e-9, tol_abs=1e-9, **kwargs):
    """
    Solve partial Gromov-Wasserstein problem.
    
    Extends partial transport to structured data by combining partial transport
    with Gromov-Wasserstein distance for comparing metric spaces.
    
    Parameters:
    - C1: array-like, shape (n1, n1)
         Intra-structure cost matrix for source space.
    - C2: array-like, shape (n2, n2)
         Intra-structure cost matrix for target space.
    - p: array-like, shape (n1,)
         Distribution over source space.
    - q: array-like, shape (n2,)
         Distribution over target space.
    - m: float, optional
         Mass to transport. If None, uses min(sum(p), sum(q)).
    - loss_fun: str or callable, default='square_loss'
         Loss function for structure comparison.
    - alpha: float, default=0.5
         Step size parameter for optimization.
    - armijo: bool, default=False
         Use Armijo line search for step size adaptation.
    - log: bool, default=False
    - max_iter: int, default=1000
    - tol_rel: float, default=1e-9
         Relative tolerance for convergence.
    - tol_abs: float, default=1e-9
         Absolute tolerance for convergence.
    - kwargs: dict
    
    Returns:
    - transport_plan: ndarray, shape (n1, n2)
         Partial Gromov-Wasserstein transport plan.
    - log: dict (if log=True)
    """

def ot.partial.partial_gromov_wasserstein2(C1, C2, p, q, m=None, loss_fun='square_loss', alpha=0.5, armijo=False, log=False, max_iter=1000, tol_rel=1e-9, tol_abs=1e-9, **kwargs):
    """
    Solve partial Gromov-Wasserstein and return cost only.
    
    Parameters: Same as partial_gromov_wasserstein()
    
    Returns:
    - cost: float
         Partial Gromov-Wasserstein cost.
    - log: dict (if log=True)
    """

def ot.partial.entropic_partial_gromov_wasserstein(C1, C2, p, q, reg, m=None, loss_fun='square_loss', G0=None, max_iter=1000, tol=1e-9, verbose=False, log=False):
    """
    Solve entropic regularized partial Gromov-Wasserstein.
    
    Combines entropic regularization with partial Gromov-Wasserstein for
    structured partial transport with computational advantages.
    
    Parameters:
    - C1: array-like, shape (n1, n1)
         Source structure matrix.
    - C2: array-like, shape (n2, n2)
         Target structure matrix.
    - p: array-like, shape (n1,)
         Source distribution.
    - q: array-like, shape (n2,)
         Target distribution.
    - reg: float
         Entropic regularization parameter.
    - m: float, optional
         Mass to transport.
    - loss_fun: str or callable, default='square_loss'
    - G0: array-like, optional
         Initial transport plan.
    - max_iter: int, default=1000
    - tol: float, default=1e-9
    - verbose: bool, default=False
    - log: bool, default=False
    
    Returns:
    - transport_plan: ndarray
    - log: dict (if log=True)
    """

def ot.partial.entropic_partial_gromov_wasserstein2(C1, C2, p, q, reg, m=None, loss_fun='square_loss', G0=None, max_iter=1000, tol=1e-9, verbose=False, log=False):
    """
    Entropic partial Gromov-Wasserstein cost only.
    
    Parameters: Same as entropic_partial_gromov_wasserstein()
    
    Returns:
    - cost: float
    - log: dict (if log=True)
    """

def ot.partial.gwgrad_partial(C1, C2, T):
    """
    Compute gradient for partial Gromov-Wasserstein problem.
    
    Specialized gradient computation for the partial GW objective function,
    accounting for the relaxed mass constraints.
    
    Parameters:
    - C1: array-like, shape (n1, n1)
         Source structure matrix.
    - C2: array-like, shape (n2, n2)
         Target structure matrix.
    - T: array-like, shape (n1, n2)
         Current transport plan.
    
    Returns:
    - gradient: ndarray, shape (n1, n2)
         Gradient of partial GW objective with respect to transport plan.
    """

def ot.partial.gwloss_partial(C1, C2, T):
    """
    Compute partial Gromov-Wasserstein loss value.
    
    Evaluates the objective function value for the partial GW problem
    given current transport plan.
    
    Parameters:
    - C1: array-like, shape (n1, n1)
         Source structure matrix.
    - C2: array-like, shape (n2, n2)
         Target structure matrix.
    - T: array-like, shape (n1, n2)
         Transport plan.
    
    Returns:
    - loss: float
         Partial GW loss value.
    """

import ot
import numpy as np

# Create distributions with different masses
np.random.seed(42)
a = np.array([0.6, 0.4])  # Source distribution
b = np.array([0.3, 0.3, 0.2])  # Target has less total mass

# Cost matrix
M = np.array([[1.0, 2.0, 3.0],
              [2.0, 1.0, 4.0]])

# Standard transport would require normalizing distributions
# Partial transport can handle different masses directly

# Transport only 70% of available mass
m = 0.7

# Compute partial transport
plan_partial = ot.partial.partial_wasserstein(a, b, M, m=m, log=True)
cost_partial = ot.partial.partial_wasserstein2(a, b, M, m=m)

print(f"Partial transport plan:\n{plan_partial[0]}")
print(f"Transported mass: {np.sum(plan_partial[0]):.3f}")
print(f"Partial transport cost: {cost_partial:.3f}")
print(f"Optimization info: {plan_partial[1]}")

# Compare partial vs full transport
masses_to_test = [0.5, 0.7, 0.9, min(np.sum(a), np.sum(b))]

print("Mass fraction -> Partial cost")
for m in masses_to_test:
    cost = ot.partial.partial_wasserstein2(a, b, M, m=m)
    print(f"{m:.1f} -> {cost:.4f}")

# Full transport for comparison (after normalization)
a_norm = a / np.sum(a)
b_norm = b / np.sum(b)
cost_full = ot.emd2(a_norm, b_norm, M)
print(f"Full transport (normalized): {cost_full:.4f}")

# Add entropic regularization to partial transport
reg = 0.1  # Regularization parameter
m = 0.8   # Partial mass

# Entropic partial transport
plan_entropic = ot.partial.entropic_partial_wasserstein(
    a, b, M, reg, m=m, verbose=True, log=True
)

print(f"Entropic partial plan:\n{plan_entropic[0]}")
print(f"Converged in {len(plan_entropic[1]['err'])} iterations")
print(f"Final error: {plan_entropic[1]['err'][-1]:.2e}")

# Demonstrate outlier robustness
# Create distributions where one sample is an outlier

n_source, n_target = 50, 60
X_source = np.random.randn(n_source, 2)
X_target = np.random.randn(n_target, 2) + [1, 1]

# Add outlier to source
X_source[-1] = [5, 5]  # Far outlier

# Create cost matrix
M_outlier = ot.dist(X_source, X_target)

# Uniform distributions
a_unif = ot.unif(n_source)
b_unif = ot.unif(n_target)

# Full transport (must transport outlier)
plan_full = ot.emd(a_unif, b_unif, M_outlier)
cost_full = ot.emd2(a_unif, b_unif, M_outlier)

# Partial transport (can ignore outlier)
m_partial = 0.9  # Transport 90% of mass
plan_partial = ot.partial.partial_wasserstein(a_unif, b_unif, M_outlier, m=m_partial)
cost_partial = ot.partial.partial_wasserstein2(a_unif, b_unif, M_outlier, m=m_partial)

print(f"Full transport cost: {cost_full:.3f}")
print(f"Partial transport cost: {cost_partial:.3f}")
print(f"Cost reduction: {(cost_full - cost_partial) / cost_full * 100:.1f}%")

# Check if outlier is transported
outlier_transported = np.sum(plan_partial[0][-1, :])  # Last source sample
print(f"Outlier transported mass: {outlier_transported:.3f}")

# Structured partial transport example
n1, n2 = 20, 25

# Create structure matrices (e.g., distance matrices)
X1 = np.random.randn(n1, 3)
X2 = np.random.randn(n2, 3)
C1 = ot.dist(X1)
C2 = ot.dist(X2)

# Distributions
p = ot.unif(n1)
q = ot.unif(n2)

# Partial mass
m_gw = 0.8

# Partial Gromov-Wasserstein
plan_pgw = ot.partial.partial_gromov_wasserstein(
    C1, C2, p, q, m=m_gw, verbose=True, log=True
)
cost_pgw = ot.partial.partial_gromov_wasserstein2(C1, C2, p, q, m=m_gw)

print(f"Partial GW cost: {cost_pgw:.4f}")
print(f"PGW transported mass: {np.sum(plan_pgw[0]):.3f}")

# Entropic version
reg_gw = 0.05
plan_epgw = ot.partial.entropic_partial_gromov_wasserstein(
    C1, C2, p, q, reg_gw, m=m_gw, verbose=True
)

print(f"Entropic PGW completed")

# Explore different mass levels systematically
mass_fractions = np.linspace(0.1, 1.0, 10)
costs = []

for frac in mass_fractions:
    m = frac * min(np.sum(a), np.sum(b))
    cost = ot.partial.partial_wasserstein2(a, b, M, m=m)
    costs.append(cost)

print("Mass fraction -> Cost")
for frac, cost in zip(mass_fractions, costs):
    print(f"{frac:.1f} -> {cost:.4f}")

# Find elbow point (good trade-off between mass and cost)
cost_gradients = np.diff(costs)
optimal_idx = np.argmin(cost_gradients) + 1
optimal_mass = mass_fractions[optimal_idx] * min(np.sum(a), np.sum(b))

print(f"Suggested optimal mass: {optimal_mass:.3f}")

# Test robustness to noise
noise_levels = [0.0, 0.1, 0.2, 0.5]
costs_full = []
costs_partial = []

for noise in noise_levels:
    # Add noise to cost matrix
    M_noisy = M + noise * np.random.randn(*M.shape)
    
    # Full transport
    cost_full_noisy = ot.emd2(a/np.sum(a), b/np.sum(b), M_noisy)
    costs_full.append(cost_full_noisy)
    
    # Partial transport (70% mass)
    cost_partial_noisy = ot.partial.partial_wasserstein2(a, b, M_noisy, m=0.7)
    costs_partial.append(cost_partial_noisy)

print("Noise level -> Full cost | Partial cost")
for noise, cost_f, cost_p in zip(noise_levels, costs_full, costs_partial):
    print(f"{noise:.1f} -> {cost_f:.4f} | {cost_p:.4f}")

# Compare computational efficiency for larger problems
sizes = [50, 100, 200]

for n in sizes:
    # Generate larger problem
    X_large_s = np.random.randn(n, 5)
    X_large_t = np.random.randn(n, 5)
    M_large = ot.dist(X_large_s, X_large_t)
    a_large = ot.unif(n)
    b_large = ot.unif(n)
    
    # Time partial transport
    import time
    
    tic = time.time()
    cost_partial_large = ot.partial.partial_wasserstein2(
        a_large, b_large, M_large, m=0.8
    )
    time_partial = time.time() - tic
    
    # Time entropic partial (faster for large problems)
    tic = time.time()
    plan_entropic_large = ot.partial.entropic_partial_wasserstein(
        a_large, b_large, M_large, reg=0.1, m=0.8
    )
    time_entropic = time.time() - tic
    
    print(f"Size {n}x{n}: Partial={time_partial:.3f}s, Entropic={time_entropic:.3f}s")