Merge pull request #49 from Schoyen/spin-tests

Spin fix

quantum_systems/basis_set.py

@@ -55,6 +55,7 @@ def __init__(
         self._spin_y = None
         self._spin_z = None
         self._spin_2 = None
+        self._spin_2_tb = None
 
         self._spf = None
         self._bra_spf = None
@@ -167,6 +168,17 @@ def spin_2(self, spin_2):
 
         self._spin_2 = spin_2
 
+    @property
+    def spin_2_tb(self):
+        return self._spin_2_tb
+
+    @spin_2_tb.setter
+    def spin_2_tb(self, spin_2_tb):
+        assert self.includes_spin
+        assert all(self.check_axis_lengths(spin_2_tb, self.l))
+
+        self._spin_2_tb = spin_2_tb
+
     @property
     def sigma_x(self):
         return self._sigma_x
@@ -258,6 +270,7 @@ def change_module(self, np):
             ("_spin_y", self.spin_y),
             ("_spin_z", self.spin_z),
             ("_spin_2", self.spin_2),
+            ("_spin_2_tb", self.spin_2_tb),
         ]:
             setattr(self, name, self.change_arr_module(arr, self.np))
 
@@ -278,6 +291,7 @@ def cast_to_complex(self):
             ("_spin_y", self.spin_y),
             ("_spin_z", self.spin_z),
             ("_spin_2", self.spin_2),
+            ("_spin_2_tb", self.spin_2_tb),
         ]:
             if arr is not None:
                 setattr(self, name, arr.astype(np.complex128))
@@ -329,7 +343,7 @@ def _change_basis_one_body_elements(self, C, C_tilde):
                 self.s, C, C_tilde=C_tilde, np=self.np
             )
 
-        for spin in [self.spin_x, self.spin_y, self.spin_z]:
+        for spin in [self.spin_x, self.spin_y, self.spin_z, self.spin_2]:
             if spin is not None:
                 spin = self.transform_one_body_elements(
                     spin, C, C_tilde=C_tilde, np=self.np
@@ -340,9 +354,9 @@ def _change_basis_two_body_elements(self, C, C_tilde):
             self.u, C, np=self.np, C_tilde=C_tilde
         )
 
-        if self.spin_2 is not None:
-            self.spin_2 = self.transform_two_body_elements(
-                self.spin_2, C, np=self.np, C_tilde=C_tilde
+        if self.spin_2_tb is not None:
+            self.spin_2_tb = self.transform_two_body_elements(
+                self.spin_2_tb, C, np=self.np, C_tilde=C_tilde
             )
 
     def _change_basis_dipole_moment(self, C, C_tilde):
@@ -470,6 +484,10 @@ def anti_symmetrize_two_body_elements(self):
         """
         if not self._anti_symmetrized_u:
             self.u = self.anti_symmetrize_u(self.u)
+
+            if self.spin_2_tb is not None:
+                self.spin_2_tb = self.anti_symmetrize_u(self.spin_2_tb)
+
             self._anti_symmetrized_u = True
 
     def change_to_general_orbital_basis(
@@ -514,6 +532,12 @@ def change_to_general_orbital_basis(
 
         self.l = 2 * self.l
 
+        overlap = self.s.copy()
+
+        self.h = self.add_spin_one_body(self.h, np=self.np)
+        self.s = self.add_spin_one_body(self.s, np=self.np)
+        self.u = self.add_spin_two_body(self.u, np=self.np)
+
         # temporary change to allow 2d representations of two-body operators, such as
         # for dvr. A 2d representation is necessary for large basis sets, in which case
         # self.spin_2 also becomes huge. Until a better representation of self.spin_2
@@ -533,18 +557,14 @@ def change_to_general_orbital_basis(
                 self.sigma_z,
             ) = self.setup_pauli_matrices(self.a, self.b, self.np)
 
-            self.spin_x = 0.5 * self.np.kron(self.s, self.sigma_x)
-            self.spin_y = 0.5 * self.np.kron(self.s, self.sigma_y)
-            self.spin_z = 0.5 * self.np.kron(self.s, self.sigma_z)
+            self.spin_x = 0.5 * self.np.kron(overlap, self.sigma_x)
+            self.spin_y = 0.5 * self.np.kron(overlap, self.sigma_y)
+            self.spin_z = 0.5 * self.np.kron(overlap, self.sigma_z)
 
-            self.spin_2 = self.setup_spin_squared_operator(
-                self.s, self.sigma_x, self.sigma_y, self.sigma_z, self.np
+            self.spin_2, self.spin_2_tb = self.setup_spin_squared_operator(
+                self.spin_x, self.spin_y, self.spin_z, self.s, self.np
             )
 
-            self.h = self.add_spin_one_body(self.h, np=self.np)
-            self.s = self.add_spin_one_body(self.s, np=self.np)
-            self.u = self.add_spin_two_body(self.u, np=self.np)
-
         if anti_symmetrize:
             self.anti_symmetrize_two_body_elements()
 
@@ -629,46 +649,56 @@ def setup_pauli_matrices(a, b, np):
         return sigma_x, sigma_y, sigma_z
 
     @staticmethod
-    def setup_spin_squared_operator(overlap, sigma_x, sigma_y, sigma_z, np):
-        r"""Static method computing the matrix elements of the two-body spin
-        squared operator, :math:`\hat{S}^2`. The spin-basis is chosen by the
-        Pauli matrices.
+    def setup_spin_squared_operator(spin_x, spin_y, spin_z, overlap, np):
+        r"""Static method computing the matrix elements of the one- and
+        two-body spin squared operator, :math:`\hat{S}^2`. The operator is
+        computed by
+
+        .. math:: \hat{S}^2 = \hat{S}_x^2 + \hat{S}_y^2 + \hat{S}_z^2,
+
+        where each squared direction :math:`\hat{S}_i^2` can be written
+
+        .. math:: \hat{S}_i^2
+            = (S_i)^{p}_{r} (S_i)^{r}_{q}
+            \hat{c}^{\dagger}_{p} \hat{c}_{q}
+            + (S_i)^{p}_{r} (S_i)^{q}_{s}
+            \hat{c}^{\dagger}_{p} \hat{c}^{\dagger}_{q}
+            \hat{c}_{s} \hat{c}_{r},
+
+        in the second quantization formulation.
 
         Parameters
         ----------
+        spin_x : np.ndarray
+            Spin-matrix in :math:`x`-direction including orbital overlap.
+        spin_y : np.ndarray
+            Spin-matrix in :math:`y`-direction including orbital overlap.
+        spin_z : np.ndarray
+            Spin-matrix in :math:`z`-direction including orbital overlap.
         overlap : np.ndarray
-            The overlap matrix elements between the spatial orbitals.
-        sigma_x : np.ndarray
-            Pauli spin-matrix in :math:`x`-direction.
-        sigma_y : np.ndarray
-            Pauli spin-matrix in :math:`y`-direction.
-        sigma_z : np.ndarray
-            Pauli spin-matrix in :math:`z`-direction.
+            The overlap matrix elements between the spin-orbitals.
         np : module
             An appropriate array and linalg module.
 
         Returns
         -------
-        np.ndarray
-            The spin-squared operator as an array on the form ``(l, l, l, l)``,
-            where ``l`` is the number of spin-orbitals.
+        (np.ndarray, np.ndarray)
+            The spin-squared operator as two arrays on the form ``(l, l)`` and
+            ``(l, l, l, l)``, where ``l`` is the number of spin-orbitals. The
+            former corresponds to the one-body part of the spin-squared
+            operator whereas the latter is the two-body part.
         """
-        overlap_2 = np.einsum("pr, qs -> pqrs", overlap, overlap)
 
-        # The 2 in sigma_*_2 (confusingly) enough does not denote the squared
-        # operator, but rather that it is a two-spin operator.
-        sigma_x_2 = np.kron(sigma_x, np.eye(2)) + np.kron(np.eye(2), sigma_x)
-        sigma_y_2 = np.kron(sigma_y, np.eye(2)) + np.kron(np.eye(2), sigma_y)
-        sigma_z_2 = np.kron(sigma_z, np.eye(2)) + np.kron(np.eye(2), sigma_z)
+        l = len(spin_x)
 
-        S_2_spin = (
-            sigma_x_2 @ sigma_x_2
-            + sigma_y_2 @ sigma_y_2
-            + sigma_z_2 @ sigma_z_2
-        ) / 4
-        S_2_spin = S_2_spin.reshape(2, 2, 2, 2)
+        spin_2 = np.zeros_like(spin_x)
+        spin_2_tb = np.zeros((l, l, l, l), dtype=spin_2.dtype)
 
-        return np.kron(overlap_2, S_2_spin)
+        for s_i in [spin_x, spin_y, spin_z]:
+            spin_2 += s_i @ overlap @ s_i
+            spin_2_tb += np.einsum("pr, qs -> pqrs", s_i, s_i)
+
+        return spin_2, spin_2_tb
 
     @staticmethod
     def add_spin_spf(spf, np):
@@ -693,13 +723,7 @@ def add_spin_one_body(h, np):
 
     @staticmethod
     def add_spin_two_body(_u, np):
-        u = _u.transpose(1, 3, 0, 2)
-        u = np.kron(u, np.eye(2))
-        u = u.transpose(2, 3, 0, 1)
-        u = np.kron(u, np.eye(2))
-        u = u.transpose(0, 2, 1, 3)
-
-        return u
+        return np.kron(_u, np.einsum("pr, qs -> pqrs", np.eye(2), np.eye(2)))
 
     @staticmethod
     def anti_symmetrize_u(_u):

quantum_systems/general_orbital_system.py

@@ -69,6 +69,10 @@ def spin_z(self):
     def spin_2(self):
         return self._basis_set.spin_2
 
+    @property
+    def spin_2_tb(self):
+        return self._basis_set.spin_2_tb
+
     def compute_reference_energy(self):
         r"""Function computing the reference energy in a general spin-orbital
         system.  This is given by

tests/test_spin.py

@@ -2,10 +2,13 @@
 import numpy as np
 
 from quantum_systems import (
+    ODQD,
     BasisSet,
     RandomBasisSet,
     GeneralOrbitalSystem,
     SpatialOrbitalSystem,
+    construct_pyscf_system_ao,
+    construct_pyscf_system_rhf,
 )
 
 
@@ -131,11 +134,14 @@ def test_overlap_squared():
 
 
 def test_spin_squared():
-    n = 4
-    l = 10
+    n = 2
+    l = 2
     dim = 2
 
     spas = SpatialOrbitalSystem(n, RandomBasisSet(l, dim))
+    spas = SpatialOrbitalSystem(
+        n, ODQD(l, 8, 1001, potential=ODQD.HOPotential(1))
+    )
     gos = spas.construct_general_orbital_system(a=[1, 0], b=[0, 1])
 
     overlap = spas.s
@@ -171,30 +177,84 @@ def test_spin_squared():
         singlet[0].T @ S_sq_spin.reshape(4, 4) @ singlet[0], 0
     )
 
-    S_sq = np.kron(overlap_sq, S_sq_spin)
-    assert S_sq.shape == gos.u.shape
+    # S_sq = np.kron(overlap_sq, S_sq_spin)
+    # assert S_sq.shape == gos.u.shape
 
-    np.testing.assert_allclose(S_sq, gos.spin_2)
+    # np.testing.assert_allclose(S_sq, 0.5 * gos.spin_2_tb)
 
-    for P in range(gos.l):
-        p = P // 2
-        sigma = P % 2
+    # for P in range(gos.l):
+    #     p = P // 2
+    #     sigma = P % 2
 
-        for Q in range(gos.l):
-            q = Q // 2
-            tau = Q % 2
+    #     for Q in range(gos.l):
+    #         q = Q // 2
+    #         tau = Q % 2
 
-            for R in range(gos.l):
-                r = R // 2
-                gamma = R % 2
+    #         for R in range(gos.l):
+    #             r = R // 2
+    #             gamma = R % 2
 
-                for S in range(gos.l):
-                    s = S // 2
-                    delta = S % 2
+    #             for S in range(gos.l):
+    #                 s = S // 2
+    #                 delta = S % 2
 
-                    np.testing.assert_allclose(
-                        overlap[p, r]
-                        * overlap[q, s]
-                        * S_sq_spin[sigma, tau, gamma, delta],
-                        S_sq[P, Q, R, S],
-                    )
+    #                 np.testing.assert_allclose(
+    #                     overlap[p, r]
+    #                     * overlap[q, s]
+    #                     * S_sq_spin[sigma, tau, gamma, delta],
+    #                     S_sq[P, Q, R, S],
+    #                 )
+
+
+def test_spin_squared_constructions():
+    # TODO: Try to make this test applicable for non-orthonomal basis sets.
+
+    # n = 2
+    # l = 10
+    # system = GeneralOrbitalSystem(n, RandomBasisSet(l, 3))
+    # system = GeneralOrbitalSystem(
+    #     n, ODQD(l, 10, 1001, potential=ODQD.HOPotential(1))
+    # )
+
+    # system = construct_pyscf_system_ao("he")
+
+    system = construct_pyscf_system_rhf("he", basis="cc-pVTZ")
+
+    spin_dir_tb_orig = []
+    spin_dir_tb_pm = []
+
+    spin_p = system.spin_x + 1j * system.spin_y
+    spin_m = system.spin_x - 1j * system.spin_y
+    spin_z = system.spin_z
+    s = system.s
+
+    np.testing.assert_allclose(
+        spin_p @ s @ spin_m - spin_m @ s @ spin_p, 2 * spin_z, atol=1e-10
+    )
+
+    # S^2 = S_- * S_+ + S_z + S_z^2
+    spin_2_mp = spin_m @ spin_p + spin_z + spin_z @ spin_z
+    spin_2_tb_mp = (
+        0.5 * np.einsum("pr, qs -> pqrs", spin_m, spin_p)
+        + 0.5 * np.einsum("pr, qs -> pqrs", spin_p, spin_m)
+        + np.einsum("pr, qs -> pqrs", spin_z, spin_z)
+    )
+    spin_2_tb_mp = system._basis_set.anti_symmetrize_u(spin_2_tb_mp)
+
+    # S^2 = S_+ * S_- - S_z + S_z^2
+    spin_2_pm = spin_p @ spin_m - spin_z + spin_z @ spin_z
+    spin_2_tb_pm = (
+        0.5 * np.einsum("pr, qs -> pqrs", spin_m, spin_p)
+        + 0.5 * np.einsum("pr, qs -> pqrs", spin_p, spin_m)
+        + np.einsum("pr, qs -> pqrs", spin_z, spin_z)
+    )
+    spin_2_tb_pm = system._basis_set.anti_symmetrize_u(spin_2_tb_pm)
+
+    np.testing.assert_allclose(spin_2_mp, spin_2_pm, atol=1e-10)
+    np.testing.assert_allclose(spin_2_tb_mp, spin_2_tb_pm)
+
+    np.testing.assert_allclose(spin_2_mp, system.spin_2, atol=1e-10)
+    np.testing.assert_allclose(spin_2_tb_mp, system.spin_2_tb)
+
+    np.testing.assert_allclose(spin_2_pm, system.spin_2, atol=1e-10)
+    np.testing.assert_allclose(spin_2_tb_pm, system.spin_2_tb)