Repository: pymc-devs/pymc
Pull request: Add a flag to LKJCorr to return the unpacked correlation matrix
## Description
Provide a wrapper around previous LKJCorr that reshapes the output to a full correllation matrix instead of the somewhat awkward vector of upper-triangular values. 
This is based on the discussion with @ricardoV94 and @jessegrabowski  in https://discourse.pymc.io/t/using-lkjcorr-together-with-mvnormal/13606/27 .

## Related Issue
- [x] Closes #7074

## Checklist
- [x] Checked that [the pre-commit linting/style checks pass](https://docs.pymc.io/en/latest/contributing/python_style.html)
- [x] Included tests that prove the fix is effective or that the new feature works
- [ ] Added necessary documentation (docstrings and/or example notebooks)
- [ ] If you are a pro: each commit corresponds to a [relevant logical change](https://wiki.openstack.org/wiki/GitCommitMessages#Structural_split_of_changes)
<!--- You may find this guide helpful: https://mainmatter.com/blog/2021/05/26/keeping-a-clean-git-history/ -->

## Type of change
- [x] New feature / enhancement
- [ ] Bug fix
- [ ] Documentation
- [ ] Maintenance
- [ ] Other (please specify):


<!-- readthedocs-preview pymc start -->
----
📚 Documentation preview 📚: https://pymc--7100.org.readthedocs.build/en/7100/

<!-- readthedocs-preview pymc end -->

Problem statement:
ENH: LKJCorr should return matrix rather than vector
### Before


`LKJCorr` returns an array corresponding to the off-diagonal elements of an upper-triangular matrix. 

```python
> pm.LKJCorr.dist(eta=1, n=3).eval()
array([ 0.06383647, -0.609944  , -0.13745336])
```

### After

`LKJCorr` should return a matrix with ones on the diagonal and the off-diagonal elements in their appropriate positions.

```python
> pm.LKJCorr.dist(eta=1, n=3).eval()
array([[ 1.        ,  0.06383647, -0.609944  ],
       [ 0.        ,  1.        , -0.13745336],
       [ 0.        ,  0.        ,  1.        ]])
```



### Context for the issue:

In the current implementations, the user is left to arrange this into a matrix themselves, which is cumbersome when there are more than one off-diagonal element. The pytensor `fill_diagonal` function only allows for scalar assignment, not array assignment as is the case with NumPy, so it must be filled by hand. Since almost every application will ultimately require the matrix, it should be returned to the user by default. This is also more consistent with `LKJCholeskyCov`, which returns an actual matrix and not just a flat array of values.

Hints from issue discussion:
Could it be An optional kwarg not to break the comparability 
Or, generalize `pytensor.fill_diagonal` to take vector-valued fills, to be consistent with NumPy and make the construction of the matrix easy.
I don't see how `fill_diagonal` helps here. The challenge is in the upper triangular
For the specific case in the example, something like this would work:
```python
import pymc as pm
import pytensor.tensor as pt
import numpy as np

n = 3
x = pm.LKJCorr.dist(eta=1, n=n)
x_diag = pt.ones(n * (n + 1) // 2)
x_diag = pt.set_subtensor(x_diag[[1, 2, 4]], x)
y = pm.math.expand_packed_triangular(n=n, packed=x_diag, lower=False)

pm.draw([x, y])
# [array([-0.06043444,  0.23111108, -0.21084051]),
#  array([[ 1.        , -0.06043444,  0.23111108],
#         [ 0.        ,  1.        , -0.21084051],
#         [ 0.        ,  0.        ,  1.        ]])]
```

However I am not sure if there's a simple rule for the indexes `[1, 2, 4]`, when `n` is larger. Could probably be done with a loop. Here is for `n=4`:
```python
n = 4
x = pm.LKJCorr.dist(eta=1, n=n)
x_diag = pt.ones(n * (n + 1) // 2)
x_diag = pt.set_subtensor(x_diag[[1, 2, 3, 5, 6, 8]], x)
y = pm.math.expand_packed_triangular(n=n, packed=x_diag, lower=False)

pm.draw([x, y])
# [array([ 0.33384753,  0.32512666,  0.62819314, -0.03447542,  0.32310409,
#         -0.1517073 ]),
#  array([[ 1.        ,  0.33384753,  0.32512666,  0.62819314],
#         [ 0.        ,  1.        , -0.03447542,  0.32310409],
#         [ 0.        ,  0.        ,  1.        , -0.1517073 ],
#         [ 0.        ,  0.        ,  0.        ,  1.        ]])]
```
Taking a step back, what kind of matrices do people usually create from the LKJCorr? I don't imagine the ones in the examples above?  Is it symmetric from the LKJCorr with ones in the diagonal?
Something like this?

```python
import pymc as pm
import pytensor.tensor as pt
import numpy as np

n = 3
length = n * (n - 1) // 2
tri_index = np.zeros((n, n), dtype=int)
tri_index[np.triu_indices(n, k=1)] = np.arange(length)
tri_index[np.triu_indices(n, k=1)[::-1]] = np.arange(length)

x = pm.LKJCorr.dist(eta=1, n=n)
y = x[tri_index]
y = pt.fill_diagonal(y, 1)

pm.draw([x, y])
# [array([-0.62563319, -0.75617993,  0.6832039 ]),
#  array([[ 1.        , -0.62563319, -0.75617993],
#         [-0.62563319,  1.        ,  0.6832039 ],
#         [-0.75617993,  0.6832039 ,  1.        ]])]
```
Yes, at least this output would be useful indeed!
We could add it as an optional argument like we do for the LKJCholeskyCov (and if we want to make it a default start adding a warning)
Some useful discussion on this topic took place at:
https://discourse.pymc.io/t/using-lkjcorr-together-with-mvnormal/13606/22

Commit message: LKJCorr now has an option of returning n by n matrix instead of just the upper triangular part as vector

Patch to review:
diff --git a/pymc/distributions/multivariate.py b/pymc/distributions/multivariate.py
index 1e5a9567a4..a1268ba32f 100644
--- a/pymc/distributions/multivariate.py
+++ b/pymc/distributions/multivariate.py
@@ -1505,6 +1505,7 @@ def rng_fn(cls, rng, n, eta, size):
     def _random_corr_matrix(cls, rng, n, eta, flat_size):
         # original implementation in R see:
         # https://github.com/rmcelreath/rethinking/blob/master/R/distributions.r
+
         beta = eta - 1.0 + n / 2.0
         r12 = 2.0 * stats.beta.rvs(a=beta, b=beta, size=flat_size, random_state=rng) - 1.0
         P = np.full((flat_size, n, n), np.eye(n))
@@ -1531,48 +1532,8 @@ def log_jac_det(self, *args):
         return super().log_jac_det(*args).sum(-1)
 
 
-class LKJCorr(BoundedContinuous):
-    r"""
-    The LKJ (Lewandowski, Kurowicka and Joe) log-likelihood.
-
-    The LKJ distribution is a prior distribution for correlation matrices.
-    If eta = 1 this corresponds to the uniform distribution over correlation
-    matrices. For eta -> oo the LKJ prior approaches the identity matrix.
-
-    ========  ==============================================
-    Support   Upper triangular matrix with values in [-1, 1]
-    ========  ==============================================
-
-    Parameters
-    ----------
-    n : tensor_like of int
-        Dimension of the covariance matrix (n > 1).
-    eta : tensor_like of float
-        The shape parameter (eta > 0) of the LKJ distribution. eta = 1
-        implies a uniform distribution of the correlation matrices;
-        larger values put more weight on matrices with few correlations.
-
-    Notes
-    -----
-    This implementation only returns the values of the upper triangular
-    matrix excluding the diagonal. Here is a schematic for n = 5, showing
-    the indexes of the elements::
-
-        [[- 0 1 2 3]
-         [- - 4 5 6]
-         [- - - 7 8]
-         [- - - - 9]
-         [- - - - -]]
-
-
-    References
-    ----------
-    .. [LKJ2009] Lewandowski, D., Kurowicka, D. and Joe, H. (2009).
-        "Generating random correlation matrices based on vines and
-        extended onion method." Journal of multivariate analysis,
-        100(9), pp.1989-2001.
-    """
-
+# Returns list of upper triangular values
+class _LKJCorr(BoundedContinuous):
     rv_op = lkjcorr
 
     @classmethod
@@ -1631,11 +1592,95 @@ def logp(value, n, eta):
         )
 
 
-@_default_transform.register(LKJCorr)
+@_default_transform.register(_LKJCorr)
 def lkjcorr_default_transform(op, rv):
     return MultivariateIntervalTransform(floatX(-1.0), floatX(1.0))
 
 
+# Thin wrapper around _LKJCorr
+class LKJCorr:
+    r"""
+    The LKJ (Lewandowski, Kurowicka and Joe) log-likelihood.
+
+    The LKJ distribution is a prior distribution for correlation matrices.
+    If eta = 1 this corresponds to the uniform distribution over correlation
+    matrices. For eta -> oo the LKJ prior approaches the identity matrix.
+
+    ========  ==============================================
+    Support   Upper triangular matrix with values in [-1, 1]
+    ========  ==============================================
+
+    Parameters
+    ----------
+    n : tensor_like of int
+        Dimension of the covariance matrix (n > 1).
+    eta : tensor_like of float
+        The shape parameter (eta > 0) of the LKJ distribution. eta = 1
+        implies a uniform distribution of the correlation matrices;
+        larger values put more weight on matrices with few correlations.
+    return_matrix : bool, default=False
+        If True, returns the full correllation matrix.
+        False only returns the values of the upper triangular matrix excluding
+        diagonal in a single vector of length n(n-1)/2 for backwards compatibility
+
+    Notes
+    -----
+    This is mainly useful if you want the standard deviations to be fixed, as
+    LKJCholsekyCov is optimized for the case where they come from a distribution.
+
+    Examples
+    --------
+    .. code:: python
+
+        with pm.Model() as model:
+
+            # Define the vector of fixed standard deviations
+            sds = 3*np.ones(10)
+
+            corr = pm.LKJCorr(
+                'corr', eta=4, n=10, return_matrix=True
+            )
+
+            # Define a new MvNormal with the given correllation matrix
+            vals = sds*pm.MvNormal('vals', mu=np.zeros(10), cov=corr, shape=10)
+
+            # Or transform an uncorrelated normal distribution:
+            vals_raw = pm.Normal('vals_raw', shape=10)
+            chol = pt.linalg.cholesky(corr)
+            vals = sds*pt.dot(chol,vals_raw)
+
+
+    References
+    ----------
+    .. [LKJ2009] Lewandowski, D., Kurowicka, D. and Joe, H. (2009).
+        "Generating random correlation matrices based on vines and
+        extended onion method." Journal of multivariate analysis,
+        100(9), pp.1989-2001.
+    """
+
+    def __new__(cls, name, n, eta, *, return_matrix=False, **kwargs):
+        if not return_matrix:
+            return _LKJCorr(name, eta=eta, n=n, **kwargs)
+        else:
+            c_vec = _LKJCorr(name + "_raw", eta=eta, n=n, **kwargs)
+            return pm.Deterministic(name, cls.vec_to_corr_mat(c_vec, n))
+
+    @classmethod
+    def dist(cls, n, eta, *, return_matrix=True, **kwargs):
+        # compute Cholesky decomposition
+        c_vec = _LKJCorr.dist(eta=eta, n=n, **kwargs)
+        if not return_matrix:
+            return c_vec
+        else:
+            return cls.vec_to_corr_mat(c_vec, n)
+
+    @classmethod
+    def vec_to_corr_mat(cls, vec, n):
+        tri = pt.zeros(pt.concatenate([vec.shape[:-1], (n, n)]))
+        tri = pt.subtensor.set_subtensor(tri[(...,) + np.triu_indices(n, 1)], vec)
+        return tri + pt.moveaxis(tri, -2, -1) + pt.diag(pt.ones(n))
+
+
 class MatrixNormalRV(RandomVariable):
     name = "matrixnormal"
     ndim_supp = 2
diff --git a/tests/distributions/test_multivariate.py b/tests/distributions/test_multivariate.py
index 3bee558124..4312874107 100644
--- a/tests/distributions/test_multivariate.py
+++ b/tests/distributions/test_multivariate.py
@@ -34,6 +34,7 @@
 from pymc.distributions.multivariate import (
     MultivariateIntervalTransform,
     _LKJCholeskyCov,
+    _LKJCorr,
     _OrderedMultinomial,
     posdef,
     quaddist_matrix,
@@ -558,7 +559,7 @@ def test_wishart(self, n):
     @pytest.mark.parametrize("x,eta,n,lp", LKJ_CASES)
     def test_lkjcorr(self, x, eta, n, lp):
         with pm.Model() as model:
-            pm.LKJCorr("lkj", eta=eta, n=n, transform=None)
+            pm.LKJCorr("lkj", eta=eta, n=n, transform=None, return_matrix=False)
 
         point = {"lkj": x}
         decimals = select_by_precision(float64=6, float32=4)
@@ -1330,7 +1331,7 @@ def test_kronecker_normal_moment(self, mu, covs, size, expected):
     )
     def test_lkjcorr_moment(self, n, eta, size, expected):
         with pm.Model() as model:
-            pm.LKJCorr("x", n=n, eta=eta, size=size)
+            pm.LKJCorr("x", n=n, eta=eta, size=size, return_matrix=False)
         assert_moment_is_expected(model, expected)
 
     @pytest.mark.parametrize(
@@ -1469,6 +1470,20 @@ def test_with_cov_rv(
 
         assert prior["mv"].shape == (10, 4, 3)
 
+    def test_with_lkjcorr_matrix(
+        self,
+    ):
+        with pm.Model() as model:
+            mu = pm.Normal("mu", 0.0, 1.0, shape=3)
+            sd = pm.Exponential("sd", 1.0, shape=3)
+
+            corr = pm.LKJCorr("corr", n=3, eta=2, return_matrix=True)
+            # pylint: enable=unpacking-non-sequence
+            mv = pm.MvNormal("mv", mu, cov=sd * (sd * corr).T, size=4)
+            prior = pm.sample_prior_predictive(samples=10, return_inferencedata=False)
+
+        assert prior["mv"].shape == (10, 4, 3)
+
     def test_issue_3758(self):
         np.random.seed(42)
         ndim = 50
@@ -2137,7 +2152,7 @@ class TestOrderedMultinomial(BaseTestDistributionRandom):
 
 
 class TestLKJCorr(BaseTestDistributionRandom):
-    pymc_dist = pm.LKJCorr
+    pymc_dist = _LKJCorr
     pymc_dist_params = {"n": 3, "eta": 1.0}
     expected_rv_op_params = {"n": 3, "eta": 1.0}
 
@@ -2165,7 +2180,7 @@ def ref_rand(size, n, eta):
             return (st.beta.rvs(size=(size, shape), a=beta, b=beta) - 0.5) * 2
 
         continuous_random_tester(
-            pm.LKJCorr,
+            _LKJCorr,
             {
                 "n": Domain([2, 10, 50], edges=(None, None)),
                 "eta": Domain([1.0, 10.0, 100.0], edges=(None, None)),
@@ -2190,7 +2205,7 @@ def ref_rand(size, n, eta):
 )
 def test_LKJCorr_default_transform(shape):
     with pm.Model() as m:
-        x = pm.LKJCorr("x", n=2, eta=1, shape=shape)
+        x = pm.LKJCorr("x", n=2, eta=1, shape=shape, return_matrix=False)
     assert isinstance(m.rvs_to_transforms[x], MultivariateIntervalTransform)
     assert m.logp(sum=False)[0].type.shape == shape[:-1]
 


Code review comments:
diff_hunk: @@ -1631,11 +1592,95 @@ def logp(value, n, eta):
         )
 
 
-@_default_transform.register(LKJCorr)
+@_default_transform.register(_LKJCorr)
 def lkjcorr_default_transform(op, rv):
     return MultivariateIntervalTransform(floatX(-1.0), floatX(1.0))
 
 
+# Thin wrapper around _LKJCorr
path: pymc/distributions/multivariate.py
text: @user1:
Comments like this have a tendency to get left behind. I would exclude or add it inside the class definition
```suggestion
```

diff_hunk: @@ -1631,11 +1592,95 @@ def logp(value, n, eta):
         )
 
 
-@_default_transform.register(LKJCorr)
+@_default_transform.register(_LKJCorr)
 def lkjcorr_default_transform(op, rv):
     return MultivariateIntervalTransform(floatX(-1.0), floatX(1.0))
 
 
+# Thin wrapper around _LKJCorr
+class LKJCorr:
+    r"""
+    The LKJ (Lewandowski, Kurowicka and Joe) log-likelihood.
+
+    The LKJ distribution is a prior distribution for correlation matrices.
+    If eta = 1 this corresponds to the uniform distribution over correlation
+    matrices. For eta -> oo the LKJ prior approaches the identity matrix.
+
+    ========  ==============================================
+    Support   Upper triangular matrix with values in [-1, 1]
+    ========  ==============================================
+
+    Parameters
+    ----------
+    n : tensor_like of int
+        Dimension of the covariance matrix (n > 1).
+    eta : tensor_like of float
+        The shape parameter (eta > 0) of the LKJ distribution. eta = 1
+        implies a uniform distribution of the correlation matrices;
+        larger values put more weight on matrices with few correlations.
+    return_matrix : bool, default=False
+        If True, returns the full correllation matrix.
+        False only returns the values of the upper triangular matrix excluding
+        diagonal in a single vector of length n(n-1)/2 for backwards compatibility
+
+    Notes
+    -----
+    This is mainly useful if you want the standard deviations to be fixed, as
+    LKJCholsekyCov is optimized for the case where they come from a distribution.
+
+    Examples
+    --------
+    .. code:: python
+
+        with pm.Model() as model:
+
+            # Define the vector of fixed standard deviations
+            sds = 3*np.ones(10)
+
+            corr = pm.LKJCorr(
+                'corr', eta=4, n=10, return_matrix=True
+            )
+
+            # Define a new MvNormal with the given correllation matrix
+            vals = sds*pm.MvNormal('vals', mu=np.zeros(10), cov=corr, shape=10)
+
+            # Or transform an uncorrelated normal distribution:
+            vals_raw = pm.Normal('vals_raw', shape=10)
+            chol = pt.linalg.cholesky(corr)
+            vals = sds*pt.dot(chol,vals_raw)
+
+
+    References
+    ----------
+    .. [LKJ2009] Lewandowski, D., Kurowicka, D. and Joe, H. (2009).
+        "Generating random correlation matrices based on vines and
+        extended onion method." Journal of multivariate analysis,
+        100(9), pp.1989-2001.
+    """
+
+    def __new__(cls, name, n, eta, *, return_matrix=False, **kwargs):
+        if not return_matrix:
+            return _LKJCorr(name, eta=eta, n=n, **kwargs)
+        else:
+            c_vec = _LKJCorr(name + "_raw", eta=eta, n=n, **kwargs)
+            return pm.Deterministic(name, cls.vec_to_corr_mat(c_vec, n))
path: pymc/distributions/multivariate.py
text: @user1:
Not sure about the name change and wrapping in a Deterministic by default, but not completely against it either

@user1:
The argument for having the packed form is that it saves memory. This actually more than doubles it because we save both the packed form and the dense one. We could leave to the users to wrap in a Deterministic if they need to have it after sampling.

OTOH this is similar to what LKJ does by default now.

@user2:
Maybe check the discussion #6828 about setting internal deterministics/coords?

Reminds me I need to finish that PR