Add a flag to LKJCorr to return the unpacked correlation matrix #7100

velochy · 2024-01-11T19:37:03Z

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

Closes ENH: LKJCorr should return matrix rather than vector #7074

Checklist

Checked that the pre-commit linting/style checks pass
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

Type of change

📚 Documentation preview 📚: https://pymc--7100.org.readthedocs.build/en/7100/

…the upper triangular part as vector

welcome · 2024-01-11T19:37:06Z

💖 Thanks for opening this pull request! 💖 The PyMC community really appreciates your time and effort to contribute to the project. Please make sure you have read our Contributing Guidelines and filled in our pull request template to the best of your ability.

ricardoV94 · 2024-01-11T20:43:32Z

tests/distributions/test_multivariate.py

+            sd = pm.Exponential("sd", 1.0, shape=3)
+
+            corr = pm.LKJCorr("corr", n=3, eta=2, return_matrix=True)
+            # pylint: enable=unpacking-non-sequence


Suggested change

# pylint: enable=unpacking-non-sequence

ricardoV94 · 2024-01-11T20:44:13Z

tests/distributions/test_multivariate.py

+            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)


I would like to directly assert the values of corr are correct

ricardoV94 · 2024-01-11T20:47:44Z

pymc/distributions/multivariate.py

+            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))


Not sure about the name change and wrapping in a Deterministic by default, but not completely against it either

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.

Maybe check the discussion #6828 about setting internal deterministics/coords?

Reminds me I need to finish that PR

ricardoV94 · 2024-01-11T20:48:33Z

pymc/distributions/multivariate.py

 def lkjcorr_default_transform(op, rv):
    return MultivariateIntervalTransform(floatX(-1.0), floatX(1.0))


+# Thin wrapper around _LKJCorr


Comments like this have a tendency to get left behind. I would exclude or add it inside the class definition

Suggested change

# Thin wrapper around _LKJCorr

ricardoV94 · 2024-01-11T20:51:12Z

Thanks for picking this up @velochy!

codecov · 2024-01-11T20:59:47Z

Codecov Report

Attention: 4 lines in your changes are missing coverage. Please review.

Comparison is base (7bb2ccd) 92.21% compared to head (ab63678) 92.21%.
Report is 3 commits behind head on main.

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #7100   +/-   ##
=======================================
  Coverage   92.21%   92.21%           
=======================================
  Files         101      101           
  Lines       16912    16901   -11     
=======================================
- Hits        15595    15586    -9     
+ Misses       1317     1315    -2

Files	Coverage Δ
pymc/distributions/multivariate.py	`93.80% <77.77%> (+0.26%)`	⬆️

... and 52 files with indirect coverage changes

twiecki · 2024-01-12T02:39:08Z

pymc/distributions/multivariate.py

+        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.


Suggested change

If True, returns the full correllation matrix.

If True, returns the full correlation matrix.

velochy · 2024-01-12T06:39:43Z

Made changes corresponding to all the comments so far.

ricardoV94

Looks pretty good! Left some minor suggestions

ricardoV94 · 2024-01-12T07:01:01Z

pymc/distributions/multivariate.py

+
+    @classmethod
+    def dist(cls, n, eta, *, return_matrix=True, **kwargs):
+        # compute Cholesky decomposition


Suggested change

# compute Cholesky decomposition

ricardoV94 · 2024-01-12T07:01:27Z

pymc/distributions/multivariate.py

+
+    def __new__(cls, name, n, eta, *, return_matrix=False, **kwargs):
+        if not return_matrix:
+            return _LKJCorr(name, eta=eta, n=n, **kwargs)


Can just call once and then branch in the return now that they're the same

ricardoV94 · 2024-01-12T07:03:08Z

pymc/distributions/multivariate.py

+    return_matrix : bool, default=False
+        If True, returns the full correlation 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


Suggested change

diagonal in a single vector of length n(n-1)/2 for backwards compatibility

diagonal in a single vector of length n(n-1)/2 for memory efficiency

ricardoV94

Thanks, this looks like a great improvement in UX

welcome · 2024-01-23T14:10:46Z

]
Congrats on merging your first pull request! 🎉 We here at PyMC are proud of you! 💖 Thank you so much for your contribution 🎁

LKJCorr now has an option of returning n by n matrix instead of just …

1129cf8

…the upper triangular part as vector

ricardoV94 reviewed Jan 11, 2024

View reviewed changes

ricardoV94 added the enhancements label Jan 11, 2024

ricardoV94 changed the title ~~Add a flag to LKJCorr to let it return fully formed correllation matrix~~ Add a flag to LKJCorr to let it return fully formed correlation matrix Jan 11, 2024

ricardoV94 changed the title ~~Add a flag to LKJCorr to let it return fully formed correlation matrix~~ Add a flag to LKJCorr to let it return unpacked correlation matrix Jan 11, 2024

ricardoV94 changed the title ~~Add a flag to LKJCorr to let it return unpacked correlation matrix~~ Add a flag to LKJCorr to return unpacked correlation matrix Jan 11, 2024

ricardoV94 changed the title ~~Add a flag to LKJCorr to return unpacked correlation matrix~~ Add a flag to LKJCorr to return the unpacked correlation matrix Jan 11, 2024

twiecki reviewed Jan 12, 2024

View reviewed changes

velochy force-pushed the lkj_corr_return_matrix branch from 9a64569 to 98a7f48 Compare January 12, 2024 06:38

ricardoV94 reviewed Jan 12, 2024

View reviewed changes

Fixes to PR according to comments

ab63678

velochy force-pushed the lkj_corr_return_matrix branch from 98a7f48 to ab63678 Compare January 12, 2024 10:43

ricardoV94 approved these changes Jan 12, 2024

View reviewed changes

ricardoV94 merged commit ef1f91f into pymc-devs:main Jan 23, 2024
22 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add a flag to LKJCorr to return the unpacked correlation matrix #7100

Add a flag to LKJCorr to return the unpacked correlation matrix #7100

velochy commented Jan 11, 2024 •

edited by ricardoV94

Loading

welcome bot commented Jan 11, 2024

ricardoV94 Jan 11, 2024

ricardoV94 Jan 11, 2024

ricardoV94 Jan 11, 2024 •

edited

Loading

ricardoV94 Jan 11, 2024 •

edited

Loading

jessegrabowski Jan 11, 2024

ricardoV94 Jan 11, 2024

ricardoV94 commented Jan 11, 2024

codecov bot commented Jan 11, 2024 •

edited

Loading

twiecki Jan 12, 2024

velochy commented Jan 12, 2024

ricardoV94 left a comment

ricardoV94 Jan 12, 2024

ricardoV94 Jan 12, 2024

ricardoV94 Jan 12, 2024

ricardoV94 left a comment

welcome bot commented Jan 23, 2024

	If True, returns the full correllation matrix.
	If True, returns the full correlation matrix.

	diagonal in a single vector of length n(n-1)/2 for backwards compatibility
	diagonal in a single vector of length n(n-1)/2 for memory efficiency

Add a flag to LKJCorr to return the unpacked correlation matrix #7100

Add a flag to LKJCorr to return the unpacked correlation matrix #7100

Conversation

velochy commented Jan 11, 2024 • edited by ricardoV94 Loading

Description

Related Issue

Checklist

Type of change

welcome bot commented Jan 11, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ricardoV94 Jan 11, 2024 • edited Loading

Choose a reason for hiding this comment

ricardoV94 Jan 11, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ricardoV94 commented Jan 11, 2024

codecov bot commented Jan 11, 2024 • edited Loading

Codecov Report

Choose a reason for hiding this comment

velochy commented Jan 12, 2024

ricardoV94 left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ricardoV94 left a comment

Choose a reason for hiding this comment

welcome bot commented Jan 23, 2024

velochy commented Jan 11, 2024 •

edited by ricardoV94

Loading

ricardoV94 Jan 11, 2024 •

edited

Loading

ricardoV94 Jan 11, 2024 •

edited

Loading

codecov bot commented Jan 11, 2024 •

edited

Loading