r""" Apply Bregman projection to a batch of affinity matrices :math:`\mathbf{K}`. The function solves the following optimization problem: .. math:: \begin{aligned} \mathbf{T} = \mathop{\arg \min}_\mathbf{T} \quad & \text{KL}(\mathbf{T} \| \mathbf{K}) \\
(
K, a=None, b=None, nx=None, max_iter=10000, tol=1e-5, grad="detach"
)
| 44 | |
| 45 | |
| 46 | def bregman_projection_batch( |
| 47 | K, a=None, b=None, nx=None, max_iter=10000, tol=1e-5, grad="detach" |
| 48 | ): |
| 49 | r""" |
| 50 | Apply Bregman projection to a batch of affinity matrices :math:`\mathbf{K}`. |
| 51 | |
| 52 | The function solves the following optimization problem: |
| 53 | |
| 54 | .. math:: |
| 55 | \begin{aligned} |
| 56 | \mathbf{T} = \mathop{\arg \min}_\mathbf{T} \quad & \text{KL}(\mathbf{T} \| \mathbf{K}) \\ |
| 57 | \text{s.t.} \quad & \mathbf{T} \mathbf{1} = \mathbf{a} \\ |
| 58 | & \mathbf{T}^T \mathbf{1} = \mathbf{b} \\ |
| 59 | & \mathbf{T} \geq 0 |
| 60 | \end{aligned} |
| 61 | |
| 62 | This is equivalent to: |
| 63 | |
| 64 | .. math:: |
| 65 | \begin{aligned} |
| 66 | \mathbf{T} = \mathop{\arg \max}_\mathbf{T} \quad & \langle \mathbf{T}, \log(\mathbf{K}) \rangle_F \\ |
| 67 | \text{s.t.} \quad & \mathbf{T} \mathbf{1} = \mathbf{a} \\ |
| 68 | & \mathbf{T}^T \mathbf{1} = \mathbf{b} \\ |
| 69 | & \mathbf{T} \geq 0 |
| 70 | \end{aligned} |
| 71 | |
| 72 | The optimal solution has the form :math:`\mathbf{T} = \text{diag}(\mathbf{f}) \mathbf{K} \text{diag}(\mathbf{g})`, |
| 73 | where the dual variables :math:`\mathbf{u}` and :math:`\mathbf{v}` are found iteratively using: |
| 74 | |
| 75 | .. math:: |
| 76 | \mathbf{f}^{(k+1)} = \frac{\mathbf{a}}{\sum \mathbf{K} \mathbf{g}^{(k)}} |
| 77 | |
| 78 | \mathbf{g}^{(k+1)} = \frac{\mathbf{b}}{\sum \mathbf{K}^T \mathbf{f}^{(k)}} |
| 79 | |
| 80 | Parameters |
| 81 | ---------- |
| 82 | K : array-like, shape (B, n, m) |
| 83 | Affinity matrix for each problem in the batch. |
| 84 | a : array-like, shape (B, n), optional |
| 85 | Source distribution for each problem. If None, uniform distribution is used. |
| 86 | b : array-like, shape (B, m), optional |
| 87 | Target distribution for each problem. If None, uniform distribution is used. |
| 88 | nx : backend object, optional |
| 89 | Numerical backend to use for computations. If None, the default backend is used. |
| 90 | max_iter : int, optional |
| 91 | Maximum number of iterations. |
| 92 | tol : float, optional |
| 93 | Tolerance for convergence. The solver stops when the maximum change in |
| 94 | the dual variables is below this value. |
| 95 | grad : str, optional |
| 96 | Gradient computation mode: 'detach', 'autodiff', or 'last_step'. |
| 97 | |
| 98 | Returns |
| 99 | ------- |
| 100 | dict |
| 101 | Dictionary containing: |
| 102 | - 'T' : array-like, shape (B, n, m) |
| 103 | Optimal transport plan for each problem. |