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 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(self):

        pdf = HO()

        return HamiltonianSysInfo(pdf.log_prob, pdf.gradient)
    def sys_after(self):

        pdf = HO()

        return HamiltonianSysInfo(
            ReducedHamiltonian(pdf.log_prob, pdf.gradient))