text: #!/usr/bin/env python
  """
  sinkhorn_pointcloud.py
  
  Discrete OT : Sinkhorn algorithm for point cloud marginals.
  
  """
  
  import torch
  from torch.autograd import Variable
  
  # Decide which device we want to run on
  device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
  
  
  def sinkhorn_normalized(x, y, epsilon, niter, mass_x=None, mass_y=None):
  
      Wxy = sinkhorn_loss(x, y, epsilon, niter, mass_x, mass_y)
      Wxx = sinkhorn_loss(x, x, epsilon, niter, mass_x, mass_x)
      Wyy = sinkhorn_loss(y, y, epsilon, niter, mass_y, mass_y)
      return 2 * Wxy - Wxx - Wyy
  
  def sinkhorn_loss(x, y, epsilon, niter, mass_x=None, mass_y=None):
      """
      Given two emprical measures with n points each with locations x and y
      outputs an approximation of the OT cost with regularization parameter epsilon
      niter is the max. number of steps in sinkhorn loop
      """
  
      # The Sinkhorn algorithm takes as input three variables :
      C = cost_matrix(x, y)  # Wasserstein cost function
  
      nx = x.shape[1]
      ny = y.shape[1]
      batch_size = x.shape[0]
  
      if mass_x is None:
          # assign marginal to fixed with equal weights
          mu = 1. / nx * torch.ones([batch_size, nx]).to(device)
      else: # normalize
          mass_x.data = torch.clamp(mass_x.data, min=0, max=1e9)
          mass_x = mass_x + 1e-9
          mu = (mass_x / mass_x.sum(dim=-1, keepdim=True)).to(device)
  
      if mass_y is None:
          # assign marginal to fixed with equal weights
          nu = 1. / ny * torch.ones([batch_size, ny]).to(device)
      else: # normalize
          mass_y.data = torch.clamp(mass_y.data, min=0, max=1e9)
          mass_y = mass_y + 1e-9
          nu = (mass_y / mass_y.sum(dim=-1, keepdim=True)).to(device)
  
      def M(u, v):
          "Modified cost for logarithmic updates"
          "$M_{ij} = (-c_{ij} + u_i + v_j) / \epsilon$"
          return (-C + u.unsqueeze(2) + v.unsqueeze(1)) / epsilon
  
      def lse(A):
          "log-sum-exp"
          return torch.log(torch.exp(A).sum(2, keepdim=True) + 1e-6)  # add 10^-6 to prevent NaN
  
      # Actual Sinkhorn loop ......................................................................
      u, v, err = 0. * mu, 0. * nu, 0.
  
      for i in range(niter):
          u = epsilon * (torch.log(mu) - lse(M(u, v)).squeeze()) + u
          v = epsilon * (torch.log(nu) - lse(M(u, v).transpose(dim0=1, dim1=2)).squeeze()) + v
  
      U, V = u, v
      pi = torch.exp(M(U, V))  # Transport plan pi = diag(a)*K*diag(b)
      cost = torch.sum(pi * C, dim=[1, 2])  # Sinkhorn cost
  
      return torch.mean(cost)
  
  
  def cost_matrix(x, y, p=2):
      "Returns the matrix of $|x_i-y_j|^p$."
      x_col = x.unsqueeze(2)
      y_lin = y.unsqueeze(1)
      c = torch.sum((torch.abs(x_col - y_lin)) ** p, -1)
      return c
  