Optional reference energy

quantum_systems/general_orbital_system.py

@@ -73,15 +73,31 @@ def spin_2(self):
     def spin_2_tb(self):
         return self._basis_set.spin_2_tb
 
-    def compute_reference_energy(self):
+    def compute_reference_energy(self, h=None, u=None):
         r"""Function computing the reference energy in a general spin-orbital
         system.  This is given by
 
         .. math:: E_0 = \langle \Phi_0 \rvert \hat{H} \lvert \Phi_0 \rangle
-            = h^{i}_{i} + \frac{1}{2} u^{ij}_{ij},
+            = h^{i}_{i} + \frac{1}{2} u^{ij}_{ij} + E_n,
 
-        where :math:`\lvert \Phi_0 \rangle` is the reference determinant, and
-        :math:`i, j` are occupied indices.
+        where :math:`\lvert \Phi_0 \rangle` is the reference determinant,
+        :math:`i, j` are occupied indices, :math:`h^{i}_{j}` the matrix
+        elements of the one-body Hamiltonian, :math:`u^{ij}_{kl}` the matrix
+        elements of the two-body Hamiltonian, and :math:`E_n` the nuclear
+        repulsion energy, i.e., the constant term in the Hamiltonian.
+
+        Parameters
+        ----------
+        h : np.ndarray
+            The one-body matrix elements. When `h=None` `self.h` is used.
+            If a custom `h` is passed in, the function assumes that `h` is at
+            least a `(n, n)`-array, where `n` is the number of occupied
+            indices. Default is `h=None`.
+        u : np.ndarray
+            The two-body matrix elements. When `u=None` `self.u` is used.
+            If a custom `u` is passed in, the function assumes that `u` is at
+            least a `(n, n, n, n)`-array, where `n` is the number of occupied
+            indices. Default is `u=None`.
 
         Returns
         -------
@@ -91,10 +107,13 @@ def compute_reference_energy(self):
 
         o, v = self.o, self.v
 
+        h = self.h if h is None else h
+        u = self.u if u is None else u
+
         return (
-            self.np.trace(self.h[o, o])
+            self.np.trace(h[o, o])
             + 0.5
-            * self.np.trace(self.np.trace(self.u[o, o, o, o], axis1=1, axis2=3))
+            * self.np.trace(self.np.trace(u[o, o, o, o], axis1=1, axis2=3))
             + self.nuclear_repulsion_energy
         )
 

quantum_systems/quantum_dots/one_dim/one_dim_potentials.py

@@ -16,10 +16,10 @@ def __init__(self, omega):
         self.omega = omega
 
     def __call__(self, x):
-        return 0.5 * self.omega ** 2 * x ** 2
+        return 0.5 * self.omega**2 * x**2
 
     def derivative(self, x):
-        return self.omega ** 2 * x
+        return self.omega**2 * x
 
 
 class DWPotential(HOPotential):
@@ -28,16 +28,16 @@ def __init__(self, omega, l):
         self.l = l
 
     def __call__(self, x):
-        return super().__call__(x) + 0.5 * self.omega ** 2 * (
-            0.25 * self.l ** 2 - self.l * abs(x)
+        return super().__call__(x) + 0.5 * self.omega**2 * (
+            0.25 * self.l**2 - self.l * abs(x)
         )
 
     def derivative(self, x):
         """
         Uses Heaviside function to avoid division by zero. Is ill defined in x=0
         anyways
         """
-        return super().derivative(x) - self.l * self.omega ** 2 * (
+        return super().derivative(x) - self.l * self.omega**2 * (
             np.heaviside(x, 0.5) - 0.5
         )
 
@@ -53,7 +53,7 @@ def __init__(self, a=4):
 
     def __call__(self, x):
         return (
-            (1.0 / (2 * self.a ** 2))
+            (1.0 / (2 * self.a**2))
             * (x + 0.5 * self.a) ** 2
             * (x - 0.5 * self.a) ** 2
         )
@@ -62,7 +62,7 @@ def derivative(self, x):
         a = self.a
         return (
             1
-            / a ** 2
+            / a**2
             * (
                 (x + 0.5 * a) * (x - 0.5 * a) ** 2
                 + (x - 0.5 * a) * (x + 0.5 * a) ** 2
@@ -82,10 +82,10 @@ def __init__(self, a=0.5, b=1, c=-7):
         self.c = c
 
     def __call__(self, x):
-        return self.a * x ** 6 + self.b * x ** 4 + self.c * x ** 2
+        return self.a * x**6 + self.b * x**4 + self.c * x**2
 
     def derivative(self, x):
-        return 6 * self.a * x ** 5 + 3 * self.b * x ** 3 + 2 * self.c * x
+        return 6 * self.a * x**5 + 3 * self.b * x**3 + 2 * self.c * x
 
 
 class AsymmetricDWPotential(OneDimPotential):
@@ -100,10 +100,10 @@ def __init__(self, a=1, b=1, c=-2.5):
         self.c = c
 
     def __call__(self, x):
-        return self.a * x ** 4 + self.b * x ** 3 + self.c * x ** 2
+        return self.a * x**4 + self.b * x**3 + self.c * x**2
 
     def derivative(self, x):
-        return 4 * self.a * x ** 3 + 3 * self.b * x ** 2 + 2 * self.c * x
+        return 4 * self.a * x**3 + 3 * self.b * x**2 + 2 * self.c * x
 
 
 class GaussianPotential(OneDimPotential):
@@ -115,11 +115,11 @@ def __init__(self, weight, center, deviation, np):
 
     def __call__(self, x):
         return -self.weight * self.np.exp(
-            -((x - self.center) ** 2) / (2.0 * self.deviation ** 2)
+            -((x - self.center) ** 2) / (2.0 * self.deviation**2)
         )
 
     def derivative(self, x):
-        return -(x - self.center) / self.deviation ** 2 * self(x)
+        return -(x - self.center) / self.deviation**2 * self(x)
 
 
 class GaussianPotentialHardWall(OneDimPotential):
@@ -143,7 +143,7 @@ def __call__(self, x):
 
         return (
             -self.weight
-            * np.exp(-((x - self.center) ** 2) / (2.0 * self.deviation ** 2))
+            * np.exp(-((x - self.center) ** 2) / (2.0 * self.deviation**2))
             + wall
         )
 
@@ -154,7 +154,7 @@ def __init__(self, Za=2, c=0.54878464):
         self.c = c
 
     def __call__(self, x):
-        return -self.Za / np.sqrt(x ** 2 + self.c)
+        return -self.Za / np.sqrt(x**2 + self.c)
 
     def derivative(self, x):
-        return self.Za * x / (x ** 2 + self.c) ** (3 / 2)
+        return self.Za * x / (x**2 + self.c) ** (3 / 2)

quantum_systems/quantum_dots/one_dim/one_dim_qd.py

@@ -31,7 +31,7 @@ def _trapz(f, x):
 
 @numba.njit(cache=True)
 def _shielded_coulomb(x_1, x_2, alpha, a):
-    return alpha / np.sqrt((x_1 - x_2) ** 2 + a ** 2)
+    return alpha / np.sqrt((x_1 - x_2) ** 2 + a**2)
 
 
 @numba.njit(cache=True)
@@ -134,14 +134,14 @@ def setup_basis(self):
     def ho_function(self, x, n):
         return (
             self.normalization(n)
-            * np.exp(-0.5 * self.omega * x ** 2)
+            * np.exp(-0.5 * self.omega * x**2)
             * spec.hermite(n)(np.sqrt(self.omega) * x)
         )
 
     def normalization(self, n):
         return (
             1.0
-            / np.sqrt(2 ** n * spec.factorial(n))
+            / np.sqrt(2**n * spec.factorial(n))
             * (self.omega / np.pi) ** 0.25
         )
 
@@ -156,7 +156,7 @@ def construct_position_integrals(self):
                 * Nn_up
                 * (n + 1)
                 * np.sqrt(np.pi)
-                * 2 ** n
+                * 2**n
                 * spec.factorial(n)
                 / self.omega
             )
@@ -257,8 +257,8 @@ def __init__(
     def setup_basis(self):
         dx = self.grid[1] - self.grid[0]
 
-        h_diag = 1.0 / (dx ** 2) + self.potential(self.grid[1:-1])
-        h_off_diag = -1.0 / (2 * dx ** 2) * np.ones(self.num_grid_points - 3)
+        h_diag = 1.0 / (dx**2) + self.potential(self.grid[1:-1])
+        h_off_diag = -1.0 / (2 * dx**2) * np.ones(self.num_grid_points - 3)
 
         h = (
             np.diag(h_diag)
@@ -299,7 +299,7 @@ def construct_position_integrals(self):
             for q in range(self.l):
                 self.position[0, p, q] = np.trapz(
                     self.spf[p].conj()
-                    * (self.grid + self.beta * self.grid ** 2)
+                    * (self.grid + self.beta * self.grid**2)
                     * self.spf[q],
                     self.grid,
                 )

quantum_systems/quantum_dots/two_dim/two_dim_helper.py

@@ -44,8 +44,8 @@ def spf_theta(theta, m):
 def spf_radial(r, n, m, mass, omega):
     a = bohr_radius(mass, omega)
 
-    laguerre = scipy.special.assoc_laguerre(a ** 2 * r ** 2, n, abs(m))
-    radial_dep = np.exp(-(a ** 2) * r ** 2 / 2.0)
+    laguerre = scipy.special.assoc_laguerre(a**2 * r**2, n, abs(m))
+    radial_dep = np.exp(-(a**2) * r**2 / 2.0)
 
     return (a * r) ** abs(m) * laguerre * radial_dep
 
@@ -55,8 +55,8 @@ def spf_radial_function(n, m, mass, omega):
 
     radial_function = (
         lambda r: (a * r) ** abs(m)
-        * sympy.assoc_laguerre(n, abs(m), a ** 2 * r ** 2)
-        * sympy.exp(-(a ** 2) * r ** 2 / 2.0)
+        * sympy.assoc_laguerre(n, abs(m), a**2 * r**2)
+        * sympy.exp(-(a**2) * r**2 / 2.0)
     )
 
     return radial_function
@@ -66,7 +66,7 @@ def radial_integral(r_p, r_q, order=1):
     r = sympy.Symbol("r")
 
     return sympy.integrate(
-        r * r ** order * r_p(r).conjugate() * r_q(r), (r, 0, sympy.oo)
+        r * r**order * r_p(r).conjugate() * r_q(r), (r, 0, sympy.oo)
     )
 
 
@@ -315,7 +315,7 @@ def get_double_well_one_body_elements(
 
         h[p, p] += (
             omega * get_shell_energy(n_p, m_p)
-            + omega ** 2 * barrier_strength ** 2 / 8.0
+            + omega**2 * barrier_strength**2 / 8.0
         )
 
         for q in range(num_orbitals):
@@ -327,7 +327,7 @@ def get_double_well_one_body_elements(
 
             h[p, q] -= (
                 0.5
-                * omega ** 2
+                * omega**2
                 * barrier_strength
                 * spf_norm(n_p, m_p, mass, omega)
                 * spf_norm(n_q, m_q, mass, omega)
@@ -347,21 +347,21 @@ def get_smooth_double_well_one_body_elements(
 ):
     h = np.zeros((num_orbitals, num_orbitals), dtype=dtype)
 
-    prefactor = omega ** 2 / 4
+    prefactor = omega**2 / 4
 
     for p in range(num_orbitals):
         n_p, m_p = get_indices_nm(p)
         r_p = spf_radial_function(n_p, m_p, mass, omega)
 
-        h[p, p] += omega * get_shell_energy(n_p, m_p) + omega ** 2 * a ** 2 / 64
+        h[p, p] += omega * get_shell_energy(n_p, m_p) + omega**2 * a**2 / 64
 
         for q in range(num_orbitals):
             n_q, m_q = get_indices_nm(q)
             r_q = spf_radial_function(n_q, m_q, mass, omega)
 
             h[p, q] += (
                 prefactor
-                * (1 / a ** 2)
+                * (1 / a**2)
                 * spf_norm(n_p, m_p, mass, omega)
                 * spf_norm(n_q, m_q, mass, omega)
                 * radial_integral(r_p, r_q, order=4)

quantum_systems/quantum_dots/two_dim/two_dim_ho.py

@@ -256,7 +256,7 @@ def __init__(self, *args, omega_c=0, **kwargs):
         super().__init__(*args, **kwargs)
 
     def setup_basis(self):
-        self.omega = np.sqrt(self.omega ** 2 + self.omega_c ** 2 / 4)
+        self.omega = np.sqrt(self.omega**2 + self.omega_c**2 / 4)
 
         num_orbitals = self.l
         n_array = np.arange(num_orbitals)

quantum_systems/sinc_dvr/one_dim/sinc_dvr.py

@@ -109,10 +109,10 @@ def setup_basis(self, u_repr):
         mask = np.ones(self.h.shape, dtype=bool)
         np.fill_diagonal(mask, 0)
         self.h[mask] = (-1.0) ** diff_grid[mask] / (
-            self.dx ** 2 * diff_grid[mask] ** 2
+            self.dx**2 * diff_grid[mask] ** 2
         )
         # fill diagonal
-        self.h[ind, ind] = np.pi ** 2 / (6 * self.dx ** 2)
+        self.h[ind, ind] = np.pi**2 / (6 * self.dx**2)
         self.h[ind, ind] += self.potential(self.grid)
 
         self.s = self.construct_s()
@@ -125,8 +125,8 @@ def setup_basis(self, u_repr):
 
     def set_u_repr(self, new_repr):
         """Switches representation of coloumb matrix elements between 2d and 4d"""
-        coords0 = np.arange(self.l ** 2) % self.l
-        coords1 = np.arange(self.l ** 2) // self.l
+        coords0 = np.arange(self.l**2) % self.l
+        coords1 = np.arange(self.l**2) // self.l
 
         if new_repr == self.u_repr:
             print("u repr is already {}, doing nothing".format(new_repr))
@@ -149,7 +149,7 @@ def construct_sinc_grid(self):
 
     def construct_position_integrals(self):
         self.position = np.zeros((1, self.l, self.l), dtype=self.spf.dtype)
-        self.position[0] = np.diag(self.grid + self.beta * self.grid ** 2)
+        self.position[0] = np.diag(self.grid + self.beta * self.grid**2)
 
     def construct_coulomb_elements(self, u_repr="4d"):
         """Computes Sinc-DVR matrix elements of onebody operator h and two-body

quantum_systems/spatial_orbital_system.py

@@ -103,15 +103,31 @@ def construct_general_orbital_system(
 
         return gos
 
-    def compute_reference_energy(self):
+    def compute_reference_energy(self, h=None, u=None):
         r"""Function computing the reference energy in an orbital system.
         This is given by
 
         .. math:: E_0 = \langle \Phi_0 \rvert \hat{H} \lvert \Phi_0 \rangle
-            = 2 h^{i}_{i} + 2 u^{ij}_{ij} - u^{ij}_{ji},
+            = 2 h^{i}_{i} + 2 u^{ij}_{ij} - u^{ij}_{ji} + E_n,
 
-        where :math:`\lvert \Phi_0 \rangle` is the reference determinant, and
-        :math:`i, j` are occupied indices.
+        where :math:`\lvert \Phi_0 \rangle` is the reference determinant,
+        :math:`i, j` are occupied indices, :math:`h^{i}_{j}` the matrix
+        elements of the one-body Hamiltonian, :math:`u^{ij}_{kl}` the matrix
+        elements of the two-body Hamiltonian, and :math:`E_n` the nuclear
+        repulsion energy, i.e., the constant term in the Hamiltonian.
+
+        Parameters
+        ----------
+        h : np.ndarray
+            The one-body matrix elements. When `h=None` `self.h` is used.
+            If a custom `h` is passed in, the function assumes that `h` is at
+            least a `(n, n)`-array, where `n` is the number of occupied
+            indices. Default is `h=None`.
+        u : np.ndarray
+            The two-body matrix elements. When `u=None` `self.u` is used.
+            If a custom `u` is passed in, the function assumes that `u` is at
+            least a `(n, n, n, n)`-array, where `n` is the number of occupied
+            indices. Default is `u=None`.
 
         Returns
         -------
@@ -121,11 +137,13 @@ def compute_reference_energy(self):
 
         o, v = self.o, self.v
 
+        h = self.h if h is None else h
+        u = self.u if u is None else u
+
         return (
-            2 * self.np.trace(self.h[o, o])
-            + 2
-            * self.np.trace(self.np.trace(self.u[o, o, o, o], axis1=1, axis2=3))
-            - self.np.trace(self.np.trace(self.u[o, o, o, o], axis1=1, axis2=2))
+            2 * self.np.trace(h[o, o])
+            + 2 * self.np.trace(self.np.trace(u[o, o, o, o], axis1=1, axis2=3))
+            - self.np.trace(self.np.trace(u[o, o, o, o], axis1=1, axis2=2))
             + self.nuclear_repulsion_energy
         )
 

quantum_systems/system.py

@@ -75,7 +75,7 @@ def change_to_hf_basis(self, *args, **kwargs):
         pass
 
     @abc.abstractmethod
-    def compute_reference_energy(self):
+    def compute_reference_energy(self, h=None, u=None):
         pass
 
     def compute_particle_density(self, rho_qp, C=None, C_tilde=None):