Fix in reset of proposal kernels

papamarkou · Dec 14, 2020 · d65dcc6 · d65dcc6
1 parent 3fc6eb3
commit d65dcc6
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 1 deletion.
diff --git a/eeyore/__init__.py b/eeyore/__init__.py
@@ -1 +1 @@
-__version__ = '0.0.14'
+__version__ = '0.0.15'
diff --git a/eeyore/samplers/mala.py b/eeyore/samplers/mala.py
@@ -30,6 +30,10 @@ def set_current(self, theta, data=None):
         self.current['target_val'], self.current['grad_val'] = \
             self.model.upto_grad_log_target(self.current['sample'].clone().detach(), x, y)
 
+    def reset(self, theta, data=None, reset_counter=True, reset_chain=True):
+        super().reset(theta, data=data, reset_counter=reset_counter, reset_chain=reset_chain)
+        self.set_kernel(self.current)
+
     def kernel_mean(self, state):
         return state['sample'] + 0.5 * self.step * state['grad_val']
 

diff --git a/eeyore/samplers/metropolis_hastings.py b/eeyore/samplers/metropolis_hastings.py
@@ -41,6 +41,10 @@ def set_current(self, theta, data=None):
         x, y = super().set_current(theta, data=data)
         self.current['target_val'] = self.model.log_target(self.current['sample'].clone().detach(), x, y)
 
+    def reset(self, theta, data=None, reset_counter=True, reset_chain=True):
+        super().reset(theta, data=data, reset_counter=reset_counter, reset_chain=reset_chain)
+        self.set_kernel(self.current)
+
     def set_kernel(self, state, scale=None, scale_tril=None):
         self.kernel.set_density_params(state['sample'].clone().detach())