Merge pull request #30 from lothian/autocorrelation

Adding autocorrelation function to RTCC.

pycc/rt/rtcc.py

@@ -47,9 +47,9 @@ class rtcc(object):
     f(): Returns a flattened NumPy array of cluster residuals
         The ODE defining function (the right-hand-side of a Runge-Kutta solver)
     collect_amps():
-        Collect the cluster amplitudes into a single vector
+        Collect the cluster amplitudes and phase into a single vector
     extract_amps():
-        Separate a flattened array of cluster amplitudes into the t1, t2, l1, and l2 components
+        Separate a flattened array of amplitudes (and phase) into the t1, t2, l1, and l2 components
     dipole()
         Compute the electronic or magnetic dipole moment for a given time t
     energy()
@@ -103,16 +103,17 @@ def f(self, t, y):
         Returns
         -------
         f(t, y): NumPy array
-            flattened array of cluster residuals
+            flattened array of cluster residuals and phase
         """
         # Extract amplitude tensors
-        t1, t2, l1, l2 = self.extract_amps(y)
+        t1, t2, l1, l2, phase = self.extract_amps(y)
 
         # Add the field to the Hamiltonian
         if isinstance(t1, torch.Tensor):
             F = self.ccwfn.H.F.clone() + self.mu_tot * self.V(t)
         else:
             F = self.ccwfn.H.F.copy() + self.mu_tot * self.V(t)
+
         # Compute the current residuals
         rt1, rt2 = self.ccwfn.residuals(F, t1, t2)
         rt1 = rt1 * (-1.0j)
@@ -126,42 +127,49 @@ def f(self, t, y):
         if self.ccwfn.local is not None:
             rl1, rl2 = self.ccwfn.Local.filter_res(rl1, rl2)
 
+        # Phase contribution = exp(-phase(t))
+        phase = self.phase(F, t1, t2)
+
         # Pack up the residuals
-        y = self.collect_amps(rt1, rt2, rl1, rl2)
+        y = self.collect_amps(rt1, rt2, rl1, rl2, phase)
 
         return y
 
-    def collect_amps(self, t1, t2, l1, l2):
+    def collect_amps(self, t1, t2, l1, l2, phase):
         """
         Parameters
         ----------
+        phase : scalar
+            current wave function phase
         t1, t2, l2, l2 : NumPy arrays
             current cluster amplitudes or residuals
 
         Returns
         -------
         NumPy array
-            amplitudes or residuals as a vector (flattened array)
+            amplitudes or residuals and phase as a vector (flattened array)
         """
         if isinstance(t1, torch.Tensor):
             t1 = torch.flatten(t1)
             t2 = torch.flatten(t2)
             l1 = torch.flatten(l1)
             l2 = torch.flatten(l2)
-            return torch.cat((t1, t2, l1, l2)).type(torch.complex128)
+            return torch.cat((t1, t2, l1, l2, phase)).type(torch.complex128)
         else:
-            return np.concatenate((t1, t2, l1, l2), axis=None).astype('complex128')
+            return np.concatenate((t1, t2, l1, l2, phase), axis=None).astype('complex128')
 
 
     def extract_amps(self, y):
         """
         Parameters
         ----------
         y : NumPy array
-            flattened array of cluster amplitudes or residuals
+            amplitudes or residuals and phase as a vector (flattened array)
 
         Returns
         -------
+        phase : scalar
+            current wave function phase
         t1, t2, l2, l2 : NumPy arrays
             current cluster amplitudes or residuals
         """
@@ -175,14 +183,17 @@ def extract_amps(self, y):
             t1 = torch.reshape(y[:len1], (no, nv))
             t2 = torch.reshape(y[len1:(len1+len2)], (no, no, nv, nv))
             l1 = torch.reshape(y[(len1+len2):(len1+len2+len1)], (no, nv))
-            l2 = torch.reshape(y[(len1+len2+len1):], (no, no, nv, nv))
+            l2 = torch.reshape(y[(len1+len2+len1):-1], (no, no, nv, nv))
         else:
             t1 = np.reshape(y[:len1], (no, nv))
             t2 = np.reshape(y[len1:(len1+len2)], (no, no, nv, nv))
             l1 = np.reshape(y[(len1+len2):(len1+len2+len1)], (no, nv))
-            l2 = np.reshape(y[(len1+len2+len1):], (no, no, nv, nv))
+            l2 = np.reshape(y[(len1+len2+len1):-1], (no, no, nv, nv))
 
-        return t1, t2, l1, l2
+        # Extract the phase
+        phase = y[-1]
+
+        return t1, t2, l1, l2, phase
 
     def dipole(self, t1, t2, l1, l2, withref = True, magnetic = False):
         """
@@ -197,7 +208,7 @@ def dipole(self, t1, t2, l1, l2, withref = True, magnetic = False):
 
         Returns
         -------
-        x, y, z : complex128
+        x, y, z : scalars
             Cartesian components of the dipole moment
         """
         opdm = self.ccdensity.compute_onepdm(t1, t2, l1, l2, withref=withref)
@@ -211,34 +222,6 @@ def dipole(self, t1, t2, l1, l2, withref = True, magnetic = False):
         z = ints[2].flatten().dot(opdm.flatten())
         return x, y, z
 
-    def energy(self, t, t1, t2, l1, l2):
-        """
-        Parameters
-        ----------
-        t : float
-            current time step in external ODE solver
-        t1, t2, l1, l2 : NumPy arrays
-            current cluster amplitudes
-
-        Returns
-        -------
-        ecc : complex128
-            CC correlation energy
-        """
-        o = self.ccwfn.o
-        v = self.ccwfn.v
-        if isinstance(t1, torch.Tensor):
-            F = self.ccwfn.H.F.clone() + self.mu_tot * self.V(t)
-        else:
-            F = self.ccwfn.H.F.copy() + self.mu_tot * self.V(t)
-
-        contract = self.contract 
-
-        ecc = 2.0 * contract('ia,ia->', F[o,v], t1)
-        L = self.ccwfn.H.L
-        ecc = ecc + contract('ijab,ijab->', build_tau(t1, t2), L[o,o,v,v])
-        return ecc
-
     def lagrangian(self, t, t1, t2, l1, l2):
         """
         Parameters
@@ -250,7 +233,7 @@ def lagrangian(self, t, t1, t2, l1, l2):
 
         Returns
         -------
-        ecc : complex128
+        ecc : scalars
             CC Lagrangian energy (including reference contribution, but excluding nuclear repulsion)
         """
         o = self.ccwfn.o
@@ -292,6 +275,78 @@ def lagrangian(self, t, t1, t2, l1, l2):
         etwo = oooo_energy + vvvv_energy + ooov_energy + vvvo_energy + ovov_energy + oovv_energy
         return eref + eone + etwo
 
+    def phase(self, F, t1, t2):
+        """
+        Parameters
+        ----------
+        F : NumPy array
+            current (field-dependent Fock operator
+        t1, t2: NumPy arrays
+            current cluster amplitudes
+
+        Returns
+        -------
+        phase: scalar
+            wave function quasienergy/phase-factor with contribution defined as = exp(-phase(t))
+        """
+        contract = self.contract
+        o = self.ccwfn.o
+        v = self.ccwfn.v
+        L = self.ccwfn.H.L
+
+        if isinstance(t1, torch.Tensor):
+            eref = 2.0 * torch.trace(F[o,o])
+            tmp = self.ccwfn.H.L[o,o,o,o].to(self.ccwfn.device1)
+            eref -= torch.trace(opt_einsum.contract('i...i', tmp.swapaxes(0,1), backend='torch'))
+            del tmp
+
+        else:
+            eref = 2.0 * np.trace(F[o,o])
+            eref -= np.trace(np.trace(self.ccwfn.H.L[o,o,o,o], axis1=1, axis2=3))
+
+        if self.ccwfn.model == 'CCD':
+            ecc = contract('ijab,ijab->', t2, L[o,o,v,v])
+        else:
+            ecc = 2.0 * contract('ia,ia->', F[o,v], t1)
+            ecc += contract('ijab,ijab->', self.ccwfn.build_tau(t1, t2), L[o,o,v,v])
+
+        return (eref + ecc) * (-1.0j)
+
+    def autocorrelation(self, y_left, y_right):
+        """
+        Parameters
+        ----------
+        y_left, y_right : Numpy arrays
+            amplitudes or residuals and phase as a vector (flattened array) for two different time steps
+
+        Returns
+        -------
+        float
+            the autocorrelation function, A(t1, t2) as defined in Eq. (18) of J. Chem. Phys. 150, 144106 (2019)
+        """
+        contract = opt_einsum.contract
+
+        t1_l, t2_l, l1_l, l2_l, phase_l = self.extract_amps(y_left)
+        t1_r, t2_r, l1_r, l2_r, phase_r = self.extract_amps(y_right)
+
+        A = 1
+        A += contract("ia,ia->", l1_l, (t1_r - t1_l))
+        A += 0.5*contract("ijab,ijab->", l2_l, (t2_r - t2_l))
+        A += 0.5*contract("ijab,ia,jb->", l2_l, t1_l, t1_l)
+        A += 0.5*contract("ijab,ia,jb->", l2_l, t1_r, t1_r)
+        A -= contract("ijab,ia,jb->", l2_l, t1_l, t1_r)
+        A *= np.exp(-phase_l) * np.exp(phase_r)
+
+        B = 1
+        B -= contract("ia,ia->", l1_r, (t1_r - t1_l))
+        B -= 0.5*contract("ijab,ijab->", l2_r, (t2_r - t2_l))
+        B += 0.5*contract("ijab,ia,jb->", l2_r, t1_r, t1_r)
+        B += 0.5*contract("ijab,ia,jb->", l2_r, t1_l, t1_l)
+        B -= contract("ijab,ia,jb->", l2_r, t1_l, t1_r)
+        B *= np.exp(-phase_r) * np.exp(phase_l)
+
+        return 0.5*A + 0.5*np.conj(B)
+
     def step(self,ODE,yi,t,ref=False):
         """
         A single step in the propagation
@@ -319,7 +374,7 @@ def step(self,ODE,yi,t,ref=False):
 
         # calculate properties
         ret = {}
-        t1, t2, l1, l2 = self.extract_amps(y)
+        t1, t2, l1, l2, phase = self.extract_amps(y)
         ret['ecc'] = self.lagrangian(t,t1,t2,l1,l2)
         mu_x, mu_y, mu_z = self.dipole(t1,t2,l1,l2,withref=ref,magnetic=False)
         ret['mu_x'] = mu_x
@@ -342,7 +397,7 @@ def propagate(self, ODE, yi, tf, ti=0, ref=False, chk=False, tchk=False,
         ODE : integrators object
             callable integrator with timestep attribute
         yi : NumPy array
-            flattened array of initial cluster amplitudes or residuals
+            flattened array of initial cluster amplitudes or residuals and phase
         tf : float
             final timestep
         ti : float
@@ -395,16 +450,17 @@ def propagate(self, ODE, yi, tf, ti=0, ref=False, chk=False, tchk=False,
                     ret_t = pk.load(ampf)
             else:
                 ret_t = {key: None}
-            t1,t2,l1,l2 = self.extract_amps(yi)
+            t1,t2,l1,l2,phase = self.extract_amps(yi)
             ret_t[key] = {"t1":t1,
                     "t2":t2,
                     "l1":l1,
-                    "l2":l2}
+                    "l2":l2,
+                    "phase":phase}
         else:
             save_t = False
 
         # initial properties
-        t1, t2, l1, l2 = self.extract_amps(yi)
+        t1, t2, l1, l2, phase = self.extract_amps(yi)
         ret[key]['ecc'] = self.lagrangian(ti,t1,t2,l1,l2)
         mu_x, mu_y, mu_z = self.dipole(t1,t2,l1,l2,withref=ref,magnetic=False)
         ret[key]['mu_x'] = mu_x
@@ -437,7 +493,7 @@ def propagate(self, ODE, yi, tf, ti=0, ref=False, chk=False, tchk=False,
 
             # save amplitudes if asked and correct timestep
             if save_t and (point%tchk<0.0001):
-                t1,t2,l1,l2 = self.extract_amps(y)
+                t1,t2,l1,l2,phase = self.extract_amps(y)
                 ret_t[key] = {"t1":t1,
                         "t2":t2,
                         "l1":l1,

pycc/tests/test_006_rtccsd.py

@@ -45,15 +45,16 @@ def test_rtcc_he_cc_pvdz():
     V = sine_square_laser(F_str, omega, tprime)
 
     # RT-CC Setup
+    phase = 0
     t0 = 0
     tf = 1.0
     h = 0.01
     rtcc = pycc.rtcc(cc, cclambda, ccdensity, V)
-    y0 = rtcc.collect_amps(cc.t1, cc.t2, cclambda.l1, cclambda.l2).astype('complex128')
+    y0 = rtcc.collect_amps(cc.t1, cc.t2, cclambda.l1, cclambda.l2, phase).astype('complex128')
     ODE = ode(rtcc.f).set_integrator('vode',atol=1e-13,rtol=1e-13)
     ODE.set_initial_value(y0, t0)
 
-    t1, t2, l1, l2 = rtcc.extract_amps(y0)
+    t1, t2, l1, l2, phase = rtcc.extract_amps(y0)
     mu0_x, mu0_y, mu0_z = rtcc.dipole(t1, t2, l1, l2)
     ecc0 = rtcc.lagrangian(t0, t1, t2, l1, l2)
 
@@ -62,7 +63,7 @@ def test_rtcc_he_cc_pvdz():
     while ODE.successful() and ODE.t < tf:
         y = ODE.integrate(ODE.t+h)
         t = ODE.t
-        t1, t2, l1, l2 = rtcc.extract_amps(y)
+        t1, t2, l1, l2, phase = rtcc.extract_amps(y)
         mu_x, mu_y, mu_z = rtcc.dipole(t1, t2, l1, l2)
         ecc = rtcc.lagrangian(t, t1, t2, l1, l2)
 

pycc/tests/test_007_dipole.py

@@ -38,8 +38,8 @@ def test_dipole_h2_2_cc_pvdz():
 
     # no laser
     rtcc = pycc.rtcc(cc, cclambda, ccdensity, None, magnetic = True)
-    y0 = rtcc.collect_amps(cc.t1, cc.t2, cclambda.l1, cclambda.l2).astype('complex128')
-    t1, t2, l1, l2 = rtcc.extract_amps(y0)
+    y0 = rtcc.collect_amps(cc.t1, cc.t2, cclambda.l1, cclambda.l2, ecc).astype('complex128')
+    t1, t2, l1, l2, phase = rtcc.extract_amps(y0)
 
     ref = [0, 0, 0.005371586416860086] # au
     mu_x, mu_y, mu_z = rtcc.dipole(t1, t2, l1, l2)

pycc/tests/test_016_chk.py

@@ -67,6 +67,7 @@ def test_chk(datadir):
     V = gaussian_laser(F_str, 0, sigma, center=center)
 
     # RTCC setup
+    phase = 0
     h = 0.1
     tf = 10
     rtcc = pycc.rtcc(cc,cclambda,ccdensity,V,magnetic=True,kick='z')
@@ -79,6 +80,7 @@ def test_chk(datadir):
     # propagate to 10au
     ODE = rk2(h)
     y0 = chk['y']
+    y0 = np.append(y0, phase)
     ti = chk['time']
     ofile = datadir.join(f"output.pk")
     tfile = datadir.join(f"t_out.pk")

pycc/tests/test_019_localrt.py

@@ -68,11 +68,12 @@ def test_rtpno(datadir):
     V = gaussian_laser(F_str, omega, sigma, center=center)
 
     # RT-CC Setup
+    phase = 0
     t0 = 0
     tf = 0.5
     h = 0.02
     rtcc = pycc.rtcc(cc, cclambda, ccdensity, V)
-    y0 = rtcc.collect_amps(cc.t1, cc.t2, cclambda.l1, cclambda.l2).astype('complex128')
+    y0 = rtcc.collect_amps(cc.t1, cc.t2, cclambda.l1, cclambda.l2, phase).astype('complex128')
     ODE = rk4(h)
 
     # Propagate
@@ -123,11 +124,12 @@ def test_rtpao(datadir):
     V = gaussian_laser(F_str, omega, sigma, center=center)
 
     # RT-CC Setup
+    phase = 0
     t0 = 0
     tf = 0.5
     h = 0.02
     rtcc = pycc.rtcc(cc, cclambda, ccdensity, V)
-    y0 = rtcc.collect_amps(cc.t1, cc.t2, cclambda.l1, cclambda.l2).astype('complex128')
+    y0 = rtcc.collect_amps(cc.t1, cc.t2, cclambda.l1, cclambda.l2, phase).astype('complex128')
     ODE = rk4(h)
 
     # Propagate

pycc/tests/test_021_rk4.py

@@ -48,15 +48,16 @@ def test_rtcc_water_cc_pvdz():
     V = gaussian_laser(F_str, omega, sigma, center)
 
     # RT-CC Setup
+    phase = 0
     t0 = 0
     tf = 0.1
     h = 0.01
     t = t0
     rtcc = pycc.rtcc(cc, cclambda, ccdensity, V)
-    y0 = rtcc.collect_amps(cc.t1, cc.t2, cclambda.l1, cclambda.l2).astype('complex128')
+    y0 = rtcc.collect_amps(cc.t1, cc.t2, cclambda.l1, cclambda.l2, phase).astype('complex128')
     y = y0
     ODE = rk4(h)
-    t1, t2, l1, l2 = rtcc.extract_amps(y0)
+    t1, t2, l1, l2, phase = rtcc.extract_amps(y0)
     mu0_x, mu0_y, mu0_z = rtcc.dipole(t1, t2, l1, l2)
     ecc0 = rtcc.lagrangian(t0, t1, t2, l1, l2)
 
@@ -75,7 +76,7 @@ def test_rtcc_water_cc_pvdz():
     while t < tf:
         y = ODE(rtcc.f, t, y)
         t += h 
-        t1, t2, l1, l2 = rtcc.extract_amps(y)
+        t1, t2, l1, l2, phase = rtcc.extract_amps(y)
         mu_x, mu_y, mu_z = rtcc.dipole(t1, t2, l1, l2)
         ecc = rtcc.lagrangian(t, t1, t2, l1, l2)
         """