pyotc 0.2.2__py3-none-any.whl

pyotc/otc_backend/policy_iteration/dense/approx_tce.py

@@ -0,0 +1,42 @@
+import numpy as np
+
+
+def approx_tce(P, c, L, T):
+    """
+    Approximates the Transition Coupling Evaluation (TCE) vectors g and h
+    using a truncation-based approximation of the exact TCE method.
+
+    Args:
+        P (np.ndarray): Transition matrix of shape (dx*dy, dx*dy).
+        c (np.ndarray): Cost vector of shape (dx*dy,) or (dx*dy, 1).
+        L (int): Maximum number of iterations for computing the cost vector g.
+        T (int): Maximum number of iterations for computing the bias vector h.
+
+    Returns:
+        g (np.ndarray): Approximated average cost (gain) vector of shape (dx*dy,).
+        h (np.ndarray): Approximated bias vector of shape (dx*dy,).
+    """
+
+    d = P.shape[0]
+    c = np.reshape(c, (d, -1))
+    c_max = np.max(c)
+
+    g_old = c
+    g = P @ g_old
+    l = 1
+    tol = 1e-12
+    while l <= L and np.max(np.abs(g - g_old)) > tol * c_max:
+        g_old = g
+        g = P @ g_old
+        l += 1
+
+    g = np.mean(g) * np.ones((d, 1))
+    diff = c - g
+    h = diff.copy()
+    t = 1
+    while t <= T and np.max(np.abs(P @ diff)) > tol * c_max:
+        h += P @ diff
+        diff = P @ diff
+        t += 1
+
+    return g, h

pyotc/otc_backend/policy_iteration/dense/entropic.py

@@ -0,0 +1,161 @@
+"""
+Entropic Optimal Transition Coupling (OTC) solvers.
+
+Implements variants of the OTC algorithm using entropic regularization.
+Includes both a custom Sinkhorn implementation and one based on the POT library.
+
+References:
+    - Section 5, "Optimal Transport for Stationary Markov Chains via Policy Iteration"
+      (https://www.jmlr.org/papers/volume23/21-0519/21-0519.pdf)
+
+Methods:
+    - logsinkhorn: A self-implemented log-scaled Sinkhorn solver.
+    - ot_sinkhorn: Sinkhorn solver from POT library.
+    (reference: https://pythonot.github.io/gen_modules/ot.bregman.html#ot.bregman.sinkhorn)
+    - ot_logsinkhorn: Sinkhorn solver from POT library in log scale.
+    (reference: https://pythonot.github.io/gen_modules/ot.bregman.html#ot.bregman.sinkhorn_log)
+    - ot_greenkhorn: Sinkhorn solver of greedy version from POT library.
+    (reference: https://pythonot.github.io/gen_modules/ot.bregman.html#ot.bregman.greenkhorn)
+"""
+
+import time
+import numpy as np
+import ot
+from ..utils import get_best_stat_dist
+from .approx_tce import approx_tce
+from .entropic_tci import entropic_tci
+from pyotc.otc_backend.optimal_transport.logsinkhorn import logsinkhorn
+
+
+def entropic_otc(
+    Px,
+    Py,
+    c,
+    L=100,
+    T=100,
+    xi=0.1,
+    method="logsinkhorn",
+    sink_iter=100,
+    reg_num=None,
+    get_sd=False,
+    silent=True,
+):
+    """
+    Solves the Entropic Optimal Transition Coupling (OTC) problem between two Markov chains
+    using approximate policy iteration and entropic regularization.
+
+    This method alternates between approximate coupling evaluation
+    and entropic coupling improvement (via Sinkhorn iterations), until convergence.
+
+    Args:
+        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).
+        L (int): Number of iterations for computing the cost vector g in approx_tce.
+        T (int): Number of iterations for computing the bias vector h in approx_tce.
+        xi (float): Scaling factor for entropic cost adjustment in entropic_tci.
+        method (str): Method for the Sinkhorn algorithm. Must choose from ['logsinkhorn', 'ot_sinkhorn', 'ot_logsinkhorn', 'ot_greenkhorn']. Default is 'logsinkhorn'. See 'Methods' above for details.
+        sink_iter (int): Number of iterations for 'logsinkhorn' method. Maximum number of Sinkhorn iterations for other methods from POT library. Used in the entropic TCI step.
+        reg_num (float): Entropic regularization term, used only for methods from POT package.
+        get_sd (bool): If True, compute best stationary distribution using linear programming.
+        silent (bool): If False, print convergence info during iterations and running time
+
+    Returns:
+        exp_cost (float): Expected transport cost under the optimal transition coupling.
+        P (np.ndarray): Optimal transition coupling matrix of shape (dx*dy, dx*dy).
+        stat_dist (Optional[np.ndarray]): Stationary distribution of the optimal transition coupling of shape (dx, dy),
+                                            or None if get_sd is False.
+    """
+    if not silent:
+        start_time = time.time()
+        print(f"Starting entropic otc with {method} method...")
+
+    dx, dy = Px.shape[0], Py.shape[0]
+    max_c = np.max(c)
+    tol = 1e-5 * max_c
+
+    g_old = max_c * np.ones(dx * dy)
+    g = g_old - 10 * tol
+    P = np.kron(Px, Py)
+
+    if method == "logsinkhorn":
+
+        def solver_fn(A, a, b):
+            return logsinkhorn(A, a, b, sink_iter)
+    elif method == "ot_sinkhorn":
+        if reg_num is None:
+            raise ValueError("reg_num must be specified for 'ot_sinkhorn'")
+
+        def solver_fn(A, a, b):
+            return ot.sinkhorn(a, b, A, reg=reg_num, numItermax=sink_iter)
+    elif method == "ot_logsinkhorn":
+        if reg_num is None:
+            raise ValueError("reg_num must be specified for 'ot_logsinkhorn'")
+
+        def solver_fn(A, a, b):
+            return ot.bregman.sinkhorn_log(a, b, A, reg=reg_num, numItermax=sink_iter)
+    elif method == "ot_greenkhorn":
+        if reg_num is None:
+            raise ValueError("reg_num must be specified for 'ot_greenkhorn'")
+
+        def solver_fn(A, a, b):
+            return ot.bregman.greenkhorn(a, b, A, reg=reg_num, numItermax=sink_iter)
+    else:
+        raise ValueError(f"Unknown method: {method}")
+
+    iter_ctr = 0
+    while g_old[0] - g[0] > tol:
+        iter_ctr += 1
+        P_old = P
+        g_old = g
+        if not silent:
+            print("Iteration:", iter_ctr)
+            start_iter = time.time()
+
+        # Approximate transition coupling evaluation
+        if not silent:
+            print("Computing entropic TCE...")
+        g, h = approx_tce(P, c, L, T)
+
+        # Entropic transition coupling improvement (passing solver function to entropic_tci)
+        if not silent:
+            print("Computing entropic TCE...")
+        P = entropic_tci(h=h, P0=P_old, Px=Px, Py=Py, xi=xi, solver_fn=solver_fn)
+
+        if not silent:
+            iter_time = time.time() - start_iter
+            elapsed = time.time() - start_time
+            g0 = float(np.ravel(g)[0])
+            g0_old = float(np.ravel(g_old)[0])
+            diff = g0_old - g0
+            ratio = diff / g0 if g0 != 0 else float("inf")
+            print(
+                f"[Iter {iter_ctr} taking {iter_time:.2f}s] Δg={diff:.3e}, g[0]={g0:.6f}, Δg/g[0]={ratio:.3e}, total elapsed={elapsed:.2f}s"
+            )
+
+    # In case of numerical instability, make non-negative and normalize.
+    P = np.maximum(P, 0)
+    row_sums = np.sum(P, axis=1, keepdims=True)
+    P = P / np.where(row_sums > 0, row_sums, 1)
+
+    if get_sd:
+        if not silent:
+            print(
+                f"Convergence reached in {iter_ctr} iterations. Computing stationary distribution..."
+            )
+        stat_dist, exp_cost = get_best_stat_dist(P, c)
+        stat_dist = np.reshape(stat_dist, (dx, dy))
+    else:
+        if not silent:
+            print(
+                f"Convergence reached in {iter_ctr} iterations. No stationary distribution computation requested."
+            )
+        stat_dist = None
+        exp_cost = g[0].item()
+
+    if not silent:
+        print(
+            f"[entropic_otc] Finished. Total time elapsed: {time.time() - start_time:.3f} seconds."
+        )
+
+    return exp_cost, P, stat_dist

pyotc/otc_backend/policy_iteration/dense/entropic_tci.py

@@ -0,0 +1,49 @@
+import numpy as np
+import ot
+
+
+def entropic_tci(h, P0, Px, Py, xi, solver_fn):
+    """
+    Performs entropic Transition Coupling Improvement (TCI) using log-domain Sinkhorn algorithm.
+
+    For each (i, j) state pair from the product space of two Markov chains, this function solves
+    a local entropic optimal transport problem based on the bias vector h.
+
+    Args:
+        h (np.ndarray): Bias vector of shape (dx*dy,).
+        P0 (np.ndarray): Previous 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).
+        xi (float): Scaling factor for entropic cost adjustment.
+        solver_fn (callable): A function solves the optimization and provides a transport plan. Specified in 'entropic_otc'.
+
+    Returns:
+        np.ndarray: Updated transition coupling matrix of shape (dx*dy, dx*dy).
+    """
+
+    dx, dy = Px.shape[0], Py.shape[0]
+    P = P0.copy()
+    h_mat = np.reshape(h, (dx, dy))
+    K = -xi * h_mat
+
+    for i in range(dx):
+        for j in range(dy):
+            dist_x = Px[i, :]
+            dist_y = Py[j, :]
+            x_idxs = np.where(dist_x > 0)[0]
+            y_idxs = np.where(dist_y > 0)[0]
+
+            if len(x_idxs) == 1 or len(y_idxs) == 1:
+                P[dy * i + j, :] = P0[dy * i + j, :]
+            else:
+                A_matrix = K[np.ix_(x_idxs, y_idxs)]
+                sub_dist_x = dist_x[x_idxs]
+                sub_dist_y = dist_y[y_idxs]
+
+                sol = solver_fn(A_matrix, sub_dist_x, sub_dist_y)
+
+                sol_full = np.zeros((dx, dy))
+                sol_full[np.ix_(x_idxs, y_idxs)] = sol
+                P[dy * i + j, :] = sol_full.flatten()
+
+    return P

pyotc/otc_backend/policy_iteration/dense/exact.py

@@ -0,0 +1,127 @@
+import numpy as np
+import time
+
+from .exact_tce import exact_tce
+from .exact_tci_lp import exact_tci as exact_tci_lp
+from .exact_tci_pot import exact_tci as exact_tci_pot
+from ..utils import get_stat_dist
+
+
+def exact_otc_lp(Px, Py, c, stat_dist="best"):
+    start = time.time()
+    print("Starting exact_otc_dense...")
+
+    dx = Px.shape[0]
+    dy = Py.shape[0]
+    P_old = np.ones((dx * dy, dx * dy))
+    P = np.kron(Px, Py)
+
+    while np.max(np.abs(P - P_old)) > 1e-10:
+        P_old = np.copy(P)
+
+        print("Computing exact TCE...")
+        g, h = exact_tce(P, c)
+
+        print("Computing exact TCI...")
+        P = exact_tci_lp(g, h, P_old, Px, Py)
+
+        # Check for convergence.
+        if np.all(P == P_old):
+            if stat_dist is None:
+                print(
+                    "Convergence reached. No stationary distribution computation requested."
+                )
+                exp_cost = g[0].item()
+                end = time.time()
+                print(
+                    f"[exact_otc] Finished. Total time elapsed: {end - start:.3f} seconds."
+                )
+                return float(exp_cost), P, None
+            else:
+                print("Convergence reached. Computing stationary distribution...")
+                stat_dist = get_stat_dist(P, method=stat_dist, c=c)
+                stat_dist = np.reshape(stat_dist, (dx, dy))
+                exp_cost = g[0].item()
+                end = time.time()
+                print(
+                    f"[exact_otc] Finished. Total time elapsed: {end - start:.3f} seconds."
+                )
+                return float(exp_cost), P, stat_dist
+
+    return None, None, None
+
+
+def exact_otc(Px, Py, c, stat_dist="best"):
+    """
+    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.
+
+    Args:
+        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'.
+
+    Returns:
+        exp_cost (float): Expected transport cost under the optimal transition coupling.
+        R (np.ndarray): 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).
+
+        Returns (None, None, None) if the algorithm fails to converge.
+    """
+
+    start = time.time()
+    print("Starting exact_otc_dense...")
+
+    dx, dy = Px.shape[0], Py.shape[0]
+    R_old = np.ones((dx * dy, dx * dy))
+    R = np.kron(Px, Py)
+    iter = 0
+
+    while np.max(np.abs(R - R_old)) > 1e-10:
+        print("Iteration:", iter)
+        R_old = np.copy(R)
+
+        print("Computing exact TCE...")
+        g, h = exact_tce(R, c)
+
+        print("Computing exact TCI...")
+        R = exact_tci_pot(g, h, R_old, Px, Py)
+
+        # Check if the transition coupling matrix has converged
+        if np.all(R == R_old):
+            if stat_dist is None:
+                print(
+                    f"Convergence reached in {iter + 1} iterations. No stationary distribution computation requested."
+                )
+                exp_cost = g[0].item()
+                end = time.time()
+                print(
+                    f"[exact_otc] Finished. Total time elapsed: {end - start:.3f} seconds."
+                )
+                return float(exp_cost), R, None
+            else:
+                print(
+                    f"Convergence reached in {iter + 1} iterations. Computing stationary distribution..."
+                )
+                stat_dist = get_stat_dist(R, method=stat_dist, c=c)
+                stat_dist = np.reshape(stat_dist, (dx, dy))
+                exp_cost = g[0].item()
+                end = time.time()
+                print(
+                    f"[exact_otc] Finished. Total time elapsed: {end - start:.3f} seconds."
+                )
+                return float(exp_cost), R, stat_dist
+
+        iter += 1
+
+    return None, None, None

pyotc/otc_backend/policy_iteration/dense/exact_tce.py

@@ -0,0 +1,56 @@
+"""
+Original Transition Coupling Evaluation (TCE) methods from:
+https://www.jmlr.org/papers/volume23/21-0519/21-0519.pdf
+"""
+
+import numpy as np
+from numpy.linalg import pinv
+
+
+def exact_tce(R, c):
+    """
+    Computes the exact Transition Coupling Evaluation (TCE) vectors g and h
+    using the linear system described 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).
+
+    The method solves a block linear system involving the transition matrix R and cost vector c.
+    If the system is not full rank, a pseudo-inverse (pinv) is used as fallback.
+
+    Args:
+        R (np.ndarray): Transition matrix of shape (dx*dy, dx*dy).
+        c (np.ndarray): Cost vector of shape (dx*dy, dx*dy).
+
+    Returns:
+        g (np.ndarray): Average cost (gain) vector of shape (dx*dy,).
+        h (np.ndarray): Total extra cost (bias) vector of shape (dx*dy,).
+
+    Notes:
+        - If the matrix A is singular or ill-conditioned, the solution uses `np.linalg.pinv`,
+          which may lead to numerical instability.
+        - Make sure Pz is a proper stochastic matrix (rows sum to 1).
+    """
+    d = R.shape[0]
+    c = np.reshape(c, (d, -1))
+
+    # Construct the block matrix A and right-hand side vector b
+    A = np.block(
+        [
+            [np.eye(d) - R, np.zeros((d, d)), np.zeros((d, d))],
+            [np.eye(d), np.eye(d) - R, np.zeros((d, d))],
+            [np.zeros((d, d)), np.eye(d), np.eye(d) - R],
+        ]
+    )
+    b = np.concatenate([np.zeros((d, 1)), c, np.zeros((d, 1))])
+
+    # Solve the linear system Ax = b
+    try:
+        sol = np.linalg.solve(A, b)
+    except:
+        sol = np.matmul(pinv(A), b)
+
+    # Extract g and h from the solution
+    g = sol[0:d].flatten()
+    h = sol[d : 2 * d].flatten()
+
+    return g, h

pyotc/otc_backend/policy_iteration/dense/exact_tci_lp.py

@@ -0,0 +1,65 @@
+"""
+Original Transition Coupling Improvements (TCI) methods from:
+https://jmlr.csail.mit.edu/papers/volume23/21-0519/21-0519.pdf
+
+Use scipy.linprog (LP solver) library to solve optimal transport problem.
+"""
+
+import numpy as np
+import copy
+from pyotc.otc_backend.optimal_transport.native import computeot_lp
+
+
+def check_constant(f, Px, threshold=1e-3):
+    dx = Px.shape[0]
+    g_const = True
+    for i in range(dx):
+        for j in range(i + 1, dx):
+            if abs(f[i] - f[j]) > threshold:
+                g_const = False
+                break
+        if not g_const:
+            break
+    return g_const
+
+
+def setup_ot(f, Px, Py, Pz):
+    dx = Px.shape[0]
+    dy = Py.shape[0]
+    f_mat = np.reshape(f, (dx, dy))
+    for x_row in range(dx):
+        for y_row in range(dy):
+            dist_x = Px[x_row, :]
+            dist_y = Py[y_row, :]
+            # Check if either distribution is degenerate.
+            if any(dist_x == 1) or any(dist_y == 1):
+                sol = np.outer(dist_x, dist_y)
+            # If not degenerate, proceed with OT.
+            else:
+                sol, val = computeot_lp(f_mat, dist_x, dist_y)
+            idx = dy * (x_row) + y_row
+            Pz[idx, :] = np.reshape(sol, (-1, dx * dy))
+    return Pz
+
+
+def exact_tci(g, h, P0, Px, Py):
+    # Check if g is constant.
+    dx = Px.shape[0]
+    dy = Py.shape[0]
+    Pz = np.zeros((dx * dy, dx * dy))
+    g_const = check_constant(f=g, Px=Px)
+
+    # If g is not constant, improve transition coupling against g.
+    if not g_const:
+        Pz = setup_ot(f=g, Px=Px, Py=Py, Pz=Pz)
+        if np.max(np.abs(np.matmul(P0, g) - np.matmul(Pz, g))) <= 1e-7:
+            Pz = copy.deepcopy(P0)
+        else:
+            return Pz
+
+    # Try to improve with respect to h.
+    Pz = setup_ot(f=h, Px=Px, Py=Py, Pz=Pz)
+    if np.max(np.abs(np.matmul(P0, h) - np.matmul(Pz, h))) <= 1e-4:
+        Pz = copy.deepcopy(P0)
+
+    return Pz

pyotc/otc_backend/policy_iteration/dense/exact_tci_pot.py

@@ -0,0 +1,90 @@
+"""
+Original Transition Coupling Improvements (TCI) method from:
+https://www.jmlr.org/papers/volume23/21-0519/21-0519.pdf
+
+Use the python optimal transport (POT) library to solve optimal transport problem.
+"""
+
+import numpy as np
+import copy
+from pyotc.otc_backend.optimal_transport.pot import computeot_pot
+
+
+def setup_ot(f, Px, Py, R):
+    """
+    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.
+
+    Args:
+        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:
+        R (np.ndarray): Updated transition coupling matrix of shape (dx*dy, dx*dy).
+    """
+
+    dx, dy = Px.shape[0], Py.shape[0]
+    f_mat = np.reshape(f, (dx, dy))
+
+    for x_row in range(dx):
+        for y_row in range(dy):
+            dist_x = Px[x_row, :]
+            dist_y = Py[y_row, :]
+
+            # Check if either distribution is degenerate.
+            if any(dist_x == 1) or any(dist_y == 1):
+                sol = np.outer(dist_x, dist_y)
+            # If not degenerate, proceed with OT.
+            else:
+                sol, _ = computeot_pot(f_mat, dist_x, dist_y)
+            idx = dy * (x_row) + y_row
+            R[idx, :] = np.reshape(sol, (-1, dx * dy))
+
+    return R
+
+
+def exact_tci(g, h, R0, Px, Py):
+    """
+    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).
+
+    Args:
+        g (np.ndarray): Gain vector from TCE of shape (dx*dy,).
+        h (np.ndarray): Bias vector from TCE of shape (dx*dy,).
+        R0 (np.ndarray): 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:
+        R (np.ndarray): Improved transition coupling matrix of shape (dx*dy, dx*dy).
+    """
+
+    # Check if g is constant.
+    dx, dy = Px.shape[0], Py.shape[0]
+    R = np.zeros((dx * dy, dx * dy))
+    g_const = np.max(g) - np.min(g) <= 1e-3
+
+    # If g is not constant, improve transition coupling against g.
+    if not g_const:
+        R = setup_ot(g, Px, Py, R)
+        if np.max(np.abs(np.matmul(R0, g) - np.matmul(R, g))) <= 1e-7:
+            R = copy.deepcopy(R0)
+        else:
+            return R
+
+    # Try to improve with respect to h.
+    R = setup_ot(h, Px, Py, R)
+    if np.max(np.abs(np.matmul(R0, h) - np.matmul(R, h))) <= 1e-4:
+        R = copy.deepcopy(R0)
+
+    return R

pyotc/otc_backend/policy_iteration/sparse/exact.py

@@ -0,0 +1,89 @@
+import numpy as np
+import scipy.sparse as sp
+import time
+
+from .exact_tce import exact_tce
+from .exact_tci import exact_tci
+from ..utils import get_stat_dist
+
+
+def exact_otc(Px, Py, c, stat_dist="best", max_iter=100):
+    """
+    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.
+
+    Args:
+        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:
+        exp_cost (float): 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).
+    """
+
+    start = time.time()
+    print("Starting exact_otc_sparse...")
+    dx, dy = Px.shape[0], Py.shape[0]
+
+    # Initial coupling matrix using Kronecker product
+    R = sp.kron(sp.csr_matrix(Px), sp.csr_matrix(Py), format="csr")
+
+    for iter in range(max_iter):
+        print("Iteration:", iter)
+        R_old = R.copy()
+
+        print("Computing exact TCE...")
+        g, h = exact_tce(R, c)
+
+        print("Computing exact TCI...")
+        R = exact_tci(g, h, R_old, Px, Py)
+
+        # Check if the transition coupling matrix has converged
+        if (R != R_old).nnz == 0:
+            if stat_dist is None:
+                print(
+                    f"Convergence reached in {iter + 1} iterations. No stationary distribution computation requested."
+                )
+                exp_cost = g[0].item()
+                end = time.time()
+                print(
+                    f"[exact_otc] Finished. Total time elapsed: {end - start:.3f} seconds."
+                )
+                return float(exp_cost), R, None
+            else:
+                print(
+                    f"Convergence reached in {iter + 1} iterations. Computing stationary distribution..."
+                )
+                stat_dist = get_stat_dist(R, method=stat_dist, c=c)
+                stat_dist = np.reshape(stat_dist, (dx, dy))
+                exp_cost = g[0].item()
+                end = time.time()
+                print(
+                    f"[exact_otc] Finished. Total time elapsed: {end - start:.3f} seconds."
+                )
+                return float(exp_cost), R, stat_dist
+
+    # Return None if convergence is not achieved
+    print(f"Convergence not achieved after {iter} iterations. Returning None.")
+    return None, None, None