Correct way to measure divergence between similar point clouds

[arturtoshev]

Here is what I want to do:
Given are two 2D point clouds of N*N particles of the same weight in a periodic box $[0,1]^2$. What is the correct way to measure the distance in their distribution? By distance, I mean that if both particle sets are on a Cartesian grid (say slightly shifted to each other), then the distance should be close to zero and increase the more shifted the clouds are, and if one of the particle sets populates only a subset of the domain, then the distance should increase.

We have been using the ott.tools.sinkhorn_divergence.sinkhorn_diverngence().divergence function so far, but I experimented with some edge cases and now I doubt that I understand what this function does (I mean that I didn't expect negative divergence to be allowed and I also expected that the more shifted the point clouds are, the larger the divergence gets). Then, I also looked at ott.solvers.linear.sinkhorn.Sinkhorn(ott.problems.linear.linear_problem.LinearProblem(ott.geometry.pointcloud.PointCloud())).reg_ot_cost as I expected this to be somehow proportional to the divergence, but this quantity is extremely dependent on an epsilon parameter, which I don't understand.

The code I experimented with is here and the plots I got look like that:

If the negative divergence is only because we are in the range of numerical precision errors (I actually ran the whole thing also with double precision, but still got negative divergences), then what is the correct way to capture slight deviations in the actual particle distribution, e.g. the onset of particle clustering?

Best, Artur

[michalk8]

Hi @arturtoshev ,

yes, I believe the small negative values are just coming from numerical imprecision.
As for why it doesn't have the desired behavior, I'm not familiar with the ground cost used in the snippet you provided. Using SqEuclidean as a ground cost, it looks like this (between x/y):

Whereas the one from jax-md looks like this (and has similar behavior to mse):

Then, I also looked at ott.solvers.linear.sinkhorn.Sinkhorn(ott.problems.linear.linear_problem.LinearProblem(ott.geometry.pointcloud.PointCloud())).reg_ot_cost

The regularized OT cost is not a divergence, so if need to have a divergence, I wouldn't use it and would use Sinkhorn divergence instead.

as I expected this to be somehow proportional to the divergence, but this quantity is extremely dependent on an epsilon parameter, which I don't understand.

Epsilon the entropy regularization parameter - as epsilon -> 0, it recovers the original OT problem (can be solved, e.g., via the Hungarian algorithm in $O(n^3)$ ).
Lower epsilon makes the coupling matrix more crisp, but also and also depends on the scale of the values of the cost matrix. If not specified, by default we use epsilon = 0.05 * mean(cost_matrix).

Afaik, in our notebooks we only show the effect of here here, see the animation at the bottom as epsilon increases).
For your use-case, I'd experiment with a different range of values, but definitely wouldn't go as low as 1e-8 as in the code you provided; lowest I'd go is around 1e-4.

Lastly, in your code snippet, instead of manually computing the Geometry objects, you can directly use:

@jax.tree_util.register_pytree_node_class
class MyCost(ott.geometry.costs.CostFn):
    """Squared Euclidean distance."""

    def pairwise(self, x: jnp.ndarray, y: jnp.ndarray) -> float:
        return ((displacement_fn(x, y)) ** 2).sum(axis=-1)


@jax.jit
def sinkhorn_divergence_ott(x, y):
    out = ott.tools.sinkhorn_divergence.sinkhorn_divergence(
        pointcloud.PointCloud,
        x,
        y,
        cost_fn=MyCost(),
        sinkhorn_kwargs={"threshold": 1e-6},
    )
    return out.divergence, out

[JonasErbesdobler]

Hi @michalk8,

Thanks a lot for your thorough response!

I am writing you on behalf of @arturtoshev since I am currently tasked with utilizing the Sinkhorn divergence as an error measurement for particle distribution.

Using SqEuclidean as a ground cost, it looks like this (between x/y):

As for the ground cost, we also use the squared Euclidean norm, see distance_matrix(x, y) in the original post. We are just using a box with periodic boundaries.

For your use-case, I'd experiment with a different range of values, but definitely wouldn't go as low as 1e-8 as in the code you provided; lowest I'd go is around 1e-4.

Let's say we want to use the Sinkhorn divergence to quantify the error between the particle distribution of two point clouds that are arbitrarily distributed but exactly the same, with the only difference between them being the shift by a vector dr (in our periodic box). For the sake of having a sufficiently accurate error estimation, we figured it would be beneficial to keep epsilon reasonably low. All we actually want to detect is the difference in particle density between the point clouds, i.e., clustering. We were thinking of 1e-6, or would you say that's still unnecessarily low?

Lastly, in your code snippet, instead of manually computing the Geometry objects, you can directly use: [...]

Regarding the use of pointcloud.PointCloud, I have two questions as well.

First, while experimenting with directly using pointcloud.PointCloud and computing the Geometry objects ourselves, we found that the computational time when using PointCloud is considerably larger (around 2x with 2.5k points for 100 executions). Do you have an idea why that is the case?
For computing the Geometry objects, the sinkhorn_divergence_ott(x, y) function is used,

@jax.jit
def sinkhorn_divergence_ott(x, y, eps):
    loss_matrix_xy = distance_matrix(x, y)
    loss_matrix_yy = distance_matrix(y, y)
    loss_matrix_xx = distance_matrix(x, x)
    return ott.tools.sinkhorn_divergence.sinkhorn_divergence(
        ott.geometry.geometry.Geometry,
        loss_matrix_xy,
        loss_matrix_xx,
        loss_matrix_yy,
        # uniform weights
        a=jnp.ones((x.shape[0],)) / x.shape[0],
        b=jnp.ones((y.shape[0],)) / y.shape[0],
        sinkhorn_kwargs={"threshold": 1e-6},
        share_epsilon=True,
        epsilon=eps,
    ).divergence

where the distance matrices are computed from the squared Euclidean norm, and for the PointCloud approach, the following function is defined using the MyCost() ground cost from your comment.

@jax.jit
def sinkhorn_divergence_ott_point(x, y, eps):
    return ott.tools.sinkhorn_divergence.sinkhorn_divergence(
        ott.geometry.pointcloud.PointCloud,
        x,
        y,
        # uniform weights
        a=jnp.ones((x.shape[0],)) / x.shape[0],
        b=jnp.ones((y.shape[0],)) / y.shape[0],
        cost_fn=MyCost(),
        sinkhorn_kwargs={"threshold": 1e-6},
        share_epsilon=True,
        epsilon=eps,
    ).divergence

Second, comparing these two methods, the divergence differs notably when using the same epsilon and said squared Euclidean norm for both. Is there a difference in how the Sinkhorn divergence is computed, changing from Geometry to PointCloud?

Furthermore, since we aim to define a threshold for the divergence to quantify whether a particle distribution is acceptable compared to a reference distribution, we want to normalize the divergence. For this, it must be independent of the number of particles. Considering this and using the squared Euclidean norm as the ground cost, we propose multiplying the divergence by the number of particles. In your opinion, does this make sense? Unfortunately, it seems to work using PointCloud but not Geometry.

Best, Jonas