cleanup and simplifications

Author: jfeist

jftools/short_iterative_lanczos.py

@@ -7,11 +7,11 @@
 from scipy import sparse as sp
 
 try:
-    from . import short_iterative_lanczos_cython as _sil_cython
+    from .short_iterative_lanczos_cython import _lanczos_timeprop_cython
 
     have_cython_backend = True
 except ImportError:
-    _sil_cython = None
+    _lanczos_timeprop_cython = None
     have_cython_backend = False
 
 try:
@@ -57,6 +57,17 @@ def _matvec(H, phi):
     return H @ phi
 
 
+def _as_hfun(H):
+    """Return an in-place Hfun(t, phi, Hphi) callable for H."""
+    if callable(H):
+        return H
+
+    H_f = H.dot if hasattr(H, "dot") else H.__matmul__
+    def Hfun(t, phi, Hphi):
+        Hphi[:] = H_f(phi)
+
+    return Hfun
+
 def _qobj_state_io(phi0):
     outdims = phi0.dims
     outshape = phi0.full().shape
@@ -73,17 +84,7 @@ class _lanczos_timeprop_reference:
     def __init__(self, H, maxsteps, target_convg, debug=0, do_full_order=False):
         if have_qutip and isinstance(H, qutip.Qobj):
             H = _qobj_to_matrix(H)
-
-        if not callable(H):
-            # time-independent operator
-            # assume it supports dot or matmul for matrix-vector multiplication
-            def Hfun(t, phi, Hphi):
-                Hphi[:] = _matvec(H, phi)
-                return Hphi
-
-            self.Hfun = Hfun
-        else:
-            self.Hfun = H
+        self.Hfun = _as_hfun(H)
 
         self.maxsteps = maxsteps
         self.target_convg = target_convg
@@ -95,6 +96,7 @@ def Hfun(t, phi, Hphi):
 
         self.curr_coeff = zeros(maxsteps + 1, dtype=complex)
         self.prev_coeff = self.curr_coeff.copy()
+        self.breakdown_tol = 1e-14
 
     def propagate(self, phi0, ts, maxHT=None):
         phi0 = np.asarray(phi0).view(normdotndarray)
@@ -129,6 +131,7 @@ def _step(self, t, HT):
         prev_coeff = self.prev_coeff
         debug = self.debug
         Hfun = self.Hfun
+        max_lanczos_steps = min(self.maxsteps, phia[0].shape[0])
 
         HT_done = HT
 
@@ -140,9 +143,12 @@ def _step(self, t, HT):
         # doesn't converge at first step
         curr_coeff[:] = 0.0
 
-        for step in range(1, self.maxsteps + 1):
+        convg = np.inf
+        exact_complete = False
+
+        for step in range(1, max_lanczos_steps + 1):
             # set |phia(step)> to H|phia(step-1)>
-            phia[step] = Hfun(t, phia[step - 1], phia[step])
+            Hfun(t, phia[step - 1], phia[step])
             prefacs[step] = prefacs[step - 1]
             phinorm = prefacs[step] * phia[step].norm()
             # phinorm = sqrt(<q(step-1)|H H|q(step-1)>)
@@ -177,18 +183,23 @@ def _step(self, t, HT):
             # i.e. to prefac = 1.d0 / sqrt(<phi|phi>)
             phinorm = phia[step].norm()
             beta[step - 1] = prefacs[step] * phinorm
-            prefacs[step] = 1.0 / phinorm
-            if abs(log10(prefacs[step])) > 4.0:
-                phia[step] *= prefacs[step]
+            if phinorm <= self.breakdown_tol:
                 prefacs[step] = 1.0
-            if abs(beta[step - 1]) < 1e-2 and debug > 2:
-                print("WARNING! beta[%d]=%g is very small - there seems to be a linearly dependent vector!" % (step, beta[step - 1]))
-            if debug > 1:
-                # check if new vector is orthogonal to all others
-                for ii in range(step):
-                    dotpr = prefacs[ii] * prefacs[step] * phia[ii].dot(phia[step])
-                if abs(dotpr) > 1e-12:
-                    print("WARNING! vectors not orthogonal. dotpr(%d,%d) = %g" % (ii, step, dotpr))
+                phia[step][:] = 0.0
+                exact_complete = True
+            else:
+                prefacs[step] = 1.0 / phinorm
+                if abs(log10(prefacs[step])) > 4.0:
+                    phia[step] *= prefacs[step]
+                    prefacs[step] = 1.0
+                if abs(beta[step - 1]) < 1e-2 and debug > 2:
+                    print("WARNING! beta[%d]=%g is very small - there seems to be a linearly dependent vector!" % (step, beta[step - 1]))
+                if debug > 1:
+                    # check if new vector is orthogonal to all others
+                    for ii in range(step):
+                        dotpr = prefacs[ii] * prefacs[step] * phia[ii].dot(phia[step])
+                    if abs(dotpr) > 1e-12:
+                        print("WARNING! vectors not orthogonal. dotpr(%d,%d) = %g" % (ii, step, dotpr))
 
             # check convergence
             prev_coeff[:] = curr_coeff[:]
@@ -201,6 +212,9 @@ def _step(self, t, HT):
             if debug > 5:
                 print("convg:", convg)
 
+            if exact_complete or step == max_lanczos_steps:
+                break
+
             if not self.do_full_order and convg < self.target_convg:
                 break
 
@@ -209,7 +223,7 @@ def _step(self, t, HT):
             print(beta[: step - 1])
 
         # if convergence was reached in lanczos_loop, convg < target_convg, and this loop is never entered
-        while convg > self.target_convg:
+        while (not exact_complete) and convg > self.target_convg:
             # error (~convg) should be O(HT**maxsteps)
             # convg = a * HT**maxsteps
             # target_convg = a * HT_new**maxsteps
@@ -258,40 +272,35 @@ def _is_csr_matrix(H):
     return csr_array is not None and isinstance(H, csr_array)
 
 
-def _select_backend(H, backend=None):
-    if backend is None:
-        backend = "auto"
-    else:
-        backend = backend.strip().lower()
+def _select_backend(H, backend):
+    backend = backend.strip().lower()
 
     if backend == "python":
         return "python"
 
     if backend == "cython":
         if not have_cython_backend:
             raise ValueError("backend='cython' requested but Cython backend extension is not available.")
-        if _is_dense_matrix(H) or _is_csr_matrix(H):
-            return "cython"
-        raise ValueError("backend='cython' requested but Hamiltonian type is unsupported for cython backend.")
+        return "cython"
 
     if backend != "auto":
         raise ValueError("Unknown backend value '%s'. Valid values are 'python', 'cython', 'auto'." % backend)
 
-    if callable(H):
-        return "python"
-    if have_cython_backend and (_is_dense_matrix(H) or _is_csr_matrix(H)):
+    if have_cython_backend:
         return "cython"
     return "python"
 
 
 class lanczos_timeprop:
-    def __init__(self, H, maxsteps, target_convg, debug=0, do_full_order=False, backend=None):
+    def __init__(self, H, maxsteps, target_convg, debug=0, do_full_order=False, backend="auto"):
         if have_qutip and isinstance(H, qutip.Qobj):
             H = _qobj_to_matrix(H)
         self.backend = _select_backend(H, backend)
         if self.backend == "cython":
-            self._impl = _sil_cython.CythonLanczosPropagator(H, maxsteps, target_convg, debug, 
-                                                             do_full_order)
+            if not (_is_dense_matrix(H) or _is_csr_matrix(H)):
+                H = _as_hfun(H)
+
+            self._impl = _lanczos_timeprop_cython(H, maxsteps, target_convg, debug, do_full_order)
         else:
             self._impl = _lanczos_timeprop_reference(H, maxsteps, target_convg, debug, do_full_order)
 
@@ -320,6 +329,6 @@ def __getattr__(self, name):
         return getattr(self._impl, name)
 
 
-def sesolve_lanczos(H, phi0, ts, maxsteps, target_convg, maxHT=None, debug=0, do_full_order=False, backend=None):
+def sesolve_lanczos(H, phi0, ts, maxsteps, target_convg, maxHT=None, debug=0, do_full_order=False, backend="auto"):
     prop = lanczos_timeprop(H, maxsteps, target_convg, debug, do_full_order, backend)
     return prop.propagate(phi0, ts, maxHT)

jftools/short_iterative_lanczos_cython.pyx

@@ -49,18 +49,20 @@ cdef void _csr_matvec(c128[::1] data, i64[::1] indices, i64[::1] indptr,
         y[i] = acc
 
 
-cdef class CythonLanczosPropagator:
+cdef class _lanczos_timeprop_cython:
     cdef int maxsteps
     cdef double target_convg
     cdef int debug
     cdef bint do_full_order
     cdef bint is_csr
     cdef int dim
+    cdef bint is_callable
 
     cdef c128[::1, :] H_dense_f
     cdef c128[::1] H_data
     cdef i64[::1] H_indices
     cdef i64[::1] H_indptr
+    cdef object H_callable
 
     cdef f64[::1] alpha
     cdef f64[::1] beta
@@ -73,32 +75,39 @@ cdef class CythonLanczosPropagator:
     cdef f64[:, ::1] coeff_z_buf
     cdef f64[::1] coeff_work_buf
     cdef int[::1] coeff_iwork_buf
+    cdef double breakdown_tol
 
     def __cinit__(self, H, int maxsteps, double target_convg, int debug=0, bint do_full_order=False):
         self.maxsteps = maxsteps
         self.target_convg = target_convg
         self.debug = debug
         self.do_full_order = do_full_order
+        self.dim = 0
 
-        if sp.isspmatrix_csr(H) or (hasattr(sp, "csr_array") and isinstance(H, sp.csr_array)):
+        if callable(H):
+            self.is_callable = True
+            self.is_csr = False
+            self.H_callable = H
+        elif sp.isspmatrix_csr(H) or (hasattr(sp, "csr_array") and isinstance(H, sp.csr_array)):
+            self.is_callable = False
             self.is_csr = True
             self.dim = H.shape[0]
             H_csr = sp.csr_matrix(H).astype(np.complex128)
             self.H_data = np.asarray(H_csr.data, dtype=np.complex128)
             self.H_indices = np.asarray(H_csr.indices, dtype=np.int64)
             self.H_indptr = np.asarray(H_csr.indptr, dtype=np.int64)
         else:
+            self.is_callable = False
             self.is_csr = False
             self.H_dense_f = np.asfortranarray(H, dtype=np.complex128)
             self.dim = self.H_dense_f.shape[0]
 
-        # Allocate working arrays as Cython-owned buffers.
-        self.phia = view.array(shape=(maxsteps + 1, self.dim), itemsize=sizeof(c128), format="Zd", mode="c")
         self.alpha = view.array(shape=(maxsteps + 1,), itemsize=sizeof(f64), format="d")
         self.beta = view.array(shape=(maxsteps + 1,), itemsize=sizeof(f64), format="d")
         self.prefacs = view.array(shape=(maxsteps + 1,), itemsize=sizeof(f64), format="d")
         self.curr_coeff = view.array(shape=(maxsteps + 1,), itemsize=sizeof(c128), format="Zd")
         self.prev_coeff = view.array(shape=(maxsteps + 1,), itemsize=sizeof(c128), format="Zd")
+        self.breakdown_tol = 1e-14
 
         # Allocate LAPACK work buffers once and reuse in _calc_coeff.
         self.coeff_d_buf = view.array(shape=(maxsteps,), itemsize=sizeof(f64), format="d")
@@ -115,9 +124,12 @@ cdef class CythonLanczosPropagator:
         cdef c128[:, ::1] out_v
         cdef c128[::1] phi0_v = phi0
 
-        if self.dim != n:
+        if self.is_callable:
+            self.dim = n
+        elif self.dim != n:
             raise ValueError("State dimension does not match Hamiltonian dimension")
 
+        self.phia = view.array(shape=(self.maxsteps + 1, n), itemsize=sizeof(c128), format="Zd", mode="c")
         self.phia[0, :] = phi0_v
 
         tt = ts[0]
@@ -131,7 +143,7 @@ cdef class CythonLanczosPropagator:
                 HT = tf - tt
                 if maxHT is not None and HT > maxHT:
                     HT = maxHT
-                HT_done = self._step(HT)
+                HT_done = self._step(tt, HT)
                 tt += HT_done
             out_v[its, :] = self.phia[0, :]
             its += 1
@@ -173,10 +185,11 @@ cdef class CythonLanczosPropagator:
             for j in range(n):
                 coeff[j] += self.coeff_z_buf[i, j] * fac
 
-    cdef double _step(self, double HT):
-        cdef int step, ii, n
+    cdef double _step(self, double t, double HT):
+        cdef int step, ii, n, max_lanczos_steps
         cdef double HT_done = HT
         cdef double phinorm, convg, scale
+        cdef bint exact_complete = False
         cdef c128 dotpr
         cdef c128 s
         cdef c128 alpha_blas, beta_blas
@@ -191,9 +204,12 @@ cdef class CythonLanczosPropagator:
         cdef c128[::1] prev_v = self.prev_coeff
 
         n = phia_v.shape[1]
+        max_lanczos_steps = self.maxsteps
+        if max_lanczos_steps > n:
+            max_lanczos_steps = n
         vec_n = n
         incx = incy = one = 1
-        if not self.is_csr:
+        if not self.is_csr and not self.is_callable:
             m = ncol = lda = n
             trans = 'N'
             alpha_blas = 1.0 + 0.0j
@@ -204,8 +220,12 @@ cdef class CythonLanczosPropagator:
         for ii in range(self.maxsteps + 1):
             curr_v[ii] = 0.0 + 0.0j
 
-        for step in range(1, self.maxsteps + 1):
-            if self.is_csr:
+        convg = 1e300
+
+        for step in range(1, max_lanczos_steps + 1):
+            if self.is_callable:
+                self.H_callable(t, np.asarray(phia_v[step - 1]), np.asarray(phia_v[step]))
+            elif self.is_csr:
                 _csr_matvec(self.H_data, self.H_indices, self.H_indptr, phia_v[step - 1], phia_v[step])
             else:
                 blas.zgemm(&trans, &trans, &m, &one, &ncol, &alpha_blas, &self.H_dense_f[0, 0], &lda, &phia_v[step - 1, 0], &ncol, &beta_blas, &phia_v[step, 0], &m)
@@ -230,20 +250,28 @@ cdef class CythonLanczosPropagator:
 
             phinorm = _vnorm(phia_v[step])
             beta_v[step] = prefacs_v[step] * phinorm
-            prefacs_v[step] = 1.0 / phinorm
-
-            if fabs(log10(prefacs_v[step])) > 4.0:
-                s = prefacs_v[step]
-                blas.zscal(&vec_n, &s, &phia_v[step, 0], &incx)
+            if phinorm <= self.breakdown_tol:
                 prefacs_v[step] = 1.0
+                exact_complete = True
+                for ii in range(n):
+                    phia_v[step, ii] = 0.0 + 0.0j
+            else:
+                prefacs_v[step] = 1.0 / phinorm
+
+                if fabs(log10(prefacs_v[step])) > 4.0:
+                    s = prefacs_v[step]
+                    blas.zscal(&vec_n, &s, &phia_v[step, 0], &incx)
+                    prefacs_v[step] = 1.0
 
             self.prev_coeff[:] = self.curr_coeff
             self._calc_coeff(step, HT_done, curr_v)
             convg = _coeff_diff_norm(curr_v, prev_v, step + 1)
+            if exact_complete or step == max_lanczos_steps:
+                break
             if (not self.do_full_order) and convg < self.target_convg:
                 break
 
-        while convg > self.target_convg:
+        while (not exact_complete) and convg > self.target_convg:
             scale = 0.95 * pow(self.target_convg / convg, 1.0 / step)
             if scale < 0.5:
                 scale = 0.5