pymc-devs / pymc

Bayesian Modeling and Probabilistic Programming in Python
https://docs.pymc.io/
Other
8.66k stars 2k forks source link

Implement unconstraining transform for LKJCorr #7380

Open johncant opened 3 months ago

johncant commented 3 months ago

I've ported this bijector from tensorflow and added to LKJCorr. This ensures that initial samples drawn from LKJCorr are positive definite, which fixes #7101 . Sampling now completes successfully with no divergences.

There are several parts I'm not comfortable with:

@fonnesbeck @twiecki @jessegrabowski @velochy - please could you take a look? I would like to make sure that this fix makes sense before adding tests and making the linters pass.

Notes:

Description

Backward method

  1. Start with identity matrix and fill lower triangular elements with unconstrained real numbers.
  2. Normalize each row so the L-2 norm is 1
  3. This is now a Cholesky factor that will always result in positive definite correlation matrices

Forward method

  1. Reconstruct the correlation matrix from its upper triangular elements
  2. Perform cholesky decomposition to obtain L
  3. The diagonal elements of L are multipliers we used to normalize the other elements.
  4. Extract those diagonal elements and divide to undo the backward method

log_jac_det

This was quite complicated to implement, so I used the symbolic jacobian.

Related Issue

Checklist

Type of change


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

welcome[bot] commented 3 months ago

Thank You Banner] :sparkling_heart: Thanks for opening this pull request! :sparkling_heart: 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.

johncant commented 2 months ago

Hi, It's unlikely I'm going to have any time to work on this for the next 6 months. The hardest part is coming up with a closed form solution for log_det_jac, which I don't think I'm very close to doing.

twiecki commented 2 months ago

Thanks for the update @johncant and for pushing this as far as you did.