class TestHMCSampler(test.Case):
def setUp(self):
self._mass_matrix = InvertibleMatrix(np.array([[2.0, 0.0], [1.0, 3.0]]))
self._initstate = State(np.array([1.0, 2.0]))
self._sampler = MockedHMCSampler(pdf=SamplePDF(), state=self._initstate.clone(),
gradient=SamplePDF().grad,
timestep=0.3, nsteps=25, mass_matrix=self._mass_matrix)
def testPropose(self):
np.random.seed(5)
initmom = np.random.multivariate_normal(mean=np.zeros(len(self._initstate.position)),
cov=self._mass_matrix)
self._sampler.state = self._initstate.clone()
np.random.seed(5)
res = self._sampler._propose()
assert(np.all(res.current_state.position == self._initstate.position))
assert(np.all(res.current_state.momentum == initmom))
assert(np.all(res.proposal_state.position == self._initstate.position * 2.0))
assert(np.all(res.proposal_state.momentum == initmom * 2))
def testHamiltonian(self):
state = State(np.array([1.0, 2.0]), np.array([2.0, 1.0]))
assert(self._sampler._hamiltonian(state) == 3.5 - np.log(1.0 / (2.0 * np.pi)))
def testCalcPacc(self):
istate = State(np.array([1.0, 2.0]), np.array([2.0, 1.0]))
fstate = State(np.array([1.0, 2.0]) * 2.0, np.array([2.0, 1.0]) * 2.0)
pcom_fstate = fstate.clone()
pcom = SimpleProposalCommunicator(istate.clone(), pcom_fstate)
self._sampler.state.momentum = None
pacc = self._sampler._calc_pacc(pcom)
dH = self._sampler._hamiltonian(fstate) - self._sampler._hamiltonian(istate)
assert(csb.numeric.exp(-dH / self._sampler.temperature) == pacc)
assert(pcom_fstate.momentum == None)
pcom_fstate = fstate.clone()
pcom = SimpleProposalCommunicator(istate.clone(), pcom_fstate)
self._sampler.state.momentum = np.array([1.0, 4.0])
pacc = self._sampler._calc_pacc(pcom)
dH = self._sampler._hamiltonian(fstate) - self._sampler._hamiltonian(istate)
assert(csb.numeric.exp(-dH / self._sampler.temperature) == pacc)
assert(np.all(pcom_fstate.momentum == np.array([1.0, 4.0])))
def _propose_swap(self, param_info):
return MockSwapCommunicator(
param_info,
Trajectory([State(np.array([1.0])),
State(np.array([2.0]))]),
Trajectory([State(np.array([2.0])),
State(np.array([1.0]))]))
def setUp(self):
self.samplers = [
MockSampler(None, State(np.array([3.0]))),
MockSampler(None, State(np.array([5.0])))
]
self.param_info = MockSwapParameterInfo(self.samplers[0],
self.samplers[1])
self.algo = MockedAbstractExchangeMC(self.samplers, [self.param_info])
def setUp(self):
self._mass_matrix = InvertibleMatrix(np.array([[2.0, 0.0], [1.0,
3.0]]))
self._initstate = State(np.array([1.0, 2.0]))
self._sampler = MockedHMCSampler(pdf=SamplePDF(),
state=self._initstate.clone(),
gradient=SamplePDF().grad,
timestep=0.3,
nsteps=25,
mass_matrix=self._mass_matrix)
def _perform_prop_pert(self, state, extra_info=None):
if self._return_momentum == True:
final = State(state.position * 2, state.momentum * 2)
else:
final = State(state.position * 2)
res = NonequilibriumTrajectory([state, final],
heat=42.0,
work=-42.0,
jacobian=1.1)
return res, None, None
def testCalcPacc(self):
istate = State(np.array([1.0, 2.0]), np.array([2.0, 1.0]))
fstate = State(np.array([1.0, 2.0]) * 2.0, np.array([2.0, 1.0]) * 2.0)
pcom_fstate = fstate.clone()
pcom = SimpleProposalCommunicator(istate.clone(), pcom_fstate)
self._sampler.state.momentum = None
pacc = self._sampler._calc_pacc(pcom)
dH = self._sampler._hamiltonian(fstate) - self._sampler._hamiltonian(
istate)
assert (csb.numeric.exp(-dH / self._sampler.temperature) == pacc)
assert (pcom_fstate.momentum == None)
pcom_fstate = fstate.clone()
pcom = SimpleProposalCommunicator(istate.clone(), pcom_fstate)
self._sampler.state.momentum = np.array([1.0, 4.0])
pacc = self._sampler._calc_pacc(pcom)
dH = self._sampler._hamiltonian(fstate) - self._sampler._hamiltonian(
istate)
assert (csb.numeric.exp(-dH / self._sampler.temperature) == pacc)
assert (np.all(pcom_fstate.momentum == np.array([1.0, 4.0])))
def setUp(self):
pdf1 = HO()
pdf2 = HO(k1=2.0, k2=2.0)
self.samplers = [
MockSampler(pdf1, State(np.array([3.0]))),
MockSampler(pdf2, State(np.array([5.0])))
]
self.param_info = RESwapParameterInfo(self.samplers[0],
self.samplers[1])
self.algo = ReplicaExchangeMC(self.samplers, [self.param_info])
def setUp(self):
self._mass_matrix = InvertibleMatrix(np.array([[2.0, 0.0], [1.0, 3.0]]))
self._initstate = State(np.array([1.0, 2.0]))
self._sampler = MockedHMCSampler(pdf=SamplePDF(), state=self._initstate.clone(),
gradient=SamplePDF().grad,
timestep=0.3, nsteps=25, mass_matrix=self._mass_matrix)
def testHMCPropagatorMM(self):
mm = InvertibleMatrix(np.array([[1., 0.], [0., 2.]]))
init_state = State(np.random.normal(size=2))
gen = HMCPropagator(self.pdf, self.gradient, self.timestep * 1.5, self.nsteps, mass_matrix=mm)
self.checkResult(gen.generate(init_state, self.nits))
def testProposeWithoutMomentum(self):
self.protocol = MockProtocol(False)
self.reverse_protocol = MockProtocol(False)
for s in self.reverse_protocol.steps:
s.set_propagation_first()
## Test without momentum
init = State(np.array([2.0]))
sampler = MockedNCMCSampler(init, self.protocol, self.reverse_protocol)
## Test _propose
# Make sure the first random number is < 0.5
np.random.seed(5)
result = sampler._propose()
self.assertEqual(result.traj.heat, -2 * 42)
self.assertEqual(result.traj.work, 2 * 42)
self.assertEqual(result.traj.initial.position[0], init.position[0])
self.assertEqual(result.traj.final.position[0], init.position[0] * 4)
self.assertEqual(result.traj.initial.momentum, None)
self.assertEqual(result.traj.final.momentum, None)
# Make sure the first random number is > 0.5
np.random.seed(4)
result = sampler._propose()
self.assertEqual(result.traj.heat, 2 * 42)
self.assertEqual(result.traj.work, -2 * 42)
self.assertEqual(result.traj.initial.position[0], init.position[0])
self.assertEqual(result.traj.final.position[0], init.position[0] * 4)
self.assertEqual(result.traj.initial.momentum, None)
self.assertEqual(result.traj.final.momentum, None)
def setUp(self):
super(ReferenceRegressions, self).setUp()
self.dt = 0.1
self.grad = self._createGradient(1.)
self.nsteps = 100
self.state = State(np.array([1.]), np.array([0.]))
def setUp(self):
super(TestIntegrators, self).setUp()
self.dt = 0.1
self.grad = self._createGradient(1.)
self.nsteps = 100
self.state = State(np.array([1.]), np.array([0.]))
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 testAcceptSwap(self):
swapcom = MockSwapCommunicator(
self.param_info,
Trajectory([State(np.array([1.0])),
State(np.array([2.0]))]),
Trajectory([State(np.array([2.0])),
State(np.array([1.0]))]))
np.random.seed(5)
swapcom.acceptance_probability = 0.75
res = self.algo._accept_swap(swapcom)
assert (res)
swapcom.acceptance_probability = 0.15
res = self.algo._accept_swap(swapcom)
assert (not res)
def propose(self, local_replica, partner_state, partner_energy, params):
n_steps = params.n_steps
propagator = self._propagator_factory(local_replica.pdf, params)
ps_pos = partner_state.position
traj = propagator.generate(State(ps_pos))
traj = GeneralTrajectory([traj.initial, traj.final], work=traj.work)
return traj
def setUp(self):
super(TestMCPropagators, self).setUp()
self.pdf = SamplePDF()
self.gradient = self._createGradient(1.)
self.timestep = 1.2
self.stepsize = 1.2
self.nsteps = 15
self.nits = 10000
self.state = State(np.random.normal(size=1))
def set1pParams(self):
init_state = State(np.random.uniform(low=-3.0, high=3.0, size=1))
self.temperatures = [0.4, 2.0]
self.samplers = [RWMCSampler(MultimodalPDF(), init_state, 0.5,
temperature=self.temperatures[0]),
RWMCSampler(MultimodalPDF(), init_state, 5.5,
temperature=self.temperatures[1])]
self.grad = self.samplers[0]._pdf.grad
self.nits = 10000
self.Ts = [lambda l: l * self.temperatures[i+1] + (1. - l) * self.temperatures[i]
for i in range(len(self.samplers) - 1)]
def testCalcPacc(self):
istate = State(np.array([1.0, 2.0]), np.array([2.0, 1.0]))
fstate = State(np.array([1.0, 2.0]) * 2.0, np.array([2.0, 1.0]) * 2.0)
pcom_fstate = fstate.clone()
pcom = SimpleProposalCommunicator(istate.clone(), pcom_fstate)
self._sampler.state.momentum = None
pacc = self._sampler._calc_pacc(pcom)
dH = self._sampler._hamiltonian(fstate) - self._sampler._hamiltonian(istate)
assert(csb.numeric.exp(-dH / self._sampler.temperature) == pacc)
assert(pcom_fstate.momentum == None)
pcom_fstate = fstate.clone()
pcom = SimpleProposalCommunicator(istate.clone(), pcom_fstate)
self._sampler.state.momentum = np.array([1.0, 4.0])
pacc = self._sampler._calc_pacc(pcom)
dH = self._sampler._hamiltonian(fstate) - self._sampler._hamiltonian(istate)
assert(csb.numeric.exp(-dH / self._sampler.temperature) == pacc)
assert(np.all(pcom_fstate.momentum == np.array([1.0, 4.0])))
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 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 testSample(self):
np.random.seed(5)
ipos = np.array([2.0])
self._sampler.state = State(ipos)
res = self._sampler.sample()
assert (res.position == ipos * 2.0)
assert (res.momentum == None)
assert (self._sampler.state.position == ipos * 2.0)
assert (self._sampler.state.momentum == None)
def set2pParams(self):
init_state = State(np.random.uniform(low=-3.0, high=3.0, size=2))
pdf = Multimodal2DPDF()
self.temperatures = [0.4, 1.0, 2.0]
self.samplers = [RWMCSampler(pdf, init_state, 0.2,
temperature=self.temperatures[0]),
RWMCSampler(pdf, init_state, .8,
temperature=self.temperatures[1]),
RWMCSampler(pdf, init_state, 2.,
temperature=self.temperatures[2])]
self.grad = self.samplers[0]._pdf.grad
self.nits = 20000
self.Ts = [lambda l: l * self.temperatures[i+1] + (1. - l) * self.temperatures[i]
for i in range(len(self.samplers) - 1)]
def testNonequilibriumStepPropagator(self):
protocol = Protocol([MockStep(True) for i in range(10)])
gen = NonequilibriumStepPropagator(protocol)
## Test generate()
init = State(np.array([2.0]), np.array([2.0]))
res = gen.generate(init)
self.assertEqual(res.final.position, init.position * (2**10))
self.assertEqual(res.final.momentum, init.momentum * (2**10))
self.assertEqual(res.work, 10 * 42)
self.assertEqual(res.heat, -10 * 42)
self.assertEqual(res.jacobian, 1.1**10)
def propose(self, local_replica, partner_state, partner_energy, params):
pdf = self._pdf_factory(local_replica, params)
n_steps = params.n_steps
propagator = self._propagator_factory(local_replica.pdf, params)
ps_pos = partner_state.position
traj = propagator.generate(State(ps_pos), return_trajectory=True)
E_remote = partner_energy
E_local = -local_replica.pdf.log_prob(traj.final.position)
deltaE = (E_local - E_remote)
traj = GeneralTrajectory([traj.initial, traj[len(traj)/2], traj.final], work=traj.work, heat=traj.heat, delta_Epot=deltaE)
return traj
def propose(self, local_replica, partner_state, partner_energy, params):
pdf = self._pdf_factory(local_replica, params)
propagator = self._propagator_factory(pdf, params)
ps_pos = partner_state.position
traj = propagator.generate(State(ps_pos, numpy.random.normal(size=ps_pos.shape)), params.n_steps, True)
traj.work = self._calculate_work(local_replica, partner_energy, traj)
E_remote = partner_energy
E_local = -local_replica.pdf.log_prob(traj.final.position)
deltaE = (E_local - E_remote) + 0.5 * numpy.sum(traj.final.momentum ** 2) - 0.5 * numpy.sum(traj.initial.momentum ** 2)
traj = GeneralTrajectory([traj.initial, traj[len(traj)/2], traj.final],
work=traj.work, heat=traj.heat,
delta_Epot=E_local - E_remote, delta_Ekin=0.5 * numpy.sum(traj.final.momentum ** 2) - 0.5 * numpy.sum(traj.initial.momentum ** 2))
return 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 testStep(self):
step = Step(MockPerturbation(), MockPropagation())
init = State(np.array([2.0]), np.array([2.0]))
## Test step with first perturbation, then propagation
res = step.perform(init)[0]
self.assertEqual(res.final.position, init.position * 4)
self.assertEqual(res.final.momentum, init.momentum * 4)
self.assertEqual(res.heat, -42.0)
self.assertEqual(res.work, 42.0)
self.assertEqual(res.jacobian, 1.1)
## Test step with first perturbation, then propagation
step.set_propagation_first()
res = step.perform(init)[0]
self.assertEqual(step.perform, step._perform_prop_pert)
self.assertEqual(res.final.position, init.position * 4)
self.assertEqual(res.final.momentum, init.momentum * 4)
self.assertEqual(res.heat, -42.0)
self.assertEqual(res.work, 42.0)
self.assertEqual(res.jacobian, 1.1)
def generate(self, init_state, length, return_trajectory=False):
final_state = State(init_state.position * 2, init_state.momentum * 2)
return Trajectory([init_state, final_state])
def setUp(self):
self._sampler = MockSampler(pdf=SamplePDF(),
state=State(np.array([1.0])))
def _run_propagator(self, state):
final = State(state.position * 2, state.momentum * 2)
return Trajectory([state, final])
def testHamiltonian(self):
state = State(np.array([1.0, 2.0]), np.array([2.0, 1.0]))
assert (self._sampler._hamiltonian(state) == 3.5 -
np.log(1.0 / (2.0 * np.pi)))
def testAcceptProposal(self):
proposal_state = State(np.array([1.234]))
res = self._sampler._accept_proposal(proposal_state)
assert (self._sampler.state == proposal_state)
def _propose(self):
pcom = SimpleProposalCommunicator(self._state,
State(self._state.position * 2.0))
return pcom
