diff --git a/mclmc/dynamics.py b/mclmc/dynamics.py index dab1f84..4b00f61 100644 --- a/mclmc/dynamics.py +++ b/mclmc/dynamics.py @@ -86,10 +86,10 @@ def step(x, u, g, eps, sigma): # V T V T V uu, r1 = V(eps * lambda_c, u, g * sigma) xx, ll, gg, reflect = T(eps, x, 0.5*uu*sigma) - uu = (1 - 2 * reflect) * uu + uu *= reflect uu, r2 = V(eps * (1 - 2 * lambda_c), uu, gg * sigma) xx, ll, gg, reflect = T(eps, xx, 0.5*uu*sigma) - uu = (1 - 2 * reflect) * uu + uu *= reflect uu, r3 = V(eps * lambda_c, uu, gg * sigma) #kinetic energy change @@ -108,9 +108,9 @@ def step(x, u, g, eps, sigma): # V T V uu, r1 = V(eps * 0.5, u, g * sigma) xx, l, gg, reflect = T(eps, x, uu*sigma) - uu = (1 - 2 * reflect) * uu + uu *= reflect uu, r2 = V(eps * 0.5, uu, gg * sigma) - + # kinetic energy change kinetic_change = (r1 + r2) * (d-1) diff --git a/mclmc/sampler.py b/mclmc/sampler.py index a73fb7b..eb4edc8 100644 --- a/mclmc/sampler.py +++ b/mclmc/sampler.py @@ -307,9 +307,11 @@ def single_chain_sample(self, num_steps, x_initial, random_key, output, thinning ### sampling ### + hyp = Parameters(L, eps, sigma) + self.hyp = hyp if output == OutputType.normal or output == OutputType.detailed: - X, _, E = self.sample_normal(num_steps, MCLMCState(x, u, l, g, key), Parameters(L, eps, sigma), thinning) + X, _, E = self.sample_normal(num_steps, MCLMCState(x, u, l, g, key), hyp, thinning) if output == OutputType.detailed: return X, E, L, eps else: @@ -537,8 +539,11 @@ def positive(self, where): """ mask = self.to_mask(where) - self.map = lambda x: (jnp.abs(x) * mask + x * (1- mask), x < 0.) + self.map = lambda x: (jnp.abs(x) * mask + x * (1- mask), self.to_reflect((x < 0.) * mask)) + def to_reflect(self, mask): + return 1 - 2 * mask + def rectangular(self, where, a, b): """Used if some parameters are rectangularly constrained (a< x < b)