search

ericphanson / UnbalancedOptimalTransport.jl

Sinkhorn divergences for measures of unequal mass
Other
14 stars 1 forks source link

Support precomputed cost matrices #5

Closed baggepinnen closed 4 years ago

baggepinnen commented 4 years ago

This PR allows the user to supply either a cost function like before, or a precomputed cost matrix. If the matrix is supplied, the support points of the measures are not required. This PR is non-breaking, but it introduces some extra allocations in case the user does not supply a precomputed cost matrix. Despite the extra allocations, this PR should be a net positive for speed since it saves a lot of computations. Addresses #4

baggepinnen commented 4 years ago

It appears that the number of allocations is very unstable between julia versions. I upped the threshold to 64 bytes from the previous 32, but on nightly there are some 240 bytes allocated.

ericphanson commented 4 years ago

Ah ok, maybe it’s not such a good idea to test allocations. Some of them could be dynamic dispatches since Julia doesn’t dispatch on keyword arguments; what do you think of making the cost the 2nd argument instead of a keyword argument? Could add an extra method to provide the default value of norm, and replace handle_C with just a dispatch to precompute_cost for a non-AbstractMatrix method.

It’s breaking but seems like it might be a better interface.

Edit: and maybe rename precompute_cost to compute_cost_matrix? I’m thinking if we use it in the tutorial it’ll be more obvious what it’s for with that name.

baggepinnen commented 4 years ago

C is now the second positional argument, and one can provide either a precomputed matrix, or a cost function. If nothing is provided, norm(x-y) is used. sinkhorn_divergence is an outlier here in that it does currently not support a precomputed cost matrix, since it requires three different matrices.

I also loosened the allocation threshold so that it accepts some wiggle room for whatever happens on julia#master

baggepinnen commented 4 years ago

Great! I used github's feature to apply all your suggested changes. I've never used that before so I hope I didn't crew up :P

Good catch with the potential silent failure, hadn't considered that.

ericphanson commented 4 years ago

I haven’t used it much either but looks like it worked! Mind adding a test for the new error path? I’m not sure why we aren’t getting a codecov comment but I’ve never had 100% coverage before this repo and would like to keep it 😅.

ericphanson commented 4 years ago

Awesome, thanks! I'll tag a release shortly. edit: https://github.com/JuliaRegistries/General/pull/12940

baggepinnen commented 4 years ago

Awesome, I'm already using this in https://github.com/baggepinnen/SpectralDistances.jl 😍

ericphanson commented 4 years ago

Cool! That looks like a really nice package. Please feel free to open issues if you run into any problems down the line!