def testReducedHamiltonian(self):
pdf = HO(k1=2.0, k2=2.0)
init = State(np.array([2.0]), np.array([-2.0]))
ham = ReducedHamiltonian(lambda x: pdf.log_prob(x, 0.0), pdf.gradient, temperature=4.0)
self.assertEqual(4.0, ham.E(init.position))
self.assertEqual(2.0, ham.kinetic_energy(init.momentum))
self.assertEqual(0.0, ham.kinetic_energy(None))
self.assertEqual(-1.0, ham.rlog_prob(init.position))
self.assertEqual(0.5, ham.rkinetic_energy(init.momentum))
self.assertEqual(1.5, ham(init))
def testNCMCPropagator(self):
Nhalf = 5
dt = 0.1
md_tl = 5
ks = np.linspace(1.0, 0.2, Nhalf).tolist()
sigmas = [1 / np.sqrt(k) for k in ks]
sigmas += sigmas[::-1][1:]
N = len(sigmas)
pdfs = [SamplePDF(sigma=s) for s in sigmas]
hamiltonians = [
ReducedHamiltonian(pdfs[i].log_prob, pdfs[i].grad)
for i in range(N)
]
sys_infos = [HamiltonianSysInfo(hamiltonians[i]) for i in range(N)]
steps = [
Step(
ReducedHamiltonianPerturbation(sys_infos[i],
sys_infos[i + 1],
evaluate_work=False),
PlainMDPropagation(sys_infos[i + 1],
PlainMDPropagationParam(
dt, md_tl, pdfs[i + 1].grad),
evaluate_heat=False)) for i in range(N - 1)
]
rv_steps = [
Step(
ReducedHamiltonianPerturbation(sys_infos[i],
sys_infos[i + 1],
evaluate_work=False),
PlainMDPropagation(sys_infos[i],
PlainMDPropagationParam(
dt, md_tl, pdfs[i].grad),
evaluate_heat=False)) for i in range(N - 1)
]
for s in rv_steps:
s.set_propagation_first()
protocol = Protocol(steps)
rv_protocol = Protocol(rv_steps)
class MDProbStepNCMCSampler(AbstractNCMCSampler):
def _calc_pacc(self, proposal_communicator):
return np.exp(-proposal_communicator.traj.deltaH)
class MDPropStepNCMCPropagator(AbstractNCMCPropagator):
def _init_sampler(self, init_state):
self._sampler = MDProbStepNCMCSampler(init_state,
self.protocol,
self.reverse_protocol)
gen = MDPropStepNCMCPropagator(protocol, rv_protocol)
init_state = State(np.array([1.0]))
traj = gen.generate(init_state, self.nits, return_trajectory=True)
self.checkResult(traj)
def testReducedHamiltonianPerturbation(self):
pdf = HO(k1=1.0, k2=2.0)
redham1 = ReducedHamiltonian(lambda x: pdf.log_prob(x, 0.0))
redham2 = ReducedHamiltonian(lambda x: pdf.log_prob(x, 1.0))
sys1 = HamiltonianSysInfo(redham1)
sys2 = HamiltonianSysInfo(redham2)
init = State(np.array([2.0]), np.array([2.0]))
traj = Trajectory([init, init])
hampert = ReducedHamiltonianPerturbation(sys1, sys2)
## Test _calculate_work
self.assertEqual(hampert._calculate_work(traj), 2.0)
## Test __call__
result = hampert(init)
self.assertEqual(result.initial.position[0], init.position[0])
self.assertEqual(result.initial.momentum[0], init.momentum[0])
self.assertEqual(result.initial.position[0], result.final.position[0])
self.assertEqual(result.initial.momentum[0], result.final.momentum[0])
self.assertEqual(result.work, 2.0)
self.assertEqual(result.jacobian, 1.0)
def testReducedHamiltonian(self):
pdf = HO(k1=2.0, k2=2.0)
init = State(np.array([2.0]), np.array([-2.0]))
ham = ReducedHamiltonian(lambda x: pdf.log_prob(x, 0.0),
pdf.gradient,
temperature=4.0)
self.assertEqual(4.0, ham.E(init.position))
self.assertEqual(2.0, ham.kinetic_energy(init.momentum))
self.assertEqual(0.0, ham.kinetic_energy(None))
self.assertEqual(-1.0, ham.rlog_prob(init.position))
self.assertEqual(0.5, ham.rkinetic_energy(init.momentum))
self.assertEqual(1.5, ham(init))
def testHMCPropagation(self):
pdf = HO()
sys = HamiltonianSysInfo(ReducedHamiltonian(pdf.log_prob,
pdf.gradient))
param = HMCPropagationParam(None, None, None)
hmcprop = HMCPropagationMocked(sys, param)
init = State(np.array([2.0]), np.array([2.0]))
## Test _set_mass_matrix
d = len(init.position)
param = HMCPropagationParam(None,
None,
None,
mass_matrix=InvertibleMatrix(np.eye(d)))
hmcprop = HMCPropagationMocked(sys, param)
hmcprop._set_mass_matrix(init)
self.assertEqual(hmcprop.param.mass_matrix,
InvertibleMatrix(np.eye(len(init.position))))
param = HMCPropagationParam(None, None, None)
hmcprop = HMCPropagationMocked(sys, param)
hmcprop._set_mass_matrix(init)
self.assertEqual(hmcprop.param.mass_matrix,
InvertibleMatrix(np.eye(len(init.position))))
## Test _calculate_heat
final = State(init.position * 2, init.momentum * 2)
traj = Trajectory([init, final])
self.assertEqual(hmcprop._calculate_heat(traj), 6.0)
## Test __call__
result = hmcprop(init)
self.assertEqual(init.position, result.initial.position)
self.assertEqual(init.momentum, result.initial.momentum)
self.assertEqual(result.final.position, init.position * 2)
self.assertEqual(result.final.momentum, init.momentum * 2)
self.assertEqual(result.heat, 6.0)
def sys_after(self):
pdf = HO()
return HamiltonianSysInfo(
ReducedHamiltonian(pdf.log_prob, pdf.gradient))