pyotc 0.2.2__py3-none-any.whl

pyotc/otc_backend/policy_iteration/sparse/exact_tce.py

@@ -0,0 +1,78 @@
+"""
+Original Transition Coupling Evaluation (TCE) method from:
+https://www.jmlr.org/papers/volume23/21-0519/21-0519.pdf
+"""
+
+import numpy as np
+import scipy.sparse as sp
+
+
+def exact_tce(R_sparse, c):
+    """
+    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.
+
+        3. To leave open the possibility of switching from 'lsmr' to 'spsolve', the corresponding 'spsolve' code has been retained as a commented-out block.
+
+    Args:
+        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:
+        g (np.ndarray): Average cost (gain) vector of shape (dx*dy,).
+        h (np.ndarray): Total extra cost (bias) vector of shape (dx*dy,).
+    """
+    n = R_sparse.shape[0]
+    c = c.reshape(n)
+
+    # Construct the block matrix A and right-hand side vector b
+    I = sp.eye(n, format="csr")
+    zero = sp.csr_matrix((n, n))
+    A = sp.bmat(
+        [[I - R_sparse, zero, zero], [I, I - R_sparse, zero], [zero, I, I - R_sparse]],
+        format="csr",
+    )
+    b = np.concatenate([np.zeros(n), c, np.zeros(n)])
+
+    # print("Solving sparse linear system in exact tce...")
+    # permc_specs = ['COLAMD', 'MMD_ATA', 'MMD_AT_PLUS_A', 'NATURAL']
+    # solution = None
+    # for spec in permc_specs:
+    #     try:
+    #         current_solution = sp.linalg.spsolve(A, rhs, permc_spec=spec)
+    #         if not np.any(np.abs(current_solution) > 1e15):
+    #             print("spsolve successful with spec:", spec)
+    #             solution = current_solution
+    #             break
+    #         else:
+    #             print(f"Solution with {spec} contains large values, trying next spec.")
+    #     except ValueError as e:
+    #         print(f"spsolve with {spec} encountered an error: trying next spec.")
+    # if solution is None:
+    #     raise RuntimeError("Failed to find a stable solution with any of the provided permc_specs for sp.linalg.spsolve solver.")
+
+    # Solve the linear system using an iterative solver (lsmr)
+    solution = sp.linalg.lsmr(A, b, atol=1e-10, btol=1e-10)[0]
+
+    # Extract vectors g and h from the solution
+    g = solution[:n]
+    h = solution[n : 2 * n]
+
+    return g, h

pyotc/otc_backend/policy_iteration/sparse/exact_tci.py

@@ -0,0 +1,88 @@
+"""
+Original Transition Coupling Improvements (TCI) method from:
+https://www.jmlr.org/papers/volume23/21-0519/21-0519.pdf
+"""
+
+import numpy as np
+import scipy.sparse as sp
+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 np.any(dist_x == 1) or np.any(dist_y == 1):
+                sol = np.outer(dist_x, dist_y)
+            else:
+                sol, _ = computeot_pot(f_mat, dist_x, dist_y)
+            idx = dy * x_row + y_row
+            sol_flat = sol.flatten()
+            for j in np.nonzero(sol_flat)[0]:
+                R[idx, j] = sol_flat[j]
+
+    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 (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:
+        R (scipy.sparse.lil_matrix): Improved transition coupling matrix of shape (dx*dy, dx*dy).
+    """
+
+    # Check if g is constant.
+    dx, dy = Px.shape[0], Py.shape[0]
+    R = sp.lil_matrix((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(R0.dot(g) - R.dot(g))) <= 1e-7:
+            R = R0.copy()
+        else:
+            return R
+
+    # Try to improve with respect to h.
+    R = setup_ot(h, Px, Py, R)
+    if np.max(np.abs(R0.dot(h) - R.dot(h))) <= 1e-4:
+        R = R0.copy()
+
+    return R

pyotc/otc_backend/policy_iteration/utils.py

@@ -0,0 +1,112 @@
+import numpy as np
+from scipy.optimize import linprog
+
+
+def get_best_stat_dist(P, c):
+    """
+    Given a transition matrix P and a cost vector c,
+    this function computes the stationary distribution that minimizes the expected cost
+    via linear programming.
+
+    Args:
+        P (np.ndarray): Transition matrix.
+        c (np.ndarray): Cost vector.
+
+    Returns:
+        stat_dist (np.ndarray): Best stationary distribution.
+        exp_cost (float): Corresponding expected cost.
+    """
+
+    # Set up constraints.
+    n = P.shape[0]
+    c = np.reshape(c, (n, -1))
+    Aeq = np.concatenate((P.T - np.eye(n), np.ones((1, n))), axis=0)
+    beq = np.concatenate((np.zeros((n, 1)), 1), axis=None)
+    beq = beq.reshape(-1, 1)
+    bound = [[0, None]] * n
+
+    # Solve linear program.
+    res = linprog(c, A_eq=Aeq, b_eq=beq, bounds=bound)
+    stat_dist = res.x
+    exp_cost = res.fun
+
+    return stat_dist, exp_cost
+
+
+def get_stat_dist(P, method="best", c=None):
+    """
+    Computes the stationary distribution of a Markov chain given its transition matrix P.
+
+    Supports multiple methods:
+        - 'best': Solves a linear program that minimizes cost under stationarity constraints.
+        - 'eigen': Solves for the stationary distribution using the eigenvalue method.
+        - 'iterative': Uses power iteration for large or sparse matrices.
+
+    Args:
+        P (np.ndarray): Transition matrix of the Markov chain, shape (n, n).
+        method (str): Method used to compute the stationary distribution.
+                      One of 'eigen', 'iterative', or 'best'. Defaults to 'best'.
+        c (np.ndarray, optional): Cost vector of shape (n,) used only when method='best'.
+
+    Returns:
+        pi (np.ndarray): Stationary distribution vector of shape (n,), summing to 1.
+
+    Raises:
+        ValueError: If method is 'best' but cost vector `c` is not provided,
+                    or if an invalid method name is given.
+    """
+    if method == "best":
+        # 'best' method minimizes expected cost under stationary constraints
+        if c is None:
+            raise ValueError("Cost function 'c' is required when method='best'.")
+
+        n = P.shape[0]
+        c = np.reshape(c, (n, -1))
+
+        # Stationarity constraint: π^T P = π^T  ⇨  (P^T - I)^T π = 0
+        # Add additional constraint: sum(π) = 1
+        Aeq = np.concatenate((P.T - np.eye(n), np.ones((1, n))), axis=0)
+        beq = np.concatenate((np.zeros((n, 1)), 1), axis=None)
+        beq = beq.reshape(-1, 1)
+
+        bound = [[0, None]] * n
+
+        # Solve the linear program: minimize c^T π s.t. Aeq π = beq
+        res = linprog(c, A_eq=Aeq, b_eq=beq, bounds=bound)
+        pi = res.x
+        return pi
+
+    elif method == "eigen":
+        # Computes the stationary distribution using eigenvalue decomposition
+        # Calculate the eigenvalues and eigenvectors
+        eigenvalues, eigenvectors = np.linalg.eig(P.T)
+
+        # Identify the eigenvector associated with eigenvalue closest to 1 and normalize to obtain a valid distribution
+        idx = np.argmin(np.abs(eigenvalues - 1))
+        pi = np.real(eigenvectors[:, idx])
+        pi /= np.sum(pi)
+        return pi
+
+    elif method == "iterative":
+        # Computes the stationary distribution using power iteration
+        max_iter = 10000
+        tol = 1e-10
+        n = P.shape[0]
+
+        # Start from uniform distribution
+        pi = np.ones(n) / n
+
+        for _ in range(max_iter):
+            pi_new = pi @ P
+            if np.linalg.norm(pi_new - pi, ord=1) < tol:
+                break
+            pi = pi_new
+
+        # Normalize the resulting distribution
+        pi /= np.sum(pi)
+        return pi
+
+    else:
+        raise ValueError(
+            f"Invalid method '{method}'. Must be one of 'best', 'eigen', or 'iterative'."
+        )

pyotc-0.2.2.dist-info/METADATA

@@ -0,0 +1,38 @@
+Metadata-Version: 2.4
+Name: pyotc
+Version: 0.2.2
+Summary: Perform optimal transition coupling (OTC) in python.
+Author: Jay Hineman, Yuning Pan, Bongsoo Yi
+License: MIT license
+License-File: AUTHORS.rst
+License-File: LICENSE
+Requires-Python: >=3.10
+Requires-Dist: matplotlib>=3.10.0
+Requires-Dist: networkx>=3.4.2
+Requires-Dist: numpy>=2.2.1
+Requires-Dist: pot>=0.9.5
+Description-Content-Type: text/markdown
+
+# pyotc
+[![codecov](https://codecov.io/github/pyotc/pyotc/graph/badge.svg?token=52QPNW0AP7)](https://codecov.io/github/pyotc/pyotc)
+
+A python implementations of optimal transport coupling algorithms.
+
+## Documentation
+Find sphinx documentation [here](https://pyotc.github.io/pyotc/).
+
+## Install
+See [install instructions](INSTALL.md)
+
+## Run Tests
+With a `uv` setup one can simply
+```bash
+uv run pytest
+```
+Otherwise, in `pip` installed context with deps met, `pytest` should behave as expected.
+
+## Contributing
+Guidelines for contribution to `pyotc` are provided in [CONTRIBUTING.md](./CONTRIBUTING.md).
+
+## Changelog
+A summary of changes and guide to versioning are recoreded in [CHANGELOG.md](./CHANGELOG.md).

pyotc-0.2.2.dist-info/RECORD

@@ -0,0 +1,34 @@
+pyotc/__init__.py,sha256=Od-nJxr4TP_mgmXMxpwaItNDISKzR0V6X8QakVNTmr8,125
+pyotc/otc.py,sha256=_we15B7Gw-IgxVlXnXbWS_8d0LGwHbaZiSrHwBSGbM0,94
+pyotc/examples/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+pyotc/examples/edge_awareness.py,sha256=k4pAs_UkDC2ZzgvkQZD9aY-gemjd8I9UAzWA1x8Z_3E,2513
+pyotc/examples/lollipops.py,sha256=fmvPkdSL9T5ZjQU8SpFnIox5DDuiVVwikjt7w3ukxLc,1589
+pyotc/examples/stochastic_block_model.py,sha256=aF-vZ9i3L75_rguTCmF3bD5qMI8T3QmlzdCMHivrVOw,1895
+pyotc/examples/wheel.py,sha256=JkBngoyHnppNUS2uXZSyRQOBhUynPOuVcEsYQdEcBQM,4177
+pyotc/otc_backend/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+pyotc/otc_backend/graph/__init__.py,sha256=n-t13QCEVXJ3naqd6gUnUAekIpUYwBjjO-f-WXWCbkg,58
+pyotc/otc_backend/graph/utils.py,sha256=mnUfgsfSEV79ej7a-4NUcLl-b956E6zKse3aYD4hRUo,3160
+pyotc/otc_backend/optimal_transport/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+pyotc/otc_backend/optimal_transport/logsinkhorn.py,sha256=V8fwUw-8UHQAKaeBGYGyq5bXilaM9IUAv4b78FEG29g,2069
+pyotc/otc_backend/optimal_transport/native.py,sha256=vxctK2knNoiyOGPV7oYMxGnh8S0iDyNW_D076dElNeo,1577
+pyotc/otc_backend/optimal_transport/native_refactor.py,sha256=83Ev5YuzYSyUyZPJTfRDk1xtXdjpYCJvH1KzV8MX2D8,1271
+pyotc/otc_backend/optimal_transport/pot.py,sha256=zeFEGJBQgVg2eFgAYHjd1ddbHQmKxrbUN1bQ4EmCZSc,1295
+pyotc/otc_backend/policy_iteration/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+pyotc/otc_backend/policy_iteration/utils.py,sha256=GCfgXhofn6ykqXkg0CJPm0VbO3mZr0CxgRwUGA2r9xc,3802
+pyotc/otc_backend/policy_iteration/dense/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+pyotc/otc_backend/policy_iteration/dense/approx_tce.py,sha256=PA92ysxJgRFwUbHxGYx6nctxr8GkL7I3WvakWfvXdNc,1150
+pyotc/otc_backend/policy_iteration/dense/entropic.py,sha256=uwdtxYOXipI3aXfmoY7U-55XumoH9gWGzWACSAGaDXY,6421
+pyotc/otc_backend/policy_iteration/dense/entropic_tci.py,sha256=TXKcQmLhpr9vYnjhIGO8sqx5O_7jwkY0pQ3eiFthtRg,1798
+pyotc/otc_backend/policy_iteration/dense/exact.py,sha256=hZqKq7zfZ5q7v3JnX37vCrpXJVDD84U496NDZA8OEvk,5001
+pyotc/otc_backend/policy_iteration/dense/exact_tce.py,sha256=WI5I9btsnrBksYDQ9BvnHe42pBaegNNrGRdYMrEkJKw,1851
+pyotc/otc_backend/policy_iteration/dense/exact_tci_lp.py,sha256=-mnwfJj_42V3DBZIjcfSFh9uwudEBm2gOOAwpiClRBM,1911
+pyotc/otc_backend/policy_iteration/dense/exact_tci_pot.py,sha256=9lJCB3JtnabDVeYAx6Igv7AkkAqEGa99LJnOOGNpIzY,3494
+pyotc/otc_backend/policy_iteration/sparse/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+pyotc/otc_backend/policy_iteration/sparse/exact.py,sha256=2VskHrp1nDGPV2yEsD5MBagrZHkUDCdChMZ28WOb0uk,4400
+pyotc/otc_backend/policy_iteration/sparse/exact_tce.py,sha256=oSymhuok3logzt0wYSYJeMkTls0zzC0kRiQw2nHXk7c,3626
+pyotc/otc_backend/policy_iteration/sparse/exact_tci.py,sha256=mW90o-2ZVd9fjgKFDlJhR223G5JYxBSzsHBr7NdK6eg,3446
+pyotc-0.2.2.dist-info/METADATA,sha256=yTtg2-uCXeJCSpSeh-NUWJ_i2NnxLzLjFuLcO46mexY,1122
+pyotc-0.2.2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+pyotc-0.2.2.dist-info/licenses/AUTHORS.rst,sha256=FtGxBvQjLU6Vdej_-zcsdWxTwnInCsjqejFRmaeC7TU,136
+pyotc-0.2.2.dist-info/licenses/LICENSE,sha256=DfWCeVpuOnLG1KJZfBO24ry8u90eeQi7mTdjoo5pt8M,1070
+pyotc-0.2.2.dist-info/RECORD,,

pyotc-0.2.2.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.27.0
+Root-Is-Purelib: true
+Tag: py3-none-any

pyotc-0.2.2.dist-info/licenses/AUTHORS.rst

@@ -0,0 +1,12 @@
+=======
+Credits
+=======
+
+* Jay Hineman, @jhineman
+* Bongsoo Yi
+* Yuning Pan
+
+Contributors
+------------
+
+None yet. Why not be the first?

pyotc-0.2.2.dist-info/licenses/LICENSE

@@ -0,0 +1,22 @@
+MIT License
+
+Copyright (c) 2024, Jay Hineman
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+