F*, gradf*, utilitary function

wdnmf_dual/methods/dual_grad.py

@@ -9,6 +9,21 @@ import ot as ot
 import torch as tn
 import numpy as np
 
+def log_mat(v):
+    n, m = v.shape
+    w = tn.zeros(n, m, dtype=tn.float64)
+    for i in range(n):
+        for j in range(m):
+            w[i, j] = np.log(v[i, j])
+    return w
+
+def scalar_product_mat(v, w):
+    n, m = v.shape
+    temp = 0
+    for j in range(m):
+        temp += float(tn.inner(v[:,j], w[:,j]))
+    return temp
+
 def simplex_prox_mat(v, z=1):
     """
     For a given matrix V seen as a collection of vectors, return the matrix W
@@ -20,8 +35,8 @@ def simplex_prox_mat(v, z=1):
     ----------
     v : 2D tn.tensor
         collection of vectors to project on the simplex
-    z : TYPE, optional
-        DESCRIPTION. The default is 1.
+    z : Float, optional
+        The columns of v are normalized to z. The default is 1.
 
     Returns
     -------
@@ -149,7 +164,7 @@ def grad_e_dual(v):
         w[i] = np.exp(v[i])
     return -w/sum(w)
 
-def f_vec_aux(source, aim, cost, reg, method = 'sinkhorn', numItermax = 1000, thr = 1e-9, verbose = False):
+def f_vec(source, aim, cost, reg, method = 'sinkhorn', numItermax = 1000, thr = 1e-9, verbose = False):
     """
     Return the entropic Wasserstein distance between the "source" vector and the "aim" vector, 
     with respect to the ground cost "cost". The entropic regulation is T \rightarrow \langle T, \log T \rangle
@@ -180,12 +195,168 @@ def f_vec_aux(source, aim, cost, reg, method = 'sinkhorn', numItermax = 1000, th
 
     Returns
     -------
-    temp : TYPE
-        DESCRIPTION.
+    temp : float
+        Distance between the source and aim vector
 
     """
     temp = ot.bregman.sinkhorn2(a = source, b = aim, M = cost, reg = reg, method = method, 
                                 numItermax = numItermax, stopThr = thr, verbose = verbose, 
                                 log = False, warn = False)
     return temp
-    
\ No newline at end of file
+    
+def f_mat(source, aim, cost, reg, method = 'sinkhorn', numItermax = 1000, thr = 1e-9, verbose = False):
+    """
+    Return the entropic Wasserstein distance between the "source"'s vector and the "aim"'s vector, 
+    with respect to the ground cost "cost". The entropic regulation is T \rightarrow \langle T, \log T \rangle
+    See POT (pythonot.github.io) for more info, this is basically just the 
+    sinkhorn algorithm from this package.
+
+    Parameters
+    ----------
+    source : 2D tn.tensor
+        Source matrix, need to be on the simplex columnwise
+    aim : 2D tn.tensor
+        Aimed matrix, need to be on the simplex columnwise
+    cost : 2D tn.tensor
+        Ground cost, need to match source and aim dimension
+    reg : float
+        Regulation rate. The closer to 0, the sharper the transport is
+    method : string,  optional
+        ‘sinkhorn’,’sinkhorn_log’, ‘greenkhorn’, ‘sinkhorn_stabilized’ or ‘sinkhorn_epsilon_scaling’
+        depending on what solver to use. Sinkhorn stabilized should be the best for
+        sharper transport with the smaller reg rate. The default is 'sinkhorn'.
+    numItermax : int, optional
+        Number of iteration before stoping the calculation of the transport plan. The default is 1000.
+    thr : float, optional
+        Precision threshold, the calculation of the transport plan stop if the \delta between
+        two iterations of the computation is smaller than thr. The default is 1e-9.
+    verbose : Boolean, optional
+        Verbose or not verbose, this is a debug question. The default is False.
+
+    Returns
+    -------
+    temp : float
+        Sum of the distances between the source's and aim's vectors
+
+    """
+    temp = 0
+    n, m = source.shape
+    for j in range(m):
+        temp += ot.bregman.sinkhorn2(a = source[:,j], b = aim[:,j], M = cost, reg = reg, method = method, 
+                                    numItermax = numItermax, stopThr = thr, verbose = verbose, 
+                                    log = False, warn = False)
+    return temp
+
+def f_vec_aux(aim, cost, reg, method = 'sinkhorn', numItermax = 1000, thr = 1e-9, verbose = False):
+    """
+    Return an easier function to manipulate than f_vec
+    Parameters
+    ----------
+    aim : 1D tn.tensor
+        Aimed vector, need to be on the simplex
+    cost : 2D tn.tensor
+        Ground cost, need to match source and aim dimension
+    reg : float
+        Regulation rate. The closer to 0, the sharper the transport is
+    method : string,  optional
+        ‘sinkhorn’,’sinkhorn_log’, ‘greenkhorn’, ‘sinkhorn_stabilized’ or ‘sinkhorn_epsilon_scaling’
+        depending on what solver to use. Sinkhorn stabilized should be the best for
+        sharper transport with the smaller reg rate. The default is 'sinkhorn'.
+    numItermax : int, optional
+        Number of iteration before stoping the calculation of the transport plan. The default is 1000.
+    thr : float, optional
+        Precision threshold, the calculation of the transport plan stop if the \delta between
+        two iterations of the computation is smaller than thr. The default is 1e-9.
+    verbose : Boolean, optional
+        Verbose or not verbose, this is a debug question. The default is False.
+
+    Returns
+    -------
+    f : function
+        Same than f_mat, but you only have to add the source.
+
+    """
+    def f(source):
+        return f_vec(source, aim, cost, reg, method = 'sinkhorn', numItermax = 1000, thr = 1e-9, verbose = False)
+    return f
+
+def f_mat_aux(aim, cost, reg, method = 'sinkhorn', numItermax = 1000, thr = 1e-9, verbose = False):
+    """
+    Return an easier function to manipulate than f_mat
+    Parameters
+    ----------
+    aim : 2D tn.tensor
+        Aimed matrix, need to be on the simplex columnwise
+    cost : 2D tn.tensor
+        Ground cost, need to match source and aim dimension
+    reg : float
+        Regulation rate. The closer to 0, the sharper the transport is
+    method : string,  optional
+        ‘sinkhorn’,’sinkhorn_log’, ‘greenkhorn’, ‘sinkhorn_stabilized’ or ‘sinkhorn_epsilon_scaling’
+        depending on what solver to use. Sinkhorn stabilized should be the best for
+        sharper transport with the smaller reg rate. The default is 'sinkhorn'.
+    numItermax : int, optional
+        Number of iteration before stoping the calculation of the transport plan. The default is 1000.
+    thr : float, optional
+        Precision threshold, the calculation of the transport plan stop if the \delta between
+        two iterations of the computation is smaller than thr. The default is 1e-9.
+    verbose : Boolean, optional
+        Verbose or not verbose, this is a debug question. The default is False.
+
+    Returns
+    -------
+    f : function
+        Same than f_mat, but you only have to add the source.
+
+    """
+    def f(source):
+        return f_mat(source, aim, cost, reg, method = 'sinkhorn', numItermax = 1000, thr = 1e-9, verbose = False)
+    return f
+
+def matK(cost, reg):
+    n, m = cost.shape
+    w = tn.zeros(n, m, dtype=tn.float64)
+    for i in range(n):
+        for j in range(m):
+            w[i, j] = np.exp(-cost[n, m]/reg)
+    return w
+
+def vecA(g, reg):
+    n = g.shape[0]
+    w = tn.zeros(n, dtype=tn.float64)
+    for i in range(n):
+        w[i] = np.exp(g[i]/reg)
+    return w
+
+def f_dual(g, aim, reg, matK):
+    alpha = vecA(g, reg)    
+    a = e_vec(aim) + scalar_product_mat(aim, log_mat(tn.matmul(matK, alpha)))
+    return reg * a
+
+def grad_f_dual(g, aim, reg, matK):
+    """
+    Compute the gradient of $f^*$ for a given aim vector. Need the computation 
+    of matK to limit the calculation time.
+
+    Parameters
+    ----------
+    g : 1D tn.tensor
+        Vector within the simplex, this is the dual variable
+    aim : 1D tn.tensor
+        Aimed vector
+    matK : 2D tn.tensor
+        Precalculated matrix, $\matK = \exp^{-\frac{\matM}{ref}}
+    reg : float64
+        Used to compute $\alpha = \exp^{g / reg}$
+
+    Returns
+    -------
+    b : 1d tn.tensor
+        Gradient of $f^*$ 
+
+    """
+    matKT = tn.transpose(matK, 0, 1)
+    alpha = vecA(g, reg)
+    a = tn.mul(matKT, tn.div(aim, tn.matmul(matK, alpha)))
+    b = tn.mul(alpha, a)
+    return b