test_unconstrained.py 文件源码

def test_sample_independent_momentum(self):
        mom_resample_coeff = 1.
        dtype = np.float64
        for n_dim in [10, 100, 1000]:
            mass_matrix = self.prng.normal(size=(n_dim, n_dim))
            mass_matrix = mass_matrix.dot(mass_matrix.T)
            mass_matrix_chol = la.cholesky(mass_matrix, lower=True)
            sampler = uhmc.EuclideanMetricHmcSampler(
                energy_func=energy_func,
                mass_matrix=mass_matrix,
                energy_grad=energy_grad,
                prng=self.prng,
                mom_resample_coeff=mom_resample_coeff,
                dtype=dtype)
            pos = self.prng.normal(size=(n_dim,)).astype(dtype)
            mom = sampler.sample_independent_momentum_given_position(
                pos, cache={}
            )
            assert mom.ndim == pos.ndim and mom.shape[0] == pos.shape[0], (
                'Momentum sampling returning incorrect shaped array.'
            )
            assert mom.dtype == pos.dtype, (
                'Momentum sampling returning array with incorrect dtype.'
            )
            sum_std = sum(mass_matrix.diagonal()**0.5)
            assert abs(mom.mean()) < 5. * sum_std / n_dim**0.5, (
                'Mean of sampled momentum > 5 std. from expected value.'
            )