Batch OT sampling of conditional PDFs

[CaioDaumann]

Dear experts,

Thank you very much for this nice package! I am trying to use your batch OT version of flow matching to perform the matching between two distributions, but I wonder how the OT mapping should work for conditional distributions.

Suppose I have a function y(x|c) and I want to transport it to z(x|c). When I sample n_events from these distributions, they will have different c values among themselves. How can one perform OT mapping between "events" with different condition values?

I thought of something like this: we split the data into bins of c [0, 0.1, 0.2, ...] where we don't expect the PDFs to change much, and during training, sample only events of y(x|c) and z(x|c) where the events of both are in the same bin.

Does that sound reasonable to you? Or is there a clever and simpler way to do it?

Best, Caio