def sinkhornAutoTune(a, b, C):
entropic_reg = 1e-6
while True:
with np.errstate(over='raise', divide='raise'):
try:
print(f"Entropic regularizer: {entropic_reg}")
gamma = sinkhorn(a, b, C, reg=entropic_reg, verbose=True)
break
except Exception as e:
print(f"Caught Exception in SKnopp... {e}")
entropic_reg = entropic_reg * 2
def _egradv_doublyStochastic_gpu(self, vgamma):
mu, nu = np.ones(self.options.r) / self.options.r, np.ones(
self.options.r) / self.options.r
Mopt = sinkhorn(mu,
nu,
-vgamma.get().reshape(self.options.r, self.options.r),
reg=self.tikhR,
numItermax=self.options.sinkhorn_iters)
return np.einsum('id,d->i',
self.P.T,
Mopt.reshape(self.options.r**2),
dtype='double').reshape(self.m, self.n)
def stabilized_entropic_gromov_wasserstein(C1, C2, p, q, loss_fun, epsilon,
max_iter = 1000, tol=1e-9, verbose=False, log=False):
C1 = np.asarray(C1, dtype=np.float64)
C2 = np.asarray(C2, dtype=np.float64)
T = np.outer(p, q) # Initialization
constC, hC1, hC2 = init_matrix(C1, C2, p, q, loss_fun)
cpt = 0
err = 1
if log:
log = {'err': []}
while (err > tol and cpt < max_iter):
Tprev = T
# compute the gradient
tens = gwggrad(constC, hC1, hC2, T)
T = sinkhorn(p, q, tens, epsilon, method = 'sinkhorn_epsilon_scaling')
if cpt % 10 == 0:
# we can speed up the process by checking for the error only all
# the 10th iterations
err = np.linalg.norm(T - Tprev)
if log:
log['err'].append(err)
if verbose:
if cpt % 200 == 0:
print('{:5s}|{:12s}'.format(
'It.', 'Err') + '\n' + '-' * 19)
print('{:5d}|{:8e}|'.format(cpt, err))
cpt += 1
if log:
log['gw_dist'] = gwloss(constC, hC1, hC2, T)
return T, log
else:
return T
def compute_random_coupling(p, q, epsilon):
"""
Computes a random coupling based on:
KL-Proj_p,q(K) = argmin_T <-\epsilon logK, T> -\epsilon H(T)
where T is a couping matrix with marginal distributions p, and q, for rows and columns, respectively
This is solved with a Bregman Sinkhorn computation
p -- marginal distribution of rows
q -- marginal distribution of columns
epsilon -- entropy coefficient
"""
num_cells = len(p)
num_locations = len(q)
K = np.random.rand(num_cells, num_locations)
C = -epsilon * np.log(K)
return sinkhorn(p, q, C, epsilon)
def _compute_mappings(D1, D2):
try:
from ot.bregman import sinkhorn
except ModuleNotFoundError:
print("POT not found")
return
n1, n2 = D1.shape[0], D2.shape[0]
gamma = sinkhorn(a=(1 / n1) * np.ones(n1),
b=(1 / n2) * np.ones(n2),
M=pairwise_distances(D1, D2, metric="euclidean"),
reg=1e-1)
mappings = [np.zeros(n1, dtype=np.int32), np.zeros(n2, dtype=np.int32)]
for i in range(n1):
mappings[0][i] = np.argmax(gamma[i, :])
for i in range(n2):
mappings[1][i] = np.argmax(gamma[:, i])
return mappings
def generalized_conditional_gradient(self, source, target, source_labels):
'''
Algoritmo GCG: gradiente condicional generalizado. Calcula la matriz
de transporte gamma.
Argumentos
--------------------------------
source: numpy.ndarray shape(m, k).
datos del dominio fuente.
target: numpy.ndarray shape(n, k).
datos del dominio objetivo.
source_labels: numpy.ndarray shape(m, )
etiquetas del dominio fuente.
Retorno
-----------------------------------
G: numpy.ndarray shape(m, n)
Matriz gamma de transporte.
'''
loop = True
a, b = self.compute_probs(source), self.compute_probs(target)
G = np.outer(a, b)
M = self.cost_transport_matrix(source, target)
cost, cost_gradient, f, df = self.computation_graph(
source, target, source_labels
)
f_val = cost(G)
it = 0
while loop:
it += 1
old_fval = f_val
# problem linearization
Mi = M + self.eta * df(G)
# solve linear program with Sinkhorn
Gc = sinkhorn(
a, b, Mi, self.labda, numItermax=self.num_iters_sinkhorn
)
deltaG = Gc - G
# line search
# Que yo sepa la evaluación de la función se hace con G, no con
# Gc
dcost = cost_gradient(G)
alpha, fc, f_val = line_search_armijo(
cost, G, deltaG, dcost, f_val
)
if alpha is None: # No se si esto es lo correcto
return G
G = G + alpha * deltaG
# test convergence
if it >= self.num_iters_sinkhorn:
loop = False
# delta_fval = (f_val - old_fval) / abs(f_val)
delta_fval = (f_val - old_fval)
# print('fval:', f_val, 'delta_fval:', delta_fval)
if abs(delta_fval) < self.convergence_threshold:
loop = False
return G
def compute_gromov_wasserstein(C1,
C2,
p,
q,
tol,
gpu,
ot_warm,
reg,
maxiter,
loss_fun='square_loss'):
assert (len(C1) == len(p))
assert (len(C2) == len(q))
T = np.outer(p, q) # Initialization
constC, hC1, hC2 = init_matrix(C1, C2, T, p, q, loss_fun)
it = 0
err = 1
global_start = time()
# plt.ion() # To plot interactively during training
while (err > tol and it <= maxiter):
start = time()
Tprev = T
# compute the gradient
tens = gwggrad(constC, hC1, hC2, T)
# FIXME: Clean this up. Global vars and tailored imports should allow
# to not have devide gpu and non gpu cases.
if not gpu:
if ot_warm and it > 0:
T, log = sinkhorn_knopp(p,
q,
tens,
reg,
init_u=log['u'],
init_v=log['v'],
log=True)
elif ot_warm:
T, log = bregman.sinkhorn(p, q, tens, reg, log=True)
else:
T = bregman.sinkhorn(p, q, tens, reg)
else:
if ot_warm and it > 0:
T, log = bregman.sinkhorn(p,
q,
tens,
reg,
returnAsGPU=True,
log=True)
elif ot_warm:
T, log = bregman.sinkhorn(p, q, tens, reg, log=True)
else:
T = bregman.sinkhorn(p, q, tens, reg, returnAsGPU=True)
time_G = time() - start
if it % 10 == 0:
# we can speed up the process by checking for the error only all
# the 10th iterations
dist = gwloss(constC, hC1, hC2, T)
if gpu:
err = np.linalg.norm(T.copy().subtract(Tprev).asarray())
else:
err = np.linalg.norm(T - Tprev)
it += 1
# plt.ioff() # To plot interactively during training
plt.close('all')
# if log:
# log['gw_dist'] = gwloss(constC, hC1, hC2, T)
# return T, log
# else:
if gpu:
return T.asarray()
else:
return T
def gromov_wasserstein_adjusted_norm(cost_mat,
C1,
C2,
alpha_linear,
p,
q,
loss_fun,
epsilon,
max_iter=1000,
tol=1e-9,
verbose=False,
log=False):
"""
Returns the gromov-wasserstein coupling between the two measured similarity matrices
(C1,p) and (C2,q)
The function solves the following optimization problem:
.. math::
\GW = arg\min_T \sum_{i,j,k,l} L(C1_{i,k},C2_{j,l})*T_{i,j}*T_{k,l}-\epsilon(H(T))
s.t. \GW 1 = p
\GW^T 1= q
\GW\geq 0
Where :
M : cost matrix in sourceXtarget space
C1 : Metric cost matrix in the source space
C2 : Metric cost matrix in the target space
p : distribution in the source space
q : distribution in the target space
L : loss function to account for the misfit between the similarity matrices
H : entropy
Parameters
----------
M : ndarray, shape (ns, nt)
Cost matrix in the sourceXtarget space
C1 : ndarray, shape (ns, ns)
Metric cost matrix in the source space
C2 : ndarray, shape (nt, nt)
Metric costfr matrix in the target space
p : ndarray, shape (ns,)
distribution in the source space
q : ndarray, shape (nt,)
distribution in the target space
loss_fun : string
loss function used for the solver either 'square_loss' or 'kl_loss'
epsilon : float
Regularization term >0
max_iter : int, optional
Max number of iterations
tol : float, optional
Stop threshold on error (>0)
verbose : bool, optional
Print information along iterations
log : bool, optional
record log if True
Returns
-------
T : ndarray, shape (ns, nt)
coupling between the two spaces that minimizes :
\sum_{i,j,k,l} L(C1_{i,k},C2_{j,l})*T_{i,j}*T_{k,l}-\epsilon(H(T))
"""
C1 = np.asarray(C1, dtype=np.float64)
C2 = np.asarray(C2, dtype=np.float64)
cost_mat = np.asarray(cost_mat, dtype=np.float64)
T = np.outer(p, q) # Initialization
cpt = 0
err = 1
try:
cost_mat_norm = cost_mat / cost_mat.max()
except:
cost_mat_norm = cost_mat
if alpha_linear == 1:
T = sinkhorn(p, q, cost_mat_norm, epsilon)
else:
while (err > tol and cpt < max_iter):
Tprev = T
if loss_fun == 'square_loss':
tens = tensor_square_loss_adjusted(C1, C2, T)
tens_all = (1 - alpha_linear) * tens + alpha_linear * cost_mat_norm
T = sinkhorn(p, q, tens_all, epsilon)
if cpt % 10 == 0:
# We can speed up the process by checking for the error only all
# the 10th iterations
err = np.linalg.norm(T - Tprev)
if log:
log['err'].append(err)
if verbose:
if cpt % 200 == 0:
print('{:5s}|{:12s}'.format('It.', 'Err') + '\n' +
'-' * 19)
print('{:5d}|{:8e}|'.format(cpt, err))
cpt += 1
if log:
return T, log
else:
return T
def entropic_gromov_wasserstein(self,
C1,
C2,
p,
q,
m,
M=None,
loss_fun='square_loss'):
C1 = np.asarray(C1, dtype=np.float32)
C2 = np.asarray(C2, dtype=np.float32)
T0 = np.outer(p, q) # Initialization
dim_G_extended = (len(p) + self.virtual_cells,
len(q) + self.virtual_cells)
q_extended = np.append(q, [(np.sum(p) - m) / self.virtual_cells] *
self.virtual_cells)
p_extended = np.append(p, [(np.sum(q) - m) / self.virtual_cells] *
self.virtual_cells)
q_extended = q_extended / np.sum(q_extended)
p_extended = p_extended / np.sum(p_extended)
constC, hC1, hC2 = init_matrix(C1, C2, p, q, loss_fun)
cpt = 0
err = 1
while (err > self.tol and cpt < self.max_iter):
Gprev = T0
# compute the gradient
if abs(m - 1) < 1e-10: # full match
Ck = gwggrad(constC, hC1, hC2, T0)
else: # partial match
Ck = gwgrad_partial(C1, C2, T0)
if M is not None:
Ck = Ck * M
Ck_emd = np.zeros(dim_G_extended)
Ck_emd[:len(p), :len(q)] = Ck
Ck_emd[-self.virtual_cells:,
-self.virtual_cells:] = 100 * np.max(Ck_emd)
Ck_emd = np.asarray(Ck_emd, dtype=np.float64)
# T = sinkhorn(p, q, Ck, epsilon, method = 'sinkhorn')
T = sinkhorn(p_extended,
q_extended,
Ck_emd,
self.epsilon,
method='sinkhorn')
T0 = T[:len(p), :len(q)]
if cpt % 10 == 0:
err = np.linalg.norm(T0 - Gprev)
if self.verbose:
if cpt % 200 == 0:
print('{:5s}|{:12s}'.format('Epoch.', 'Loss') + '\n' +
'-' * 19)
print('{:5d}|{:8e}|'.format(cpt, err))
cpt += 1
return T
