pyotc.otc_backend.policy_iteration.sparse package

Submodules

pyotc.otc_backend.policy_iteration.sparse.exact module

pyotc.otc_backend.policy_iteration.sparse.exact.exact_otc(Px, Py, c, stat_dist='best', max_iter=100)[source]

Computes the optimal transport coupling (OTC) between two stationary Markov chains represented by transition matrices Px and Py, as described in Algorithm 1 of the paper: “Optimal Transport for Stationary Markov Chains via Policy Iteration” (https://www.jmlr.org/papers/volume23/21-0519/21-0519.pdf).

The algorithm iteratively updates the transition coupling matrix until convergence by alternating between Transition Coupling Evaluation (TCE) and Transition Coupling Improvement (TCI) steps.

For a detailed discussion of the connection between the OTC problem and Markov Decision Processes (MDPs), see Section 4 of the paper. Additional background on policy iteration methods for solving average-cost MDP problems can be found in Chapters 8 and 9 of “Markov Decision Processes: Discrete Stochastic Dynamic Programming” by Martin L. Puterman.

Note

In the TCE step (implemented in exact_tce), we solve a block linear system using functions from scipy.sparse.linalg. However, when A in Ax = b is nearly singular, we have observed a few cases where both SciPy solvers (scipy.sparse.linalg.spsolve, scipy.sparse.linalg.lsmr) can produce results that differ from NumPy’s solver (np.linalg.solve). This leads to discrepancies with the dense implementation and non-convergence. This is an issue with SciPy’s sparse solvers and remains unresolved. The best approach in such cases is to use the dense implementation.

Parameters:

  • Px (np.ndarray) – Transition matrix of the source Markov chain of shape (dx, dx).

  • Py (np.ndarray) – Transition matrix of the target Markov chain of shape (dy, dy).

  • c (np.ndarray) – Cost function of shape (dx, dy).

  • stat_dist (str, optional) – Method to compute the stationary distribution. Options include ‘best’, ‘eigen’, ‘iterative’ and None. Defaults to ‘best’.

  • max_iter (int, optional) – Maximum number of iterations for the convergence process. Defaults to 100.

Returns:

Expected transport cost under the optimal transition coupling. R (scipy.sparse.csr_matrix): Optimal transition coupling matrix of shape (dx*dy, dx*dy). stat_dist (np.ndarray): Stationary distribution of the optimal transition coupling of shape (dx, dy).

If convergence is not reached within max_iter iterations, returns (None, None, None).

Return type:

exp_cost (float)

pyotc.otc_backend.policy_iteration.sparse.exact_tce module

Original Transition Coupling Evaluation (TCE) method from: https://www.jmlr.org/papers/volume23/21-0519/21-0519.pdf

pyotc.otc_backend.policy_iteration.sparse.exact_tce.exact_tce(R_sparse, c)[source]

Computes the exact Transition Coupling Evaluation (TCE) vectors g and h for a given sparse transition matrix R_sparse and cost vector c.

Specifically, solves the block linear system outlined in Algorithm 1a of the paper: “Optimal Transport for Stationary Markov Chains via Policy Iteration” (https://www.jmlr.org/papers/volume23/21-0519/21-0519.pdf).

Due to memory constraints associated with direct solvers (e.g., sp.linalg.spsolve), an iterative solver (scipy.sparse.linalg.lsmr) is employed to efficiently handle large-scale sparse systems.

Notes

1. When A in Ax = b is close to singular, we have observed few cases that both SciPy functions (scipy.sparse.linalg.spsolve, scipy.sparse.linalg.lsmr) can produce results that differ from NumPy’s solver, leading to different results with dense implementation and non-convergence. This is an issue with SciPy solvers and remains an unresolved issue. The best approach in such cases is to fall back to the dense implementation.

2. Solving Ax = b using a direct solver (scipy.sparse.linalg.spsolve) on large networks resulted in: “Not enough memory to perform factorization.” This is likely due to excessive fill-in during LU factorization of the large sparse matrix. To address this, we switch to an iterative solver (scipy.sparse.linalg.lsmr), which is more memory-efficient and better suited for large-scale sparse systems.

  1. To leave open the possibility of switching from ‘lsmr’ to ‘spsolve’, the corresponding ‘spsolve’ code has been retained as a commented-out block.

Parameters:

  • R_sparse (scipy.sparse.csr_matrix) – Sparse transition matrix of shape (dx*dy, dx*dy).

  • c (np.ndarray) – Cost vector of shape (dx, dy).

Returns:

Average cost (gain) vector of shape (dx*dy,). h (np.ndarray): Total extra cost (bias) vector of shape (dx*dy,).

Return type:

g (np.ndarray)

pyotc.otc_backend.policy_iteration.sparse.exact_tci module

Original Transition Coupling Improvements (TCI) method from: https://www.jmlr.org/papers/volume23/21-0519/21-0519.pdf

pyotc.otc_backend.policy_iteration.sparse.exact_tci.exact_tci(g, h, R0, Px, Py)[source]

Performs the Transition Coupling Improvement (TCI) step in the OTC algorithm.

This function attempts to update the current coupling transition matrix R0 based on the evaluation vectors g and h obtained from the Transition Coupling Evaluation (TCE).

Parameters:

  • g (np.ndarray) – Gain vector from TCE of shape (dx*dy,).

  • h (np.ndarray) – Bias vector from TCE of shape (dx*dy,).

  • R0 (scipy.sparse.csr_matrix or lil_matrix) – Current transition coupling matrix of shape (dx*dy, dx*dy).

  • Px (np.ndarray) – Transition matrix of the source Markov chain of shape (dx, dx).

  • Py (np.ndarray) – Transition matrix of the target Markov chain of shape (dy, dy).

Returns:

Improved transition coupling matrix of shape (dx*dy, dx*dy).

Return type:

R (scipy.sparse.lil_matrix)

pyotc.otc_backend.policy_iteration.sparse.exact_tci.setup_ot(f, Px, Py, R)[source]

This improvement step updates the transition coupling matrix R that minimizes the product Rf element-wise. In more detail, we may select a transition coupling R such that for each state pair (x, y), the corresponding row r = R((x, y), ·) minimizes rf over couplings r in Pi(Px(x, ·), Py(y, ·)). This is done by solving the optimal transport problem for each state pair (x, y) in the source and target Markov chains. The resulting transition coupling matrix R is updated accordingly.

This function uses the POT (Python Optimal Transport) library to solve the optimal transport problem for each (x, y) state pair and updates the transition coupling matrix.

Parameters:

  • f (np.ndarray) – Cost function reshaped as of shape (dx*dy,).

  • Px (np.ndarray) – Transition matrix of the source Markov chain of shape (dx, dx).

  • Py (np.ndarray) – Transition matrix of the target Markov chain of shape (dy, dy).

  • R (np.ndarray) – Transition coupling matrix to update of shape (dx*dy, dx*dy).

Returns:

Updated transition coupling matrix of shape (dx*dy, dx*dy).

Return type:

R (np.ndarray)

Module contents