Duplicate Gaussian code

[JTT94]

Duplicate code for estimating Gaussian moments and computing transport maps for Gaussian OT problems
- In tools, Gaussian

ott/ott/tools/gaussian_mixture/gaussian.py

Line 35 in 5657e78

@classmethod

and

ott/ott/tools/gaussian_mixture/scale_tril.py

Line 175 in 5657e78

m = matrix_square_root.sqrtm_only(

- And in ICNN initialisers,

ott/ott/core/icnn.py

Line 137 in 5657e78

def compute_gaussian_map(self, inputs):

,

ott/ott/core/icnn.py

Line 161 in 5657e78

mo = sqrtm_only(jnp.dot(jnp.dot(covs_sqrt, covt), covs_sqrt))

It may be worth consolidating these. It may also be worth moving Gaussian computations to /core, or creating a hierarchical structure that does not cause cyclical import errors (see e.g. #98 (comment))

[marcocuturi]

thanks for spotting these James, and, yes, I agree, there's too many sqrtm_only in our code :)

if i can ask another comment, i think the way we store gaussians is not ideal, or requires some debate. When doing gaussian OT, we might have to use $A^{0.5}$, $A$ or $A^{-0.5}$. When computing $W_2^2$, either the first or the second are needed (if computing $A$ to $B$ then $A^{0.5}$ is used, the other way around $A$ is used) whereas to compute the map one needs, similarly, either (1,3) or (2).

I am wondering if, as a preprocessing, we should not store the 3 of them. I know we can recover (2) from (1*1) but that sounds a bit like overkill. so maybe storing the 3 triangular parts is not too complicated? and then redefining the Bures cost (or creating a new one, e.g. BuresTril) could make sense to save memory and compute.