Add Sparse Geodesic Cost Support

[selmanozleyen]

Here are the steps I planned. I first want to ensure the tests work with this version and I am currently on that step.

• make sure it works on the new commit of ottjax (here: ott-jax/ott@51a658a)
• don't densify sparse arrays for geodesic costs
• adjust and add tests

In order to comply to new version:

• created new solution files because loaded objects in test_regression_testing had corrupted format.
• No convergence with new commit of ottjax in LRGW solvers unless inner_iterations=10 is set #678
• compute_sparse_laplacian gives int32 vs int64 index mismatch when input is from scipy csr ott-jax/ott#509
• Set defaults of min_iterations and max_iterations to None #686
• Fixing ottjax 0.4.6 Errors #688 (already included in this PR)
  ping: @MUCDK , @giovp
• Unbalanced FGW doesn't converge when margins are provided ott-jax/ott#519

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 78.19%. Comparing base (3151da7) to head (f3c54be).
Report is 25 commits behind head on main.

❗ Current head f3c54be differs from pull request most recent head b054584

Please upload reports for the commit b054584 to get more accurate results.

Additional details and impacted files

@@           Coverage Diff           @@
##             main     #677   +/-   ##
=======================================
  Coverage   78.19%   78.19%           
=======================================
  Files          36       36           
  Lines        3806     3806           
  Branches      706      706           
=======================================
  Hits         2976     2976           
  Misses        524      524           
  Partials      306      306           

see 2 files with indirect coverage changes

[MUCDK]

in the failing tests, set tau_a, tau_b > 0.5

[giovp]

• try increase the taus for the failing tests for FGW/GW problems
• default to None the n iterations

[selmanozleyen]

I tried all of your suggestions and it only worked when I set the tau's to be 1.0 which is balanced. I checked again and it fails after this ott-jax/ott@41906a2 commit specifically. I created a PR separately for the setting the defaults to None #686

[giovp]

ok then i think we need to raise this to @michalk8 because it is not expected

[MUCDK]

agree, can you maybe try to run the example without moscot, but directly with ott-jax @selmanozleyen ? This might help @michalk8

[selmanozleyen]

agree, can you maybe try to run the example without moscot, but directly with ott-jax @selmanozleyen ? This might help @michalk8

@MUCDK ok I created it here ott-jax/ott#519

[MUCDK]

Can you please also update the tests with geodesic costs for other problems, e.g.

moscot/tests/problems/time/test_temporal_problem.py

Line 252 in 6cc2215

def test_geodesic_cost_set_xy_cost_dense(self, adata_time):

Also, accordingly, please update

moscot/src/moscot/utils/tagged_array.py

Line 106 in 6cc2215

Sparse arrays will be always densified.