dim>1 for ot.bregman.barycenter

[gabrieldernbach]

🚀 Feature

Extend ot.bregman.barycenter to higher dimensions (>1D data)

Motivation

Computing barycenters for color images requires at least 3 dimensions (rgb or lab).

Pitch

Recent papers have suggested this for data augmentation

[rflamary]

Hello @gabrieldernbach,

Note that the current solver ot.bregman.barycenter works in any dimensions since you provide the solver with a distance matrix that can be computed in any dimensions (it won't scale to millions of pixels though).

We also have solvers in 2d for regular grids here:
https://pythonot.github.io/auto_examples/barycenters/plot_convolutional_barycenter.html#sphx-glr-auto-examples-barycenters-plot-convolutional-barycenter-py

And the solver with free support works also very well in any dimension:
https://pythonot.github.io/auto_examples/barycenters/plot_free_support_barycenter.html#sphx-glr-auto-examples-barycenters-plot-free-support-barycenter-py

I agree that some higher dimesiond solvers for grid (3d color histograms) would be interesting and we will look into implementing it in the future.

[gabrieldernbach]

Dear rflamary thank you a lot for your quick reply. I am happy to hear that this is already possible.
I had a try, but could not figure the right input shapes. May I suggest to add a note to the docs?

In the documentation it says that

A (ndarray, shape (dim, n_hists)) – n_hists training distributions a_i of size dim
M (ndarray, shape (dim, dim)) – loss matrix for OT

Lets say I have 128 pixel in 3 dimension (lab).
Then this runs through, as expected

A = np.random.rand(128, 3)
M = np.random.rand(128, 128)
ot.bregman.barycenter_sinkhorn(A, M, reg=0.01)

If I now have 5 images for which I wish to find the barycenters

A = np.random.rand(5, 128, 3)
M = np.random.rand(5, 128, 128) # also tried (128, 128)
ot.bregman.barycenter_sinkhorn(A, M, reg=0.01)

results in a dimension misalignment error. I tried some permutations but couldn't find the appropriate version.

Following the tutorial example, where the arguments are lists, I also tried

A = [np.random.rand(128, 3) for _ in range(3)]
M = np.random.rand(128, 128)
ot.bregman.barycenter_sinkhorn(A, M, reg=0.01)

which fails by the fact that A is supposed to have a shape property.

Thanks for your help.

[wershofe]

This is something I am also struggling with. I would prefer to use ot.bregman.barycenter over the free_support version because then I can specify the cost matrix.
Thanks