Solve OT problem with a given cost matrix

[mtw33]

Hi,

we got stuck on the following problem. We would like to solve the OT problem for a given matrix C in R^{m \times n}.

To do so we defined the following cost function

@jax.tree_util.register_pytree_node_class
class DiscreteCostFn(costs.CostFn):
  def __init__(self, C):
      super().__init__()
      self._C = C

  def norm(self, x):
    return jnp.zeros_like(x).squeeze()

  def pairwise(self, x, y):
    return self._C[x, y][0]

The following code runs

C = jnp.array([[0,1.],[1.,0]])
cost = DiscreteCostFn(C)
pc_geom = pointcloud.PointCloud(X, Y, cost_fn=cost, epsilon=1)  
geom = ott.geometry.geometry.Geometry(cost_matrix=pc_geom.cost_matrix)
prob = linear_problem.LinearProblem(geom, a=a, b=b)
solver = sinkhorn.Sinkhorn(inner_iterations=10, max_iterations=10000, jit=True)
out = solver(prob)

But it's not giving the right solution to the OT problem ...... any help on what we are doing wrong here would be greatly appreciated.

Thanks a lot
Marie-Therese

[marcocuturi]

Hi Marie-Thérèse, thanks a lot for your interest in ott!!

It seems, if you do know you cost matrix C, that you should simply pass it directly to instantiate a plain geometry object, not a point cloud.

In your example, this would mean

C = jnp.array([[0,1.],[1.,0]])
geom = ott.geometry.geometry.Geometry(cost_matrix=C, epsilon = 1.)

and then continue with the other instructions. The point cloud geometry is indeed only used when you don't want to instantiate a full cost matrix, and, instead, prefer to provide points (in vector form) along with a cost function (e.g. two families of vectors and the squared euclidean distance). Then the toolbox does those computations for you.

WDYT?

[mtw33]

Thanks so much! Exactly what I was looking for :-)