FR: adding conditional OT neural solvers

[cspollard]

First off: I really like this library. I was able to get it running rather quickly and test on a few toy low-dimensional examples that close nicely using the provided neural solvers.

The main reason I can't switch to using it "wholesale" is the lack of support of conditional OT solvers along the lines of

https://arxiv.org/abs/2206.14262.

Is that on the road map? Or is it implemented somewhere that I just haven't found yet?

[michalk8]

Hi @cspollard , thanks a lot for the interest! I think this would be a really nice contributions and afaik @MUCDK is/was working on an implementation of PICNNs.

[MUCDK]

Hi @cspollard and @michalk8 ,

We have an implementation in moscot (theislab/moscot#567), including all output classes etc. I think it's up to you @michalk8 and @marcocuturi whether you would like to have this in OTT.

If we implement it, I think we should refactor a bit on the OTT-side, i.e. introduce more abstract classes for the output/cond. output classes, especially given that it's not unlikely we'll add another neural solver to OTT soon.

We could use our moscot implementation as a baseline and start from there.

[cspollard]

Thanks for the quick feedback/action!

I think you'd find a number of users from collider physics interested in this. So far we have rolled our own implementations of the PICNNs, but this is nontrivial to get right.

[marcocuturi]

Hi @cspollard and thanks for your interest in OTT!

I think it's a good idea to add PICNNs to OTT. @MUCDK @michalk8 we should discuss this!

[jannisborn]

@cspollard @michalk8 what's the state on this issue? We'd also be interested to see PICNNs supported in ott-jax. Could you describe what would be required beyond building a PICNN child class of ICNNs? The initialization scheme proposed in the CondOT paper seems to be already available.

[cspollard]

Hi @jannisborn. I’m not sure. From previous messages I understood this was not a big task. If it requires a lot of development, maybe we could “chip in” to help. I let @marcocuturi @michalk8 @MUCDK comment.

Chris