bug fix: avoid mixing up linear and quadratic in genot

[soerenab]

This fixes #514.

prepare_data() now returns a dict with the right keys for the linear or quadratic matching function.

[michalk8]

@soerenab could you please fix the failing tests?

2024-04-09T14:13:11.3993310Z FAILED tests/neural/methods/genot_test.py::TestGENOT::test_genot[lin_dl] - TypeError: data_match_fn() got an unexpected keyword argument 'x'
2024-04-09T14:13:11.3995788Z FAILED tests/neural/methods/genot_test.py::TestGENOT::test_genot[quad_dl] - TypeError: data_match_fn() got an unexpected keyword argument 'xx'
2024-04-09T14:13:11.3998067Z FAILED tests/neural/methods/genot_test.py::TestGENOT::test_genot[fused_dl] - TypeError: data_match_fn() got an unexpected keyword argument 'x'
2024-04-09T14:13:11.4000374Z FAILED tests/neural/methods/genot_test.py::TestGENOT::test_genot[lin_cond_dl] - TypeError: data_match_fn() got an unexpected keyword argument 'x'
2024-04-09T14:13:11.4002702Z FAILED tests/neural/methods/genot_test.py::TestGENOT::test_genot[quad_cond_dl] - TypeError: data_match_fn() got an unexpected keyword argument 'xx'
2024-04-09T14:13:11.4005049Z FAILED tests/neural/methods/genot_test.py::TestGENOT::test_genot[fused_cond_dl] - TypeError: data_match_fn() got an unexpected keyword argument 'x'

[michalk8]

Afterwards, LGTM and can merge, the failing
2024-04-09T14:13:12.4681618Z FAILED tests/tools/soft_sort_test.py::TestSoftSort::test_soft_sort_jacobian[True] - jaxlib.xla_extension.XlaRuntimeError: INTERNAL: Generated function failed: CpuCallback error: Traceback (most recent call last): is unrelated and I need to reproduce locally.

[soerenab]

The above mentioned tests now pass. Overall, this PR makes the following changes:

• prepare_data() in GENOT no longer returns src_quad and tgt_quad if these are None (as the linear matching function does not take them as arguments)
• the order of arguments in match_quadratic() has changed from xx, yy, x, y to x, y, xx, yy to be consistent with match_linear()
• In genot I replaced data_match_fn() with get_match_function(). The previous combination of Optional arguments and functools.partial does not seem to work -> I always got errors of missing args in the linear case even though those args were not needed and indeed optional. This new implementation is also cleaner in my opinion.

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 90.92%. Comparing base (4bba69f) to head (42b05a9).
Report is 3 commits behind head on main.

❗ Current head 42b05a9 differs from pull request most recent head 2bbf904. Consider uploading reports for the commit 2bbf904 to get more accurate results

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #517      +/-   ##
==========================================
+ Coverage   90.62%   90.92%   +0.30%     
==========================================
  Files          68       68              
  Lines        7047     7044       -3     
  Branches      685      998     +313     
==========================================
+ Hits         6386     6405      +19     
+ Misses        501      485      -16     
+ Partials      160      154       -6     

Files	Coverage Δ
src/ott/neural/methods/flows/genot.py	84.61% <100.00%> (+0.13%)	⬆️
src/ott/solvers/utils.py	90.69% <ø> (ø)

... and 1 file with indirect coverage changes