Fix issues when providing diagonal matrix sqrt

metastable_projections/misc/distTools.py

@@ -27,20 +27,23 @@ def get_mean_and_chol(p: AnyDistribution, expand=False):
 
 
 def get_mean_and_sqrt(p: UniversalGaussianDistribution, expand=False):
+    mean, chol = get_mean_and_chol(p, expand=False)
     if not hasattr(p, 'cov_sqrt'):
-        raise Exception(
-            'Distribution was not induced from sqrt. On-demand calculation is not supported.')
+        if len(p.distribution.scale.shape)==2: # In the factorized case, chol = matSqrt
+            sqrt_cov = p.distribution.scale
+        else:
+            raise Exception(
+                'Distribution was not induced from sqrt. On-demand calculation is not supported.')
     else:
-        mean, chol = get_mean_and_chol(p, expand=False)
         sqrt_cov = p.cov_sqrt
-        if mean.shape[0] != sqrt_cov.shape[0]:
-            shape = list(sqrt_cov.shape)
-            shape[0] = mean.shape[0]
-            shape = tuple(shape)
-            sqrt_cov = sqrt_cov.expand(shape)
-        if expand and len(sqrt_cov.shape) <= 2:
-            sqrt_cov = th.diag_embed(sqrt_cov)
-        return mean, sqrt_cov
+    if mean.shape[0] != sqrt_cov.shape[0]:
+        shape = list(sqrt_cov.shape)
+        shape[0] = mean.shape[0]
+        shape = tuple(shape)
+        sqrt_cov = sqrt_cov.expand(shape)
+    if expand and len(sqrt_cov.shape) <= 2:
+        sqrt_cov = th.diag_embed(sqrt_cov)
+    return mean, sqrt_cov
 
 
 def get_cov(p: AnyDistribution):
@@ -123,12 +126,14 @@ def new_dist_like_from_sqrt(orig_p: AnyDistribution, mean: th.Tensor, cov_sqrt:
 
 
 def _sqrt_to_chol(cov_sqrt, only_diag=False):
+    if only_diag:
+        return cov_sqrt
     cov = th.bmm(cov_sqrt.mT, cov_sqrt)
     cov += th.eye(cov.shape[-1]).expand(cov.shape)*(1e-6)
 
     chol = th.linalg.cholesky(cov)
 
-    if only_diag:
-        chol = th.diagonal(chol, dim1=-2, dim2=-1)
+    #if only_diag:
+    #    chol = th.diagonal(chol, dim1=-2, dim2=-1)
 
     return chol