def nlog_density(self, target, source, log_weights, scales, probas=None):
"""Negative log-likelihood of the proposal generated by the source onto the target."""
x_i = LazyTensor(target[:, None, :]) # (N,1,D)
y_j = LazyTensor(self.means[None, :, :]) # (1,M,D)
s_j = self.covariances_inv # (M, D, D)
s_j = LazyTensor(s_j.view(s_j.shape[0], -1)[None, :, :]) # (1, M, D*D)
D_ij = (x_i - y_j) | s_j.matvecmult(x_i - y_j) # (N,M,1)
det_j = LazyTensor(self.log_det_cov_half[None, :, None])
logK_ij = (-D_ij / 2 - (self.D / 2) * float(np.log(2 * np.pi)) - det_j)
log_weights = self.weights.log()
logW_j = LazyTensor(log_weights[None, :, None])
logK_ij = logK_ij + logW_j
logdensities_i = logK_ij.logsumexp(dim=1).reshape(-1) # (N,)
return -logdensities_i
def nlog_density(self, target, source, log_weights, scales, probas=None):
"""Negative log-likelihood of the proposal generated by the source onto the target."""
if self.adaptive:
x_i = LazyTensor(target[:, None, :]) # (N,1,D)
y_j = LazyTensor(source[None, :, :]) # (1,M,D)
s_j = self.covariances_inv # (M, D, D)
s_j = LazyTensor(s_j.view(s_j.shape[0],
-1)[None, :, :]) # (1, M, D*D)
D_ij = (x_i - y_j) | s_j.matvecmult(x_i - y_j) # (N,M,1)
det_j = LazyTensor(self.log_det_cov_half[None, :, None])
logK_ij = (-D_ij / 2 - (self.D / 2) * float(np.log(2 * np.pi)) -
det_j)
else:
D_ij = squared_distances(target, source)
logK_ij = (-D_ij / (2 * scales**2) -
(self.D / 2) * float(np.log(2 * np.pi)) -
self.D * scales.log())
if log_weights is None:
logK_ij = logK_ij - float(np.log(len(source)))
else:
logW_j = LazyTensor(log_weights[None, :, None])
logK_ij = logK_ij + logW_j
logdensities_i = logK_ij.logsumexp(dim=1).view(-1) # (N,)
if probas is None:
return -logdensities_i
else:
return -(logdensities_i.view(-1, len(probas)) +
probas.log()[None, :]).logsumexp(dim=1).view(-1)
def routine(self,
α,
x,
β,
y,
class_xy=None,
class_yx=None,
class_xx=None,
class_yy=None,
mask_diagonal=False,
**kwargs):
N, D = x.shape
M, _ = y.shape
if self.debias:
C_xx = self.calculate_cost(x, x.detach(), class_xx)
C_yy = self.calculate_cost(y, y.detach(), class_yy)
else:
C_xx, C_yy = Non
C_xy = self.calculate_cost(x, y.detach(), class_xy)
C_yx = self.calculate_cost(y, x.detach(), class_yx)
softmin = partial(self.softmin_online,
log_conv=self.logconv(C_xy, dtype=str(x.dtype)[6:]))
diameter, ε, ε_s, ρ = scaling_parameters(x, y, self.p, self.blur,
self.reach, self.diameter,
self.scaling)
a_x, b_y, a_y, b_x = sinkhorn_loop(softmin,
log_weights(α),
log_weights(β),
C_xx,
C_yy,
C_xy,
C_yx,
ε_s,
ρ,
debias=self.debias)
F, G = sinkhorn_cost(ε,
ρ,
α,
β,
a_x,
b_y,
a_y,
b_x,
debias=self.debias,
potentials=True)
a_i = α.view(-1, 1)
b_j = β.view(1, -1)
F_i, G_j = F.view(-1, 1), G.view(1, -1)
cost = (F_i + G_j).mean()
# coupling calculation
F_i = LazyTensor(F_i.view(-1, 1, 1))
G_j = LazyTensor(G_j.view(1, -1, 1))
a_i = LazyTensor(a_i.view(-1, 1, 1))
b_j = LazyTensor(b_j.view(1, -1, 1))
if len(C_xy) == 2:
x, y = C_xy
C_ij = self.cost(x, y)
else:
x, y, Z, L = C_xy
C_ij = self.cost(x, y, class_values=[Z, L])
coupling = ((F_i + G_j - C_ij) / self.eps).exp() * (a_i * b_j)
return cost, coupling, C_ij, [F_i, G_j, a_i, b_j]
