def test_example_output(self):
        solver = single_electron.EigenSolver(self.grids, self.potential_fn)
        solver.solve_ground_state()

        example = single_electron.solved_1dsolver_to_example(
            solver, self.potential_params)

        self.assertLen(example.features.feature['density'].float_list.value,
                       10)
        self.assertLen(example.features.feature['potential'].float_list.value,
                       10)
        np.testing.assert_allclose(
            example.features.feature['kinetic_energy'].float_list.value,
            [0.82937312])
        np.testing.assert_allclose(
            example.features.feature['total_energy'].float_list.value,
            [0.07467839])
        np.testing.assert_allclose(
            example.features.feature['dx'].float_list.value,
            # dx = (1 + 1) / (10 - 1)
            [0.2222222222])
        np.testing.assert_allclose(
            example.features.feature['mu'].float_list.value,
            self.potential_params['mu'])
        np.testing.assert_allclose(
            example.features.feature['sigma'].float_list.value,
            self.potential_params['sigma'])
        np.testing.assert_allclose(
            example.features.feature['coeff'].float_list.value,
            self.potential_params['coeff'])
 def test_solver_should_be_solved(self):
     solver = single_electron.EigenSolver(self.grids, self.potential_fn)
     # The solver is not solved.
     with self.assertRaisesRegexp(ValueError,
                                  r'Input solver is not solved.'):
         single_electron.solved_1dsolver_to_example(solver,
                                                    self.potential_params)
    def test_harmonic_oscillator(self, num_electrons, expected_total_energy,
                                 expected_kinetic_energy):
        grids = np.linspace(-5, 5, 501)
        potential_fn = functools.partial(single_electron.harmonic_oscillator,
                                         k=1)
        solver = single_electron.EigenSolver(grids=grids,
                                             potential_fn=potential_fn,
                                             num_electrons=num_electrons)

        solver.solve_ground_state()
        # Integrate the density over the grid points.
        norm = np.sum(solver.density) * solver.dx

        np.testing.assert_allclose(norm, num_electrons)
        # i-th eigen-energy of the quantum harmonic oscillator is
        # 0.5 + (i - 1).
        # The total energy for num_electrons states is
        # (0.5 + 0.5 + num_electrons - 1) / 2 * num_electrons
        # = num_electrons ** 2 / 2
        np.testing.assert_allclose(solver.total_energy,
                                   expected_total_energy,
                                   atol=1e-3)
        # Kinetic energy should equal to half of the total energy.
        # num_electrons ** 2 / 4
        np.testing.assert_allclose(solver.kinetic_energy,
                                   expected_kinetic_energy,
                                   atol=1e-3)
  def test_poschl_teller(
      self, num_electrons, lam, expected_energy, atol):
    solver = single_electron.EigenSolver(
        grids=np.linspace(-10, 10, 1001),
        potential_fn=functools.partial(single_electron.poschl_teller, lam=lam),
        num_electrons=num_electrons)

    solver.solve_ground_state()

    self.assertAlmostEqual(np.sum(solver.density) * solver.dx, num_electrons)
    np.testing.assert_allclose(solver.total_energy, expected_energy, atol=atol)
    def test_gaussian_dips(self):
        grids = np.linspace(-5, 5, 501)
        coeff = np.array([1., 1.])
        sigma = np.array([1., 1.])
        mu = np.array([-2., 2.])
        potential_fn = functools.partial(single_electron.gaussian_dips,
                                         coeff=coeff,
                                         sigma=sigma,
                                         mu=mu)
        solver = single_electron.EigenSolver(grids, potential_fn)

        solver.solve_ground_state()
        # Integrate the density over the grid points.
        norm = np.sum(solver.density) * solver.dx
        # Exact kinetic energy from von Weizsacker functional.
        t_vw = single_electron.vw_grid(solver.density, solver.dx)

        np.testing.assert_allclose(norm, 1.)
        # Kinetic energy should equal to the exact solution from
        # von Weizsacker kinetic energy functional.
        np.testing.assert_allclose(solver.kinetic_energy, t_vw, atol=1e-4)