Source code for newtonprop.propagator

"""Newton Propagator"""
import logging
import numpy as np

try:

    from numba import jit
    # the the user-supplied `A` is efficient (runs in compiled code), some
    # parts of the Newton algorithm (especially _extend_leja) can become
    # signficant. In this case, compiling those parts of the Newton propagation
    # with number can get something like an extra factor 4 in speed.

except ImportError:  # pragma: nocover

[docs] def jit(*args, **kwargs): def _noop_decorator(func): return func return _noop_decorator
__all__ = ["step"] def _arnoldi(A, dt, v0, m_max, inner=np.vdot, norm=np.linalg.norm): """Calculate the (extended) Hessenberg matrix of an operator A(t) and return it together with the m+1 (m <= m_max) Arnoldi vectors (orthonormlized Krylov vectors). Also return the combined Ritz values (all eigenvalues of the Hessenberg matrices of size 1..m ) Args A (callable): Function encoding the operator A. Calling `A(v0)` must return the result of applying A to v0. dt (float): Time step v0: Initial state m_max (int): Maximal Krylov dimension inner (callable): Function that evaluates an inner product. Must take two arguments of the type of `v0` and return a complex number. If None, defaults to :func:`numpy.vdot`. norm (callable): Function that calculates the norm of a state. Must take one argument of the type of `v0` and return a real number. If None, defaults to :func:`numpy.linalg.norm`. """ m = m_max # Hessenberg matrix (at maximum size) Hess = np.matrix( np.zeros(shape=(m_max + 1, m_max + 1), dtype=np.complex128) ) # Eigenvalues of all Hess Ritz = [] arnoldi_vecs = [] beta = norm(v0) if abs(beta - 1.0) > 1.0e-10: print("beta = ", beta) raise AssertionError( "v0 must have norm 1.0. Mismatch between `inner` and `norm`?" ) v = v0 / beta arnoldi_vecs.append(v) for j in range(m): v = A(v) # v_{j+1} for i, v_i in enumerate(arnoldi_vecs): Hess[i, j] = dt * inner(v_i, v) v = v - (Hess[i, j] / dt) * v_i # At this point, we have finished the (j+1) x (j+1) Hessenberg matrix Ritz.extend(np.linalg.eigvals(Hess[: j + 1, : j + 1])) h_next = norm(v) Hess[j + 1, j] = h_next * dt if h_next <= 1.0e-14: # abort early at convergence m = j break v *= 1 / h_next # normalize arnoldi_vecs.append(v) # At this point, arnoldi_vecs contains m+1 elements Ritz = np.array(Ritz, dtype=np.complex128) return arnoldi_vecs, Hess[: m + 1, : m + 1], Ritz, m def _normalize_points(z): """Given a set of complex points z, return the normalization radius""" r = 0.0 for z_i in z: r_i = abs(z_i) if r_i > r: r = r_i # we need to enlarge the radius a little bit to account for points that # will be added in later iterations r *= 1.2 # arbitary factor assert r > 0.0, "Radius is zero" return r @jit(nopython=True) def _expmi(x): return np.exp(-1j * x) @jit(nopython=True) def _expi(x): return np.exp(1j * x) @jit(nopython=True) def _exp(x): return np.exp(x) @jit() def _extend_newton_coeffs(func, old_a, new_leja, center, radius): """Extend a set of Newton coefficients, by using a set of new_leja points which are normalized with the given center and radius """ n_old = len(old_a) m = len(new_leja) - n_old a = np.zeros(n_old + m, dtype=np.complex128) a[:n_old] = old_a n0 = n_old if n_old == 0: a[0] = func(new_leja[0]) n0 = 1 for k in range(n0, n_old + m): d = 1.0 pn = 0.0 for n in range(1, k): # 1..k-1 zd = new_leja[k] - new_leja[n - 1] d *= zd / radius pn += a[n] * d zd = new_leja[k] - new_leja[k - 1] d *= zd / radius assert abs(d) > 1.0e-200, "Divided differences too small" a[k] = (func(new_leja[k]) - a[0] - pn) / d return a @jit(nopython=True) def _extend_leja(old_leja, new_points, n_use): """Given a set of normalized (ordered) Leja points, extract `n_use` points from the (normalized) new_points, and append them to the set of leja points """ n_old = len(old_leja) new_leja = np.zeros(n_old + n_use, dtype=np.complex128) new_leja[:n_old] = old_leja[:] i_add_start = 0 if n_old == 0: # At the very beginning, start with the point that has largest absolute # value for i in range(len(new_points) - 1): # 0 .. n_old - 2 if abs(new_points[i]) > abs(new_points[-1]): temp = new_points[i] new_points[i] = new_points[-1] new_points[-1] = temp new_leja[0] = new_points[-1] i_add_start = 1 # find the best point for new_leja[n_old+n] n_added = i_add_start ex = 1.0 / (n_old + n_use) for i_add in range(i_add_start, n_use): p_max = 0.0 i_max = 0 # the new leja are defined with index 0 .. (n_old-1)+n # the new candidates are defined with index 0 .. len(new_points)-1+n for i in range(len(new_points) - i_add): # trial points (candidates) p = 1.0 for j in range(n_old + i_add): # existing leja points p *= np.abs(new_points[i] - new_leja[j]) ** ex # at this point p is the divided difference denominator for the # candidate with index i if p > p_max: p_max = p i_max = i # XXX if p_max below limit: abort new_leja[n_old + i_add] = new_points[i_max] n_added += 1 # remove the used point by moving in the last point new_points[i_max] = new_points[len(new_points) - 1 - i_add] return new_leja, n_added
[docs]def step( A, v, dt, func='exp', m_max=10, maxrestart=100, tol=1e-12, inner=None, norm=None, zero=None, ): r"""Evaluate $f(A\,dt)\,v$, for example $e^{-i\,A\,dt}\,v$ for the propagation of a quantum state. Applies the result of an arbitrary abstract operator function $f(A\,dt)$ to an abstract state $v$, where $A$ is an abstract operator and $dt$ is a scalar factor. The intended application is computing a single time step $dt$ of the time evolution of a quantum state, where the form of the time evolution operator/dynamical map is $e^{-i\,A\,dt}$. If $v$ is a Hilbert space state, $A$ is the system Hamiltonian (constant over the duration :math:`dt`). If $v$ is a density matrix, $A = i \Liouville$ with the Liouvillian super-operator $\Liouville$. The operator $f(A\,dt)$ is never explicitly constructed, only its application to $v$ is calculated. This distinguishes the Newton propagator from functions like :func:`scipy.linalg.expm` for matrix exponentiation. Internally, it uses an iterative restarted Krylov projection of $(A\,dt)$, where the Krylov dimension in each iteration is `m_max`. Consequently, storage for only `m_max` Arnoldi-vectors of the same size as `v` is required. Args: A (callable): Function encoding the abstract operator $A$. Calling ``A(v)`` must return the result of applying $A$ to $v$. v: Initial state $v$. dt (float): scalar parameter (time step :math:`dt`) func (callable or str): The scalar version of the operator-function $f$. Must take one complex argument and return a complex value. Note that `func` will only ever be called with a scalar argument, not with the argument $(A\,dt)$. If passed as a string, `func` may have the following values, corresponding to "standard" functions: * ``'exp'``: ``lambda x: np.exp(x)`` * ``'expmi'``: ``lambda x: np.exp(-1j * x)`` * ``'expi'``: ``lambda x: np.exp(1j * x)`` m_max (int): Maximal Krylov dimension maxrestart (int): Maximal number of Newton restarts (iterations of the algorithm) tol (float): Desired precision of the result inner (callable): Function that evaluates an inner product. Must take two arguments of the type of `v0` and return a complex number. If None, defaults to :func:`numpy.vdot`. norm (callable): Function that calculates the norm of a state. Must take one argument of the type of `v0` and return a real number. If None, defaults to :func:`numpy.linalg.norm`. The norm must be induced by the inner product specified by `inner`. zero (callable): Function that takes `v` as input and allocates and returns a new zero state of the same type and shape as `v`. If None, defaults to ``lambda v: numpy.zeros(shape=v.shape, dtype=v.dtype)`` Returns: The result of the operator-function $f(A\,dt)$ applied to $v$. The return value will be of the same type as the input `v`. .. Note:: The input `state` may be of any array-like type. It should have (but is not strictly required to have) a `shape` attribute. It must support being multiplied or divided by a scalar, e.g. ``state / 2.0`` or ``2 * state``. Lastly, it must support in-place addition with another state, ``state += state2``. All other mathematical properties of `state` can be defined via custom `inner`, `norm`, and `zero` functions. Mathematically, `state` must be an element a Hilbert space (a "complete inner product space"). This includes density matrices, which can be interpreted as elements of a Hilbert space provided one chooses an appropriate norm and inner product. .. Warning:: The norm *must* be the one induced by the inner product, .. math:: \Norm{v} \equiv \sqrt{\Braket{v}{v}} The parameters `inner` and `norm` must fulfill this definition. For density matrices, they should be the Hilbert-Schmidt product and the Hilbert-Schmidt norm. The "operator norm" has no associated inner product. If the operator norm is used, density matrices are not elements of a Hilbert space, but of a C* algebra, which would not allow for the evaluation of the propagator. A mismatch between `norm` and `inner` leads to subtle errors that will not be obvious (e.g., a substantial lack of precision) """ if inner is None: inner = np.vdot if norm is None: norm = np.linalg.norm else: if norm is None: norm = lambda v: np.sqrt(np.abs(inner(v, v))) if zero is None: zero = lambda v: np.zeros(shape=v.shape, dtype=v.dtype) funcs = { 'exp': _exp, 'expmi': _expmi, 'expi': _expi, } if isinstance(func, str): func = funcs[func] try: N = np.prod(v.shape) if m_max >= N: raise ValueError("m_max must be smaller than the system dimension") except AttributeError: pass return _step(A, v, dt, func, m_max, maxrestart, tol, inner, norm, zero)
def _step(A, v, dt, func, m_max, maxrestart, tol, inner, norm, zero): logger = logging.getLogger("newton") w = zero(v) # result vector Z = np.zeros(0, dtype=np.complex128) # Leja points a = np.zeros(0, dtype=np.complex128) # Newton coeffs beta = norm(v) v = v / beta for s in range(maxrestart): arnoldi_vecs, Hess, Ritz, m = _arnoldi( A, dt, v, m_max, inner=inner, norm=norm ) if m < m_max: logger.warn( "Arnoldi only returned order %d instead of the " "requested %d", m, m_max, ) if m == 0 and s == 0: # The input state appears to be an eigenstate eig_val = beta * Hess[0, 0] phase = func(eig_val) # dt is absorbed in eig_val w = phase * v break # normalize Ritz points if s == 0: radius = _normalize_points(Ritz) center = 0.0 assert radius > 0.0, "Radius is zero" # get Leja points (i.e. Ritz points in the proper order) n_s = len(Z) Z, m = _extend_leja(Z, Ritz, m) # Z now contains m new Leja points assert m > 0, "No new Leja points" a = _extend_newton_coeffs(func, a, Z, center, radius) R = np.matrix(np.zeros(shape=(m + 1, 1), dtype=np.complex128)) R[0, 0] = beta P = a[n_s] * R for k in range(1, m): # 1..m-1 R = (np.dot(Hess, R) - Z[n_s + k - 1] * R) / radius P += a[n_s + k] * R wp = zero(v) for i in range(m): # 0 .. m-1 wp += P[i, 0] * arnoldi_vecs[i] w += wp # starting vector for next iteration R = (np.dot(Hess, R) - Z[n_s + m - 1] * R) / radius beta = np.linalg.norm(R) R /= beta # beta would be the norm of v, with the above normalization, v will now # be normalized v = zero(v) for i in range(m + 1): # 0 .. m v += R[i, 0] * arnoldi_vecs[i] if beta * abs(a[-1]) / (1 + norm(w)) < tol: if logger.isEnabledFor(logging.DEBUG): # pragma: nocover logger.debug("Converged at restart %s", s) logger.debug("norm of wp : %s", norm(wp)) logger.debug("norm of w : %s", norm(w)) logger.debug("beta : %s", beta) logger.debug( "|R*a[-1]|/|w| : %s", np.linalg.norm(R) * a[-1] / norm(w) ) logger.debug("max Leja radius: %s", np.max(np.abs(Z))) break try: assert not np.isnan(np.sum(v)), "v contains NaN" assert not np.isnan(np.sum(w)), "w contains NaN" except (AttributeError, TypeError): pass if s == maxrestart - 1: # pragma: nocover logger.warn("DID NOT REACH CONVERGENCE") logger.warn("increase number of restarts") return w