Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add MCLMC sampler #586

Merged
merged 91 commits into from
Dec 5, 2023
Merged
Changes from 1 commit
Commits
Show all changes
91 commits
Select commit Hold shift + click to select a range
fbe7f75
initial draft of mclmc
reubenharry Nov 10, 2023
3a23242
refactor
reubenharry Nov 11, 2023
86b3a90
wip
reubenharry Nov 11, 2023
e82550f
wip
reubenharry Nov 11, 2023
f0e1bec
wip
reubenharry Nov 11, 2023
4d7dc57
wip
reubenharry Nov 11, 2023
82b8466
wip
reubenharry Nov 11, 2023
a4d403b
fix pre-commit
reubenharry Nov 11, 2023
a67ecb7
remove dim from class
reubenharry Nov 11, 2023
3dd4f74
add docstrings
reubenharry Nov 11, 2023
5d8061d
add mclmc to init
reubenharry Nov 13, 2023
5428f2c
Merge branch 'main' of https://github.com/blackjax-devs/blackjax
reubenharry Nov 13, 2023
59ecc8a
Merge branch 'main' into refactor
reubenharry Nov 13, 2023
2bf639e
move minimal_norm to integrators
reubenharry Nov 13, 2023
172fee0
move update pos and momentum
reubenharry Nov 13, 2023
b710e62
remove params
reubenharry Nov 13, 2023
3cc52fd
Infer the shape from inverse_mass_matrix outside the function step
reubenharry Nov 14, 2023
57d5c3b
use tree_map
reubenharry Nov 14, 2023
7e70d78
integration now aligned with mclmc repo
reubenharry Nov 15, 2023
1343463
dE and logdensity align too (fixed sign error)
reubenharry Nov 15, 2023
e53a877
make L and step size arguments to kernel
reubenharry Nov 15, 2023
05517b6
rough draft of tuning: works
reubenharry Nov 15, 2023
d84a23d
remove inv mass matrix
reubenharry Nov 15, 2023
de1e5cf
almost correct
reubenharry Nov 15, 2023
263ab3a
almost correct
reubenharry Nov 16, 2023
777213d
move tuning to adaptation
reubenharry Nov 16, 2023
e75274a
tuning works in this commit
reubenharry Nov 16, 2023
8a89f13
clean up 1
reubenharry Nov 16, 2023
49b3bec
remove sigma from tuning
reubenharry Nov 16, 2023
81999f9
wip
reubenharry Nov 16, 2023
8ab01f2
fix linting
reubenharry Nov 17, 2023
6266bc4
rename T and V
reubenharry Nov 17, 2023
ca984e7
uniformity wip
reubenharry Nov 17, 2023
59ffb21
make uniform implementation of integrators
reubenharry Nov 17, 2023
8f9214f
make uniform implementation of integrators
reubenharry Nov 18, 2023
b2e3b8e
fix minimal norm integrator
reubenharry Nov 18, 2023
2fb2293
add warning to tune3
reubenharry Nov 18, 2023
59e4424
Refactor integrators.py to make it more general.
junpenglao Nov 19, 2023
6684413
temp: explore
reubenharry Nov 19, 2023
4284092
Refactor to use integrator generation functions
junpenglao Nov 20, 2023
4a514dd
Additional refactoring
junpenglao Nov 20, 2023
ef1f62d
Minor clean up.
junpenglao Nov 21, 2023
af43521
Use standard JAX ops
junpenglao Nov 21, 2023
0dd419d
new integrator
reubenharry Nov 23, 2023
0c8330e
add references
reubenharry Nov 23, 2023
e6fa2bb
merge
reubenharry Nov 23, 2023
40fc61c
flake
reubenharry Nov 24, 2023
6ea5320
temporarily add 'explore'
reubenharry Nov 24, 2023
c83dc1a
temporarily add 'explore'
reubenharry Nov 25, 2023
c8b43be
Adding a test for energy preservation.
junpenglao Nov 26, 2023
8894248
fix formatting
junpenglao Nov 26, 2023
9865145
wip: tests
reubenharry Nov 26, 2023
68464bc
Merge branch 'integrator_refactor' into refactor
reubenharry Nov 26, 2023
0c61412
use pytrees for partially_refresh_momentum, and add test
reubenharry Nov 26, 2023
a66af60
Merge branch 'main' into refactor
junpenglao Nov 27, 2023
be07631
update docstring
reubenharry Nov 27, 2023
71d934b
resolve conflict
reubenharry Nov 27, 2023
a170d0b
remove 'explore'
reubenharry Nov 27, 2023
8cfb75f
fix pre-commit
reubenharry Nov 27, 2023
b42e77e
adding randomized MCHMC
JakobRobnik Nov 29, 2023
2b323ce
wip checkpoint on tuning
reubenharry Dec 1, 2023
9a41cdf
align blackjax and mclmc repos, for tuning
reubenharry Dec 1, 2023
cdbb4f6
use effective_sample_size
reubenharry Dec 1, 2023
947d717
patial rename
reubenharry Dec 1, 2023
e9ab7b4
rename
reubenharry Dec 1, 2023
72d70c6
clean up tuning
reubenharry Dec 1, 2023
c121beb
clean up tuning
reubenharry Dec 1, 2023
fe99163
IN THIS COMMIT, BLACKJAX AND ORIGINAL REPO AGREE. SEED IS FIXED.
reubenharry Dec 2, 2023
c456efe
RANDOMIZE KEYS
reubenharry Dec 2, 2023
d0a008a
ADD TEST
reubenharry Dec 2, 2023
d692498
ADD TEST
reubenharry Dec 2, 2023
3e8d8ea
Merge branch 'main' of https://github.com/blackjax-devs/blackjax
reubenharry Dec 2, 2023
eda029a
Merge branch 'main' into refactor
reubenharry Dec 2, 2023
a45f58f
MERGE MAIN
reubenharry Dec 2, 2023
2a21c56
INCREASE CODE COVERAGE
reubenharry Dec 2, 2023
67f0de9
REMOVE REDUNDANT LINE
reubenharry Dec 2, 2023
3f55f5f
ADD NAME 'mclmc'
reubenharry Dec 2, 2023
666c540
SPLIT KEYS AND FIX DOCSTRING
reubenharry Dec 2, 2023
c1615f5
FIX MINOR ERRORS
reubenharry Dec 2, 2023
ae1bf30
FIX MINOR ERRORS
reubenharry Dec 2, 2023
3c2dbad
Merge branch 'main' of https://github.com/blackjax-devs/blackjax
reubenharry Dec 2, 2023
c396aa1
FIX CONFLICT IN BIB
reubenharry Dec 2, 2023
0902a1c
RANDOMIZE KEYS (reversion)
reubenharry Dec 2, 2023
2e3c80b
PRECOMMIT CLEAN UP
reubenharry Dec 2, 2023
604b5a9
ADD KWARGS FOR DEFAULT HYPERPARAMS
reubenharry Dec 3, 2023
50b1c95
Merge branch 'main' of https://github.com/blackjax-devs/blackjax
reubenharry Dec 4, 2023
fecd82b
Merge branch 'main' into refactor
reubenharry Dec 4, 2023
50a8243
UPDATE ESS
reubenharry Dec 5, 2023
a20a681
NAME CHANGES
reubenharry Dec 5, 2023
75e71de
NAME CHANGES
reubenharry Dec 5, 2023
70f1dd5
MINOR FIXES
reubenharry Dec 5, 2023
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
Adding a test for energy preservation.
Co-authored-by: Reuben Cohn-Gordon <[email protected]>
  • Loading branch information
junpenglao and reubenharry committed Nov 26, 2023

Verified

This commit was created on GitHub.com and signed with GitHub’s verified signature. The key has expired.
commit c8b43be2867729a6005c191b08e9444405954637
18 changes: 18 additions & 0 deletions blackjax/util.py
Original file line number Diff line number Diff line change
@@ -82,6 +82,24 @@ def generate_gaussian_noise(
return unravel_fn(mu + linear_map(sigma, sample))


def generate_unit_vector(rng_key: PRNGKey, position: ArrayLikeTree,) -> Array:
"""Generate a random unit vector with output structure that match a given PyTree.

Parameters
----------
rng_key:
The pseudo-random number generator key used to generate random numbers.
position:
PyTree that the structure the output should to match.

Returns
-------
Random unit vector that match the structure of position.
"""
p, unravel_fn = ravel_pytree(position)
sample = normal(rng_key, shape=p.shape, dtype=p.dtype)
return unravel_fn(sample / jnp.linalg.norm(sample))

def pytree_size(pytree: ArrayLikeTree) -> int:
"""Return the dimension of the flatten PyTree."""
return sum(jnp.size(value) for value in tree_leaves(pytree))
40 changes: 40 additions & 0 deletions tests/mcmc/test_integrators.py
Original file line number Diff line number Diff line change
@@ -10,6 +10,7 @@

import blackjax.mcmc.integrators as integrators
from blackjax.mcmc.integrators import esh_dynamics_momentum_update_one_step
from blackjax.util import generate_unit_vector


def HarmonicOscillator(inv_mass_matrix, k=1.0, m=1.0):
@@ -130,6 +131,9 @@ def kinetic_energy(p):
"velocity_verlet": {"algorithm": integrators.velocity_verlet, "precision": 1e-4},
"mclachlan": {"algorithm": integrators.mclachlan, "precision": 1e-5},
"yoshida": {"algorithm": integrators.yoshida, "precision": 1e-6},
"noneuclidean_leapfrog" : {"algorithm": integrators.noneuclidean_leapfrog},
"noneuclidean_mclachlan" : {"algorithm": integrators.noneuclidean_mclachlan},
"noneuclidean_yoshida" : {"algorithm": integrators.noneuclidean_yoshida},
}


@@ -224,6 +228,42 @@ def test_esh_momentum_update(self, dims):
next_momentum1, *_ = update_stable(momentum, gradient, step_size, 1.0)
np.testing.assert_array_almost_equal(next_momentum, next_momentum1)

@chex.all_variants(with_pmap=False)
@parameterized.parameters(
[
"noneuclidean_leapfrog",
"noneuclidean_mclachlan",
"noneuclidean_yoshida",
],
)
def test_noneuclidean_integrator(self, integrator_name):
integrator = algorithms[integrator_name]
cov = jnp.asarray([[1., .5] , [.5, 2.]])
logdensity_fn = lambda x: stats.multivariate_normal.logpdf(
x, jnp.zeros([2]), cov)

step = self.variant(integrator["algorithm"](logdensity_fn))

rng = jax.random.key(4263456)
key0, key1 = jax.random.split(rng, 2)
position_init = jax.random.normal(key0, (2,))
momentum_init = generate_unit_vector(key1, position_init)
step_size = .0001
initial_state = integrators.new_integrator_state(
logdensity_fn, position_init, momentum_init)

final_state, kinetic_energy_change = jax.lax.scan(
lambda state, _: step(state, step_size),
initial_state,
xs=None,
length=15,
)

# Check the conservation of energy.
potential_energy_change = final_state.logdensity - initial_state.logdensity
energy_change = kinetic_energy_change[-1] + potential_energy_change
self.assertAlmostEqual(energy_change, 0, delta=1e-3)


if __name__ == "__main__":
absltest.main()
Loading