Hamiltonian Leapfrog Step

Problem

One leapfrog step for Hamiltonian Monte Carlo with potential U(q) = 0.5 * stiffness * q^2: update the momentum by a half step p_half = p - 0.5 * step_size * stiffness * q, the position by a full step q_new = q + step_size * p_half / mass, then the momentum by another half step using the new position. Return [q_new, p_new], or [q, p] when mass is nonpositive.

Hint

The leapfrog integrator alternates half momentum kicks with full position drifts.

Test cases

case 1

in: [0,1,1,1,0.5]

expected: [0.5,0.875]

case 2

in: [1,0,2,1,0.5]

expected: [0.75,-0.875]

case 3

in: [0,0,1,2,1]

expected: [0,0]

case 4

in: [2,-1,1,1,1]

expected: [0,-2]

case 5

in: [0.5,0.5,4,2,0.25]

expected: [0.53125,-0.015625]

Starter codePython

def hamiltonian_leapfrog_step(q, p, stiffness, mass, step_size):
    # Your code here
    pass

5 test cases · run your code against them in the editor