Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

(feat): Linear GENOT #662

Merged
merged 334 commits into from
Aug 21, 2024
Merged

(feat): Linear GENOT #662

merged 334 commits into from
Aug 21, 2024

Conversation

ilan-gold
Copy link
Collaborator

@ilan-gold ilan-gold commented Feb 26, 2024

We add two main new classes here:

  1. GENOTLinSolver wraps the underlying OTT-Jax solver in a moscot-usable way
  2. 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

MUCDK and others added 30 commits March 15, 2023 15:07
* 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 <46717574+michalk8@users.noreply.github.com>
Copy link
Collaborator

@MUCDK MUCDK left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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
Copy link
Collaborator Author

ok @MUCDK !

@selmanozleyen
Copy link
Collaborator

@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
Copy link
Collaborator

MUCDK commented Jun 19, 2024

@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.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is this file left empty intentionally?

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

maybe a TODO comment instead of leaving empty might be better

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

we can delete this file

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Comment on lines +66 to +73
# 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"),
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Collaborator Author

@ilan-gold ilan-gold Aug 5, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Collaborator

@selmanozleyen selmanozleyen left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@selmanozleyen selmanozleyen marked this pull request as ready for review July 23, 2024 07:26
@MUCDK
Copy link
Collaborator

MUCDK commented Jul 30, 2024

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

Copy link
Collaborator

@selmanozleyen selmanozleyen left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

looks good

@MUCDK MUCDK merged commit 6e65ac1 into main Aug 21, 2024
8 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

8 participants