pyotc.otc_backend.optimal_transport package

Submodules

pyotc.otc_backend.optimal_transport.logsinkhorn module

pyotc.otc_backend.optimal_transport.logsinkhorn.logsinkhorn(A, r, c, T)[source]

Implementation of classical Sinkhorn algorithm for matrix scaling. Each iteration simply alternately updates (projects) all rows or all columns to have correct marginals.

Parameters:
  • A (np.ndarray) – Negative scaled cost matrix of shape (dx, dy), e.g., -xi * cost.

  • r (np.ndarray) – desired row sums (marginals) (shape: dx,). Should sum to 1.

  • c (np.ndarray) – desired column sums (marginals) (shape: dy,). Should sum to 1.

  • T (int) – Number of full Sinkhorn iterations.

Returns:

Final scaled matrix of shape (dx, dy).

Return type:

np.ndarray

pyotc.otc_backend.optimal_transport.logsinkhorn.logsumexp(X, axis=None)[source]

Numerically stable log-sum-exp operation.

Parameters:
  • X (np.ndarray) – Input array.

  • axis (int or tuple of ints, optional) – Axis or axes over which to operate.

Returns:

The result of log(sum(exp(X))) along the specified axis.

Return type:

np.ndarray

pyotc.otc_backend.optimal_transport.logsinkhorn.round_transpoly(X, r, c)[source]

pyotc.otc_backend.optimal_transport.native module

Native linear programming implementation for solving optimal transport (OT) problems.

pyotc.otc_backend.optimal_transport.native.computeot_lp(C, r, c)[source]

Solves the optimal transport problem using linear programming (LP) with SciPy.

Given a cost matrix C and distributions r and c, this function computes the optimal transport plan that minimizes the total transport cost.

Parameters:
  • C (np.ndarray) – Cost matrix of shape (nx, ny), where C[i, j] represents the cost of transporting mass from source i to target j.

  • r (np.ndarray) – Source distribution (shape: nx,). Should sum to 1.

  • c (np.ndarray) – Target distribution (shape: ny,). Should sum to 1.

Returns:

  • lp_sol (np.ndarray): Optimal transport plan of shape (nx, ny).

  • lp_val (float): Total transport cost under the optimal plan.

Return type:

Tuple[np.ndarray, float]

pyotc.otc_backend.optimal_transport.native_refactor module

Yuning’s other other native implementation of lp ot

pyotc.otc_backend.optimal_transport.native_refactor.computeot_lp(C: ndarray, r: ndarray, c: ndarray) tuple[Any, Any][source]

Compute optimal transport mapping via LP.

Parameters:
  • C (np.ndarray) – cost

  • r (np.ndarray) – _description_

  • c (np.ndarray) – _description_

Returns:

_description_

Return type:

tuple[Any, Any]

pyotc.otc_backend.optimal_transport.native_refactor.setup_columns(Aeq: ndarray, nx: int, ny: int) None[source]
pyotc.otc_backend.optimal_transport.native_refactor.setup_rows(Aeq: ndarray, nx: int, ny: int) None[source]

pyotc.otc_backend.optimal_transport.pot module

A wrapper for the Python Optimal Transport (POT) library.

This module provides a simplified interface for computing optimal transport plans and their associated costs using the POT library.

pyotc.otc_backend.optimal_transport.pot.computeot_pot(C, r, c)[source]

Computes the optimal transport plan and its total cost using the POT library.

Given a cost matrix C and distributions r and c, this function computes the optimal transport plan that minimizes the total transport cost.

Parameters:
  • C (np.ndarray) – Cost matrix of shape (n, m), where C[i, j] represents the cost of transporting mass from source i to target j.

  • r (np.ndarray) – Source distribution (shape: n,). Should sum to 1.

  • c (np.ndarray) – Target distribution (shape: m,). Should sum to 1.

Returns:

  • lp_sol (np.ndarray): Optimal transport plan of shape (n, m).

  • lp_val (float): Total transport cost under the optimal plan.

Return type:

Tuple[np.ndarray, float]

Module contents