(feat): Linear GENOT #662

ilan-gold · 2024-02-26T14:24:42Z

We add two main new classes here:

GENOTLinSolver wraps the underlying OTT-Jax solver in a moscot-usable way
GENOTLinProblem at the next level of abstraction up, allowing users to work with AnnData objects directly while also using other parts of the moscot API like policies

To run this PR, simply do pip install -e . as it points directly at the corresponding OTT-JAX branch.

TODOs:

Finalize status with OTT-Jax PR
Remove more dead code as this was based on the ICNN implementation (if there is any)
Maybe add more tests for certain things that I did not cover? I will need to look this over
Maybe a small notebook?
Figure out how to squash all these commits. It should be clear where I picked up and the ICNN branch stopped but in any case...would be nice on the eyes. And we definitely want to squash this PR when merged

* rename tag cost to cost_matrix * fix renaming * [CI skip], adapt callback * incorporate requested changes * add test for quad custom callback * adapt kwargs for callback * fix handle_joint_attr * incorporate requested changes

* fix test in FGWProblem * add correlation test * add first test for correlation computation * add more tests * fix tests * add tfs to compute_feature_correlation * add testing for no nans in compute_feature_correlation * incorporate requested changes * fix docstring

* fix sankey return statement * adapt test * adapt return_fig

* fix return type in mpl * change import acronyms * fix tests

* Simplify linear operator * Simplify `align`, fix test

* fix return type in mpl * change import acronyms * fix tests * add interpolation option to sankey * add test to interpolate color * define colors for pull/push` * adapt tests * introduce axes in mpl.push/pull * incorporate requested changes * change default color * adapt plotting * introduce scaling * fix scale * make start/end categorical in plot * regenerate images

* Remove `FGWSolver` * Fix `tox.ini` * Fix wrong shape check * Use pure GW in generic solver * Update tests

* fix bug in SinkhornProblem * fix tox.ini

* make push/pull always use source/target * fix bug in StarPolicy _apply * adapt plotting to source/target

* fix strip plotting in sankey * simplify code

* add spearman correlation * add tests * adapt tests

* make push/pull plot in good order * [CI skip], try setting adata.uns color explicitly * [CI skip], fix copying of adata * fix pre commits * fix bug

…pe of `temporal_key` (#449) * make marginal_kwargs explicit in temporal problems * introduce check for numeric dtype in temporal mixin * add alternative way for marginal prior * adapt tolerances in tests * correct docs * fix bug * Fix math rendering * fix test Co-authored-by: Michal Klein <[email protected]>

MUCDK

Thanks @ilan-gold .
Maybe let's do the few changes I suggested, and then next week I'll go over it with @selmanozleyen as well.

ilan-gold · 2024-06-19T08:46:26Z

ok @MUCDK !

…g/neural_solvers

selmanozleyen · 2024-06-19T17:09:15Z

@MUCDK do you think GENOT code in ottjax is final? My version in my thesis project is very different from the one in ottjax . If you like maybe we can first discuss possible changes there. Depends on what you want to prioritize.

MUCDK · 2024-06-19T19:32:23Z

@MUCDK do you think GENOT code in ottjax is final? My version in my thesis project is very different from the one in ottjax . If you like maybe we can first discuss possible changes there. Depends on what you want to prioritize.

Yeah it's not super final but I wouuld prefer to have the merges more incremental as we have been workin gon this PR for almost 2 years.

Let's talk next week in person.

selmanozleyen · 2024-07-23T07:06:40Z

tests/problems/generic/test_gap_problems.py

Is this file left empty intentionally?

maybe a TODO comment instead of leaving empty might be better

we can delete this file

selmanozleyen · 2024-07-23T07:08:24Z

src/moscot/problems/time/__init__.py

I know this might be nitpicking but I am just curious if there is a reason for this change

src/moscot/problems/generic/_mixins.py

selmanozleyen · 2024-07-23T07:12:23Z

docs/conf.py

+    # see: https://github.com/numpy/numpydoc/issues/275
+    ("py:class", "None.  Remove all items from D."),
+    ("py:class", "a set-like object providing a view on D's items"),
+    ("py:class", "a set-like object providing a view on D's keys"),
+    ("py:class", "v, remove specified key and return the corresponding value."),  # noqa: E501
+    ("py:class", "None.  Update D from dict/iterable E and F."),
+    ("py:class", "an object providing a view on D's values"),
+    ("py:class", "a shallow copy of D"),


I am not sure what these mean. Is this intentional or was it left from debugging

I think these are warnings from the RTD that should be ignored base on the static class rendering file. But they are indeed too many I'm not sure if it's wanted.

Not sure, what to do here. We can't import from MutableMapping for DistributionCollection without re-implementing all the abstract methods, which seems like a waste. The other option would be to remove the DistributionCollection class since it is used only for repr/cleanness of API. @MUCDK this code was pulled from your starting point, so perhaps you can chime in.

I'd be fine with keeping it as it is if it's not too much of an issue. @selmanozleyen wdyt?

If typing.Dict doesn't work as well I guess using those ignores is a better solution than re-implementation.

selmanozleyen

Thanks @ilan-gold I made some comments. If it passes the relevant tests and docs seem good that's enough for me. But one very important thing is the conflict in the files.

MUCDK · 2024-07-30T07:30:48Z

Greatl, let's wait for @ilan-gold 's response, and then merge.

selmanozleyen

looks good

MUCDK and others added 30 commits March 15, 2023 15:07

fix jaxsampler

f86eb0e

fix jaxsampler

ac3c781

fix jaxsampler

385bcfd

fix tests

cfebc4d

add plot_convergence

e056042

remove jit from _compute_unbalanced marginals

a7af86b

fix sinkhorn_divergence

c8cbe9d

adapt tox.ini file

ff145d1

shape mismatch fixed without precommit

9e241db

remove print statement

3cc279f

finish merge

c10c5b6

adapt callbacks and rename tag cost to cost_matrix (#426)

74adb1e

* rename tag cost to cost_matrix * fix renaming * [CI skip], adapt callback * incorporate requested changes * add test for quad custom callback * adapt kwargs for callback * fix handle_joint_attr * incorporate requested changes

fix sankey return statement (#428)

7360819

* fix sankey return statement * adapt test * adapt return_fig

Bump version: 0.1.0 → 0.1.1

e08f63f

fix return statements

728537c

add save tests

f2153e9

fix return type in mpl (#432)

4f1d26e

* fix return type in mpl * change import acronyms * fix tests

Simplify linear operator (#431)

0ffc730

* Simplify linear operator * Simplify `align`, fix test

Explicitly jit the solvers (#433)

cee2466

Remove FGWSolver (#437)

e28816f

* Remove `FGWSolver` * Fix `tox.ini` * Fix wrong shape check * Use pure GW in generic solver * Update tests

fix bug in SinkhornProblem (#442)

3ff63be

* fix bug in SinkhornProblem * fix tox.ini

fix pre commits

865ad5a

make push/pull always use source/target (#443)

2571935

* make push/pull always use source/target * fix bug in StarPolicy _apply * adapt plotting to source/target

fix strip plotting in sankey (#445)

3d22677

* fix strip plotting in sankey * simplify code

Feature/spearman correlation (#444)

b5370e0

* add spearman correlation * add tests * adapt tests

Delete logo.png

34ef904

Feature/plot order (#453)

9806854

* make push/pull plot in good order * [CI skip], try setting adata.uns color explicitly * [CI skip], fix copying of adata * fix pre commits * fix bug

MUCDK requested changes Jun 19, 2024

View reviewed changes

ilan-gold added 6 commits June 19, 2024 14:14

(chore): makerecall_target and aggregate_to_topk

46e5fc2

(chore): fix default arguments

bcd288c

(chore): project_transport_matrix -> project_to_transport_matrix

822628d

(fix): remove dead NeuralAnalysisMixin code

5835390

Merge branch 'ig/neural_solvers' of github.com:theislab/moscot into i…

d245cb0

…g/neural_solvers

Merge branch 'main' into ig/neural_solvers

90e9fc6

(feat): allow custom data_match_fn

e9cd90e

selmanozleyen reviewed Jul 23, 2024

View reviewed changes

src/moscot/problems/generic/_mixins.py Show resolved Hide resolved

selmanozleyen reviewed Jul 23, 2024

View reviewed changes

selmanozleyen marked this pull request as ready for review July 23, 2024 07:26

ilan-gold added 10 commits August 5, 2024 17:02

(fix): inherit from MutableMapping instead of dict

cf11513

Merge branch 'main' into ig/neural_solvers

0f91c4e

(Fix): docs

ab51550

(fix): notebooks

2ce5c0b

(fix): docs reference

9b38b65

(fix): remove attr

60e074c

(fix): erroneous change

c8a9cbd

(fix): remove empty

b97a2b4

(fix): notebooks again?

d3b31ec

(chore): ok?

2678705

selmanozleyen approved these changes Aug 20, 2024

View reviewed changes

MUCDK merged commit 6e65ac1 into main Aug 21, 2024
8 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

(feat): Linear GENOT #662

(feat): Linear GENOT #662

ilan-gold commented Feb 26, 2024 •

edited

Loading

MUCDK left a comment

ilan-gold commented Jun 19, 2024

selmanozleyen commented Jun 19, 2024

MUCDK commented Jun 19, 2024

selmanozleyen Jul 23, 2024

selmanozleyen Aug 20, 2024

MUCDK Aug 20, 2024

selmanozleyen Jul 23, 2024

selmanozleyen Jul 23, 2024

giovp Jul 23, 2024

ilan-gold Aug 5, 2024 •

edited

Loading

MUCDK Aug 7, 2024

selmanozleyen Aug 20, 2024

selmanozleyen left a comment •

edited

Loading

MUCDK commented Jul 30, 2024

selmanozleyen left a comment

(feat): Linear GENOT #662

(feat): Linear GENOT #662

Conversation

ilan-gold commented Feb 26, 2024 • edited Loading

MUCDK left a comment

Choose a reason for hiding this comment

ilan-gold commented Jun 19, 2024

selmanozleyen commented Jun 19, 2024

MUCDK commented Jun 19, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ilan-gold Aug 5, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

selmanozleyen left a comment • edited Loading

Choose a reason for hiding this comment

MUCDK commented Jul 30, 2024

selmanozleyen left a comment

Choose a reason for hiding this comment

ilan-gold commented Feb 26, 2024 •

edited

Loading

ilan-gold Aug 5, 2024 •

edited

Loading

selmanozleyen left a comment •

edited

Loading