Manipulation of cost matrix and infinite cost

[PaulArthurM]

Problem:
While trying to reproduce the codes from the Single-cell genomics page on my own data, I wondered if it was possible to manipulate the cost matrix before using the sinkhorn function.

If we know that some couplings are impossible, independently of their distance between the x and y arrays, is it possible to define these costs as infinite to penalize them when running the sinkhorn function ? Is it theoretically acceptable to have infinite costs when using OT with the sinkhorn function or others OT algorithms? Are there alternatives to penalize couplings that do not involve infinite values? Should i just run sinkhorn function for several separated PointCloud objects with only possible coupling observations (no infinite values then)?

Solutions considered:
In case infinite values is acceptable, how should I create this new cost matrix?

• Replacing the cost matrix with geom.cost_matrix = new_cost_matrix returns an error (AttributeError).
• Create a custom CostFn that performs this step directly, in a custom pairwise method for example?
• Or a new PointCloud method/parameter to multiply by elements the cost matrix by a matrix of 1.0 or infinite values (something similar to scale_cost but for matrix multiplication)?

[marcocuturi]

Hi Paul! thanks so much for trying out OTT.

I have migrated this to a discussion, because I think the points you mention sound more like a discussion on OTT and regularized OT than issues w.r.t the code base itself.

Here are a few answers inlined:

If we know that some couplings are impossible, independently of their distance between the x and y arrays, is it possible to define these costs as infinite to penalize them when running the sinkhorn function ? Is it theoretically acceptable to have infinite costs when using OT with the sinkhorn function or others OT algorithms?

In general this can be tricky, because setting some entries of the cost matrix to infinite might, in many cases, for the problem to have an infinite optimum. When using Sinkhorn, it can be even trickier. Basically the Sinkhorn algorithm is guaranteed to converge when the kernel matrix $K$ (which is $exp(-C/\varepsilon)$) has no zeroes (the true theorem is a bit more detailed, but for our purposes, Sinkhorn might run into issues in that case.

Are there alternatives to penalize couplings that do not involve infinite values? Should i just run sinkhorn function for several separated PointCloud objects with only possible coupling observations (no infinite values then)?

I am not sure I understand this fix.

Solutions considered:
In case infinite values is acceptable,

In general, that won't be acceptable, but you can set it to a large value indeed.

how should I create this new cost matrix?

I think that if you want to enter a custom cost matrix, the "JAX" functional way would be to consider

cost_matrix = pointcloud.PointCloud(x,y).cost_matrix
... # do your computations to threshold the matrix
cost_matrix = cost_matrix.at(...).set(...) # do replacements
geom = geometry.Geometry(cost_matrix=cost_matrix)
out = sinkhorn.sinkhorn(geom)

Replacing the cost matrix with geom.cost_matrix = new_cost_matrix returns an error (AttributeError).

Yes, usually the cost_matrix member of a Geometry object is a property, and cannot be changed in place. This is in agreement with JAX's functional principle. You should instead reinstantiate a geometry, as above.

Create a custom CostFn that performs this step directly, in a custom pairwise method for example?

That would be much better.

Or a new PointCloud method/parameter to multiply by elements the cost matrix by a matrix of 1.0 or infinite values (something similar to scale_cost but for matrix multiplication)?

You mean like a mask? I think that if your criteria to decide whether the cost is very large or multiplied by 1 only depends on x and y (and x and y only, not on the overall values), the CostFn approach above is much better. See the discussion with #164

[PaulArthurM]

Hi Marco! Thanks for your explanation, you answered all my questions and I managed to enter a custom cost function with the functional method for now.

Regarding this part:

Are there alternatives to penalize couplings that do not involve infinite values? Should i just run sinkhorn function for several separated PointCloud objects with only possible coupling observations (no infinite values then)?

A rephrasing of this question might be: if all my observations in x and y belong to multiple classes and I want to perform an OT between x and y with coupling only between observations of the same class, instead of manually penalizing couplings in the cost matrix, can I split x and y into subarrays with only observations of the same class and run the sinkhorn.sinkhorn function independently for each?

Intuitively, I expect to get similar results with both methods, but perhaps one is more "mathematically" relevant (if the penalty is high enough to make couplings between different classes impossible).

[marcocuturi]

Hi @PaulArthurM
apologies for the delayed answer to your final question.
If you know that transport should only happen between certain classes, then, indeed, it makes more sense to do them independently.
If you want to allow some tranport intra-classes, but penalize it somewhat, the simplest would be to add in the cost function another (discrete) cost, that penalizes when points are mapped to different classes. To facilitate some of these computations, using an unabalanced setting (tau_a and tau_b<1) could also help.