implement time-dependent Hamiltonian specification for numba backend

Author: jfeist

jftools/short_iterative_lanczos.py

@@ -280,6 +280,15 @@ def _is_csr_matrix(H):
     return csr_array is not None and isinstance(H, csr_array)
 
 
+def _is_numba_sum_operator(H):
+    if not isinstance(H, (tuple, list)) or len(H) < 2:
+        return False
+    for term in H[1:]:
+        if not isinstance(term, (tuple, list)) or len(term) != 2 or not callable(term[1]):
+            return False
+    return True
+
+
 def _select_backend(H, backend):
     backend = backend.strip().lower()
 
@@ -299,7 +308,7 @@ def _select_backend(H, backend):
     if backend != "auto":
         raise ValueError("Unknown backend value '%s'. Valid values are 'python', 'numba', 'cython', 'auto'." % backend)
 
-    if have_numba_backend and (_is_dense_matrix(H) or _is_csr_matrix(H)):
+    if have_numba_backend and (_is_dense_matrix(H) or _is_csr_matrix(H) or _is_numba_sum_operator(H)):
         return "numba"
     if have_cython_backend:
         return "cython"

jftools/short_iterative_lanczos_numba.py

@@ -1,18 +1,47 @@
 import numpy as np
-from numba import njit, objmode
+from numba import literal_unroll, njit, objmode
 from numba.core import types
+from numba.core.dispatcher import Dispatcher
 from numba.extending import overload
 from numba_lapack import dstevd, zgemm
 from scipy import sparse as sp
 
 _CALLABLE_OPERATOR_REGISTRY = {}
 
+
 def _register_callable_operator(H):
     handle = id(H)
     _CALLABLE_OPERATOR_REGISTRY[handle] = H
     return handle
 
 
+def _evaluate_operator_coefficient(coeff_handle, t):
+    return complex(_CALLABLE_OPERATOR_REGISTRY[coeff_handle](t))
+
+
+def _evaluate_operator_coefficient_numba(coeff_spec, t):
+    raise NotImplementedError
+
+
+@overload(_evaluate_operator_coefficient_numba)
+def _overload_evaluate_operator_coefficient_numba(coeff_spec, t):
+    if isinstance(coeff_spec, types.Integer):
+        def impl(coeff_spec, t):
+            with objmode(coeff='complex128'):
+                coeff = _evaluate_operator_coefficient(coeff_spec, t)
+            return coeff
+
+        return impl
+
+    if isinstance(coeff_spec, types.Dispatcher):
+        def impl(coeff_spec, t):
+            return complex(coeff_spec(t))
+
+        return impl
+
+    raise TypeError("Coefficient spec must be a callable handle or a Numba-jitted function")
+
+
 @njit(cache=True)
 def _vdot_numba(a, b):
     out = 0.0 + 0.0j
@@ -44,34 +73,51 @@ def _axpy_numba(scale, x, y):
         y[idx] += scale * x[idx]
 
 
-def _apply_H_operator_numba(H, t, x, y):
+@njit(cache=True)
+def _scal_numba(scale, x):
+    for idx in range(x.shape[0]):
+        x[idx] *= scale
+
+
+def _apply_H_operator_numba(H, t, x, y, alpha, beta):
     raise NotImplementedError
 
 
 @overload(_apply_H_operator_numba)
-def _overload_apply_H_operator_numba(H, t, x, y):
+def _overload_apply_H_operator_numba(H, t, x, y, alpha, beta):
     if isinstance(H, types.BaseTuple) and len(H) == 3:
-        def impl(H, t, x, y):
+        def impl(H, t, x, y, alpha, beta):
             data, indices, indptr = H
             for row in range(indptr.shape[0] - 1):
                 acc = 0.0 + 0.0j
                 for pos in range(indptr[row], indptr[row + 1]):
                     acc += data[pos] * x[indices[pos]]
-                y[row] = acc
+                y[row] = alpha * acc + beta * y[row]
 
     elif isinstance(H, types.Array):
-        def impl(H, t, x, y):
+        def impl(H, t, x, y, alpha, beta):
             charN = np.uint8(ord("N"))
-            zgemm(charN, charN, H.shape[0], 1, H.shape[1], 1.0 + 0.0j, H,
-                  H.shape[0], x, x.shape[0], 0.0 + 0.0j, y, y.shape[0])
+            zgemm(charN, charN, H.shape[0], 1, H.shape[1], alpha, H,
+                  H.shape[0], x, x.shape[0], beta, y, y.shape[0])
+
+    elif isinstance(H, types.BaseTuple) and len(H) == 2:
+        def impl(H, t, x, y, alpha, beta):
+            H0, H_terms = H
+            _apply_H_operator_numba(H0, t, x, y, alpha, beta)
+            for term in literal_unroll(H_terms):
+                Hk, coeff_spec = term
+                coeff = _evaluate_operator_coefficient_numba(coeff_spec, t)
+                _apply_H_operator_numba(Hk, t, x, y, alpha * coeff, 1.0 + 0.0j)
 
     elif isinstance(H, types.Integer):
-        def impl(H, t, x, y):
+        def impl(H, t, x, y, alpha, beta):
+            if alpha != 1.0 or beta != 0.0:
+                raise ValueError("Scaling not supported for callable operators in Numba Lanczos backend")
             with objmode():
                 _CALLABLE_OPERATOR_REGISTRY[H](t, x, y)
 
     else:
-        raise TypeError("Numba Lanczos operator must be dense array, CSR tuple, or callable handle")
+        raise TypeError("Numba Lanczos operator must be dense array, CSR tuple, specialized sum tuple, or callable handle")
 
     return impl
 
@@ -115,7 +161,7 @@ def _step_numba(operator, t, HT, config, scratch):
 
     for step in range(1, max_lanczos_steps + 1):
         step_count = step
-        _apply_H_operator_numba(operator, t, phia[step - 1], phia[step])
+        _apply_H_operator_numba(operator, t, phia[step - 1], phia[step], 1.0 + 0.0j, 0.0 + 0.0j)
         prefacs[step] = prefacs[step - 1]
         phinorm = prefacs[step] * _vnorm_numba(phia[step])
 
@@ -138,12 +184,12 @@ def _step_numba(operator, t, HT, config, scratch):
         beta[step - 1] = prefacs[step] * phinorm
         if phinorm <= breakdown_tol:
             prefacs[step] = 1.0
-            phia[step] = 0.
+            phia[step][:] = 0.0 + 0.0j
             exact_complete = True
         else:
             prefacs[step] = 1.0 / phinorm
             if abs(np.log10(prefacs[step])) > 4.0:
-                phia[step] *= prefacs[step]
+                _scal_numba(prefacs[step], phia[step])
                 prefacs[step] = 1.0
 
         prev_coeff[:] = curr_coeff
@@ -205,6 +251,27 @@ def _normalize_static_operator(H):
     return (H_dense.shape[0], H_dense)
 
 
+def _normalize_sum_operator(H):
+    dim, H0 = _normalize_static_operator(H[0])
+    H_terms = []
+    handles = []
+    for term in H[1:]:
+        if not isinstance(term, (tuple, list)) or len(term) != 2:
+            raise TypeError("Time-dependent Numba operator terms must be (H_k, f_k) pairs")
+        Hk, fk = term
+        Hk_dim, Hk_norm = _normalize_static_operator(Hk)
+        if Hk_dim != dim:
+            raise ValueError("All H_k operators must match the dimension of H_0")
+        if isinstance(fk, Dispatcher):
+            coeff_spec = fk
+        else:
+            coeff_spec = _register_callable_operator(fk)
+            handles.append(coeff_spec)
+        H_terms.append((Hk_norm, coeff_spec))
+
+    return dim, (H0, tuple(H_terms)), tuple(handles)
+
+
 def _allocate_scratch(maxsteps, dim):
     return (
         np.empty(maxsteps + 1, dtype=np.float64),  # alpha diagonal terms
@@ -231,18 +298,25 @@ def __init__(self, H, maxsteps, target_convg, debug=0, do_full_order=False):
         self.breakdown_tol = 1e-14
         self.config = (maxsteps, target_convg, do_full_order, self.breakdown_tol)
         self.H = H
+        self._registry_handles = []
 
-        if callable(H):
+        if _is_numba_sum_operator_input(H):
+            self.dim, self.operator, handles = _normalize_sum_operator(H)
+            self._registry_handles.extend(handles)
+            self.scratch = _allocate_scratch(maxsteps, self.dim)
+        elif callable(H):
             self.operator = _register_callable_operator(H)
+            self._registry_handles.append(self.operator)
             self.dim = None
             self.scratch = None
         else:
             self.dim, self.operator = _normalize_static_operator(H)
             self.scratch = _allocate_scratch(maxsteps, self.dim)
 
     def __del__(self):
-        if callable(self.H) and self.operator in _CALLABLE_OPERATOR_REGISTRY:
-            del _CALLABLE_OPERATOR_REGISTRY[self.operator]
+        for handle in self._registry_handles:
+            if handle in _CALLABLE_OPERATOR_REGISTRY:
+                del _CALLABLE_OPERATOR_REGISTRY[handle]
 
     def propagate(self, phi0, ts, maxHT=None):
         phi0 = np.asarray(phi0, dtype=np.complex128)
@@ -268,4 +342,13 @@ def propagate(self, phi0, ts, maxHT=None):
 
         _propagate_numba(self.operator, ts, use_maxht, maxht_value, self.config, self.scratch, out)
 
-        return out
+        return out
+
+
+def _is_numba_sum_operator_input(H):
+    if not isinstance(H, (tuple, list)) or len(H) < 2:
+        return False
+    for term in H[1:]:
+        if not isinstance(term, (tuple, list)) or len(term) != 2 or not callable(term[1]):
+            return False
+    return True

tests/test_all.py

@@ -2,6 +2,7 @@
 
 import numpy as np
 import pytest
+from numba import njit
 from scipy.sparse import diags
 from scipy.sparse.linalg import LinearOperator, expm_multiply
 
@@ -265,6 +266,94 @@ def Hfun(t, phi, Hphi):
         assert np.allclose(phi, phi_ref, rtol=2e-4, atol=1e-6)
 
 
+def test_short_iterative_lanczos_numba_h0_sum_fk_hk_dense_exact():
+    if not jftools.short_iterative_lanczos.have_numba_backend:
+        pytest.skip("numba backend not available")
+
+    n = 3
+    h0_diag = np.array([0.8, -0.6, 0.2], dtype=float)
+    h1_diag = np.array([0.4, 0.1, -0.3], dtype=float)
+    h2_diag = np.array([-0.15, 0.25, 0.05], dtype=float)
+    omega1 = 2.3
+    omega2 = 1.7
+    phi0 = _normalized_random_state(n, seed=53)
+    ts = np.linspace(0.0, 0.6, 5)
+
+    H0 = np.diag(h0_diag).astype(complex)
+    H1 = np.diag(h1_diag).astype(complex)
+    H2 = np.diag(h2_diag).astype(complex)
+
+    @njit(cache=True)
+    def f1(t):
+        return np.cos(omega1 * t)
+
+    @njit(cache=True)
+    def f2(t):
+        return np.sin(omega2 * t)
+
+    H_numba = (H0, (H1, f1), (H2, f2))
+    prop = jftools.short_iterative_lanczos.lanczos_timeprop(H_numba, maxsteps=8, target_convg=1e-13, backend="auto")
+    assert prop.backend == "numba"
+
+    out = prop.propagate(phi0, ts, maxHT=2e-4)
+
+    for t, phi in zip(ts, out):
+        theta = h0_diag * t
+        theta += (h1_diag / omega1) * np.sin(omega1 * t)
+        theta += (h2_diag / omega2) * (1.0 - np.cos(omega2 * t))
+        phi_ref = np.exp(-1j * theta) * phi0
+        assert np.allclose(phi, phi_ref, rtol=2e-4, atol=1e-6)
+
+
+def test_short_iterative_lanczos_numba_h0_sum_fk_hk_csr_matches_callable():
+    if not jftools.short_iterative_lanczos.have_numba_backend:
+        pytest.skip("numba backend not available")
+
+    n = 20
+    H0 = _make_chain_hamiltonian(n).astype(complex)
+    H1 = diags([np.linspace(-1.0, 1.0, n)], [0], shape=(n, n), format="csr").astype(complex)
+    H2 = diags([0.2 * np.ones(n - 1), 0.2 * np.ones(n - 1)], [-1, 1], shape=(n, n), format="csr").astype(complex)
+    phi0 = _normalized_random_state(n, seed=54)
+    ts = np.linspace(0.0, 0.4, 5)
+
+    @njit(cache=True)
+    def f1(t):
+        return np.cos(1.9 * t)
+
+    @njit(cache=True)
+    def f2(t):
+        return 0.4 * np.sin(0.8 * t)
+
+    def Hfun(t, phi, Hphi):
+        Hphi[:] = H0.dot(phi)
+        Hphi[:] += np.cos(1.9 * t) * H1.dot(phi)
+        Hphi[:] += 0.4 * np.sin(0.8 * t) * H2.dot(phi)
+        return Hphi
+
+    H_numba = (H0, (H1, f1), (H2, f2))
+    out_numba = jftools.short_iterative_lanczos.sesolve_lanczos(
+        H_numba,
+        phi0,
+        ts,
+        maxsteps=14,
+        target_convg=1e-12,
+        maxHT=0.05,
+        backend="numba",
+    )
+    out_python = jftools.short_iterative_lanczos.sesolve_lanczos(
+        Hfun,
+        phi0,
+        ts,
+        maxsteps=14,
+        target_convg=1e-12,
+        maxHT=0.05,
+        backend="python",
+    )
+
+    for phi_numba, phi_python in zip(out_numba, out_python):
+        assert np.allclose(phi_numba, phi_python, rtol=5e-9, atol=5e-10)
+
+
 @pytest.mark.parametrize("backend", ["python", "numba", "cython"])
 def test_short_iterative_lanczos_full_basis_completion_stops_cleanly(backend):
     if backend == "numba" and not jftools.short_iterative_lanczos.have_numba_backend: