Function bregman_log_projection_batch

ot/batch/_utils.py:165–290

r""" Apply Bregman projection to a batch of affinity matrices :math:`\exp(\mathbf{K})`. The function solves the following optimization problem: .. math:: \begin{aligned} \mathbf{T} = \mathop{\arg \min}_\mathbf{T} \quad & \text{KL}(\mathbf{T} \| \exp(\mathbf{K})) \\

(
    K, a=None, b=None, nx=None, max_iter=10000, tol=1e-5, grad="detach"
)

def bregman_log_projection_batch(
    K, a=None, b=None, nx=None, max_iter=10000, tol=1e-5, grad="detach"
):
    r"""
    Apply Bregman projection to a batch of affinity matrices :math:`\exp(\mathbf{K})`.

    The function solves the following optimization problem:

    .. math::
        \begin{aligned}
            \mathbf{T} = \mathop{\arg \min}_\mathbf{T} \quad & \text{KL}(\mathbf{T} \| \exp(\mathbf{K})) \\
            \text{s.t.} \quad & \mathbf{T} \mathbf{1} = \mathbf{a} \\
            & \mathbf{T}^T \mathbf{1} = \mathbf{b} \\
            & \mathbf{T} \geq 0
        \end{aligned}

    This is equivalent to:

    .. math::
        \begin{aligned}
            \mathbf{T} = \mathop{\arg \max}_\mathbf{T} \quad & \langle \mathbf{T}, \mathbf{K} \rangle_F \\
            \text{s.t.} \quad & \mathbf{T} \mathbf{1} = \mathbf{a} \\
            & \mathbf{T}^T \mathbf{1} = \mathbf{b} \\
            & \mathbf{T} \geq 0
        \end{aligned}

    The optimal solution has the form :math:`\mathbf{T} = \text{diag}(\exp(\mathbf{u})) \exp(\mathbf{K}) \text{diag}(\exp(\mathbf{v}))`,
    where the dual variables :math:`\mathbf{u}` and :math:`\mathbf{v}` are found iteratively using:

    .. math::
        \mathbf{u}^{(k+1)} = \log(\mathbf{a}) - \text{LSE}(\mathbf{K} + \mathbf{v}^{(k)})

        \mathbf{v}^{(k+1)} = \log(\mathbf{b}) - \text{LSE}(\mathbf{K}^T + \mathbf{u}^{(k)})

    where LSE denotes the log-sum-exp operation. The iterations are performed using
    the log-sum-exp trick to avoid numerical issues.

    Parameters
    ----------
    K : array-like, shape (B, n, m)
        Affinity matrix for each problem in the batch.
    a : array-like, shape (B, n), optional
        Source distribution for each problem. If None, uniform distribution is used.
    b : array-like, shape (B, m), optional
        Target distribution for each problem. If None, uniform distribution is used.
    nx : backend object, optional
        Numerical backend to use for computations. If None, the default backend is used.
    max_iter : int, optional
        Maximum number of iterations.
    tol : float, optional
        Tolerance for convergence. The solver stops when the maximum change in
        the dual variables is below this value.
    grad : str, optional
        Gradient computation mode: 'detach', 'autodiff', or 'last_step'.

    Returns
    -------
    dict