Merge pull request #38 from Schoyen/sparse_dvr

Sparse dvr

quantum_systems/basis_set.py

@@ -514,31 +514,36 @@ def change_to_general_orbital_basis(
 
         self.l = 2 * self.l
 
-        self.a = self.np.array(a).astype(self.np.complex128).reshape(-1, 1)
-        self.b = self.np.array(b).astype(self.np.complex128).reshape(-1, 1)
-
-        # Check that spin basis elements are orthonormal
-        assert abs(self.np.dot(self.a.conj().T, self.a) - 1) < 1e-12
-        assert abs(self.np.dot(self.b.conj().T, self.b) - 1) < 1e-12
-        assert abs(self.np.dot(self.a.conj().T, self.b)) < 1e-12
-
-        (
-            self.sigma_x,
-            self.sigma_y,
-            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_2 = self.setup_spin_squared_operator(
-            self.s, self.sigma_x, self.sigma_y, self.sigma_z, 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
+        # can be found, we've removed the spin functionality.
+        if getattr(self, "u_repr", "4d") != "2d":
+            self.a = self.np.array(a).astype(self.np.complex128).reshape(-1, 1)
+            self.b = self.np.array(b).astype(self.np.complex128).reshape(-1, 1)
+
+            # Check that spin basis elements are orthonormal
+            assert abs(self.np.dot(self.a.conj().T, self.a) - 1) < 1e-12
+            assert abs(self.np.dot(self.b.conj().T, self.b) - 1) < 1e-12
+            assert abs(self.np.dot(self.a.conj().T, self.b)) < 1e-12
+
+            (
+                self.sigma_x,
+                self.sigma_y,
+                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_2 = self.setup_spin_squared_operator(
+                self.s, self.sigma_x, self.sigma_y, self.sigma_z, 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)
+            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()

quantum_systems/quantum_dots/one_dim/one_dim_potentials.py

@@ -7,6 +7,9 @@ class OneDimPotential(metaclass=abc.ABCMeta):
     def __call__(self, x):
         pass
 
+    def derivative(self, x):
+        raise NotImplementedError()
+
 
 class HOPotential(OneDimPotential):
     def __init__(self, omega):
@@ -15,6 +18,9 @@ def __init__(self, omega):
     def __call__(self, x):
         return 0.5 * self.omega ** 2 * x ** 2
 
+    def derivative(self, x):
+        return self.omega ** 2 * x
+
 
 class DWPotential(HOPotential):
     def __init__(self, omega, l):
@@ -26,6 +32,15 @@ def __call__(self, x):
             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 * (
+            np.heaviside(x, 0.5) - 0.5
+        )
+
 
 class DWPotentialSmooth(OneDimPotential):
     """
@@ -43,6 +58,17 @@ def __call__(self, x):
             * (x - 0.5 * self.a) ** 2
         )
 
+    def derivative(self, x):
+        a = self.a
+        return (
+            1
+            / a ** 2
+            * (
+                (x + 0.5 * a) * (x - 0.5 * a) ** 2
+                + (x - 0.5 * a) * (x + 0.5 * a) ** 2
+            )
+        )
+
 
 class SymmetricDWPotential(OneDimPotential):
     """
@@ -58,6 +84,9 @@ def __init__(self, a=0.5, b=1, c=-7):
     def __call__(self, x):
         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
+
 
 class AsymmetricDWPotential(OneDimPotential):
     """
@@ -73,6 +102,9 @@ def __init__(self, a=1, b=1, c=-2.5):
     def __call__(self, x):
         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
+
 
 class GaussianPotential(OneDimPotential):
     def __init__(self, weight, center, deviation, np):
@@ -86,6 +118,9 @@ def __call__(self, x):
             -((x - self.center) ** 2) / (2.0 * self.deviation ** 2)
         )
 
+    def derivative(self, x):
+        return -(x - self.center) / self.deviation ** 2 * self(x)
+
 
 class GaussianPotentialHardWall(OneDimPotential):
     """
@@ -120,3 +155,6 @@ def __init__(self, Za=2, c=0.54878464):
 
     def __call__(self, x):
         return -self.Za / np.sqrt(x ** 2 + self.c)
+
+    def derivative(self, x):
+        return self.Za * x / (x ** 2 + self.c) ** (3 / 2)