Use aeppl to build the logprob

aesara-devs · Sep 1, 2021 · 48c3b98 · 48c3b98
1 parent 789d1d7
commit 48c3b98
Show file tree

Hide file tree

Showing 2 changed files with 147 additions and 1,016 deletions.
diff --git a/aehmc/trajectory.py b/aehmc/trajectory.py
@@ -35,15 +35,14 @@ def static_integration(
     integrator: Callable,
     step_size: float,
     num_integration_steps: int,
-    direction: int = 1,
 ) -> Callable:
     """Generate a trajectory by integrating several times in one direction."""
 
-    directed_step_size = direction * step_size
-
     def integrate(q_init, p_init, energy_init, energy_grad_init) -> IntegratorStateType:
-        def one_step(q, p, energy, energy_grad):
-            new_state = integrator(q, p, energy, energy_grad, directed_step_size)
+        def one_step(q, p, potential_energy, potential_energy_grad):
+            new_state = integrator(
+                q, p, potential_energy, potential_energy_grad, step_size
+            )
             return new_state
 
         [q, p, energy, energy_grad], _ = aesara.scan(