# Path: blackjax/adaptation/mass_matrix.py
# def mass_matrix_adaptation(
#     is_diagonal_matrix: bool = True,
# ) -> Tuple[Callable, Callable, Callable]:
#     """Adapts the values in the mass matrix by computing the covariance
#     between parameters.
# 
#     Parameters
#     ----------
#     is_diagonal_matrix
#         When True the algorithm adapts and returns a diagonal mass matrix
#         (default), otherwise adaps and returns a dense mass matrix.
# 
#     Returns
#     -------
#     init
#         A function that initializes the step of the mass matrix adaptation.
#     update
#         A function that updates the state of the mass matrix.
#     final
#         A function that computes the inverse mass matrix based on the current
#         state.
#     """
#     wc_init, wc_update, wc_final = welford_algorithm(is_diagonal_matrix)
# 
#     def init(n_dims: int) -> MassMatrixAdaptationState:
#         """Initialize the matrix adaptation.
# 
#         Parameters
#         ----------
#         ndims
#             The number of dimensions of the mass matrix, which corresponds to
#             the number of dimensions of the chain position.
#         """
#         if is_diagonal_matrix:
#             inverse_mass_matrix = jnp.ones(n_dims)
#         else:
#             inverse_mass_matrix = jnp.identity(n_dims)
# 
#         wc_state = wc_init(n_dims)
# 
#         return MassMatrixAdaptationState(inverse_mass_matrix, wc_state)
# 
#     def update(
#         mm_state: MassMatrixAdaptationState, position: Array
#     ) -> MassMatrixAdaptationState:
#         """Update the algorithm's state.
# 
#         Parameters
#         ----------
#         state:
#             The current state of the mass matrix adapation.
#         position:
#             The current position of the chain.
#         """
#         inverse_mass_matrix, wc_state = mm_state
#         position, _ = jax.flatten_util.ravel_pytree(position)
#         wc_state = wc_update(wc_state, position)
#         return MassMatrixAdaptationState(inverse_mass_matrix, wc_state)
# 
#     def final(mm_state: MassMatrixAdaptationState) -> MassMatrixAdaptationState:
#         """Final iteration of the mass matrix adaptation.
# 
#         In this step we compute the mass matrix from the covariance matrix computed
#         by the Welford algorithm, and re-initialize the later.
#         """
#         _, wc_state = mm_state
#         covariance, count, mean = wc_final(wc_state)
# 
#         # Regularize the covariance matrix, see Stan
#         scaled_covariance = (count / (count + 5)) * covariance
#         shrinkage = 1e-3 * (5 / (count + 5))
#         if is_diagonal_matrix:
#             inverse_mass_matrix = scaled_covariance + shrinkage
#         else:
#             inverse_mass_matrix = scaled_covariance + shrinkage * jnp.identity(
#                 mean.shape[0]
#             )
# 
#         ndims = jnp.shape(inverse_mass_matrix)[-1]
#         new_mm_state = MassMatrixAdaptationState(inverse_mass_matrix, wc_init(ndims))
# 
#         return new_mm_state
# 
#     return init, update, final


import itertools
import chex
import jax
import numpy as np
from absl.testing import absltest, parameterized
from blackjax.adaptation.mass_matrix import mass_matrix_adaptation
"""Test the welford adaptation algorithm."""




class MassMatrixAdaptationTest(chex.TestCase):
    @chex.all_variants(with_pmap=False)
    @parameterized.parameters(itertools.product([1, 3], [True, False]))
    def test_welford_adaptation(self, n_dim, is_mass_matrix_diagonal):
        num_samples = 3_000

        # Generate samples from a multivariate normal distribution whose
        # covariance matrix we are going to try to recover.
        np.random.seed(0)
        mu = np.random.randn(n_dim)
        a = np.random.randn(n_dim, n_dim)
        cov = np.matmul(a.T, a)
        samples = np.random.multivariate_normal(mu, cov, num_samples)


# Path: blackjax/adaptation/mass_matrix.py
# def mass_matrix_adaptation(
#     is_diagonal_matrix: bool = True,
# ) -> Tuple[Callable, Callable, Callable]:
#     """Adapts the values in the mass matrix by computing the covariance
#     between parameters.
# 
#     Parameters
#     ----------
#     is_diagonal_matrix
#         When True the algorithm adapts and returns a diagonal mass matrix
#         (default), otherwise adaps and returns a dense mass matrix.
# 
#     Returns
#     -------
#     init
#         A function that initializes the step of the mass matrix adaptation.
#     update
#         A function that updates the state of the mass matrix.
#     final
#         A function that computes the inverse mass matrix based on the current
#         state.
#     """
#     wc_init, wc_update, wc_final = welford_algorithm(is_diagonal_matrix)
# 
#     def init(n_dims: int) -> MassMatrixAdaptationState:
#         """Initialize the matrix adaptation.
# 
#         Parameters
#         ----------
#         ndims
#             The number of dimensions of the mass matrix, which corresponds to
#             the number of dimensions of the chain position.
#         """
#         if is_diagonal_matrix:
#             inverse_mass_matrix = jnp.ones(n_dims)
#         else:
#             inverse_mass_matrix = jnp.identity(n_dims)
# 
#         wc_state = wc_init(n_dims)
# 
#         return MassMatrixAdaptationState(inverse_mass_matrix, wc_state)
# 
#     def update(
#         mm_state: MassMatrixAdaptationState, position: Array
#     ) -> MassMatrixAdaptationState:
#         """Update the algorithm's state.
# 
#         Parameters
#         ----------
#         state:
#             The current state of the mass matrix adapation.
#         position:
#             The current position of the chain.
#         """
#         inverse_mass_matrix, wc_state = mm_state
#         position, _ = jax.flatten_util.ravel_pytree(position)
#         wc_state = wc_update(wc_state, position)
#         return MassMatrixAdaptationState(inverse_mass_matrix, wc_state)
# 
#     def final(mm_state: MassMatrixAdaptationState) -> MassMatrixAdaptationState:
#         """Final iteration of the mass matrix adaptation.
# 
#         In this step we compute the mass matrix from the covariance matrix computed
#         by the Welford algorithm, and re-initialize the later.
#         """
#         _, wc_state = mm_state
#         covariance, count, mean = wc_final(wc_state)
# 
#         # Regularize the covariance matrix, see Stan
#         scaled_covariance = (count / (count + 5)) * covariance
#         shrinkage = 1e-3 * (5 / (count + 5))
#         if is_diagonal_matrix:
#             inverse_mass_matrix = scaled_covariance + shrinkage
#         else:
#             inverse_mass_matrix = scaled_covariance + shrinkage * jnp.identity(
#                 mean.shape[0]
#             )
# 
#         ndims = jnp.shape(inverse_mass_matrix)[-1]
#         new_mm_state = MassMatrixAdaptationState(inverse_mass_matrix, wc_init(ndims))
# 
#         return new_mm_state
# 
#     return init, update, final


# Path: tests/test_mass_matrix.py
import itertools
import chex
import jax
import numpy as np
from absl.testing import absltest, parameterized
from blackjax.adaptation.mass_matrix import mass_matrix_adaptation

"""Test the welford adaptation algorithm."""




class MassMatrixAdaptationTest(chex.TestCase):
    @chex.all_variants(with_pmap=False)
    @parameterized.parameters(itertools.product([1, 3], [True, False]))
    def test_welford_adaptation(self, n_dim, is_mass_matrix_diagonal):
        num_samples = 3_000

        # Generate samples from a multivariate normal distribution whose
        # covariance matrix we are going to try to recover.
        np.random.seed(0)
        mu = np.random.randn(n_dim)
        a = np.random.randn(n_dim, n_dim)
        cov = np.matmul(a.T, a)
        samples = np.random.multivariate_normal(mu, cov, num_samples)


        init, update, final = mass_matrix_adaptation(is_mass_matrix_diagonal)