wavepacket 0.1.0__py3-none-any.whl

wavepacket/__init__.py

@@ -0,0 +1,25 @@
+"""
+A package for the propagation of quantum-mechanical wave functions.
+"""
+
+__version__ = "0.1.0"
+
+from . import builder
+from . import expression
+from . import grid
+from . import operator
+from . import solver
+from . import testing
+from . import typing
+
+
+from .exceptions import (BadFunctionCall, BadGridError, BadStateError, ExecutionError,
+                         InvalidValueError)
+from .generators import Gaussian, PlaneWave
+from .logging import log
+
+
+__all__ = ['__version__', 'log', 'BadFunctionCall', 'BadGridError', 'BadStateError',
+           'ExecutionError', 'InvalidValueError',
+           'Gaussian', 'PlaneWave',
+           'builder', 'expression', 'grid', 'operator', 'solver', 'testing', 'typing']

wavepacket/builder/__init__.py

@@ -0,0 +1,8 @@
+"""
+Functions to assemble wave functions or density operators.
+"""
+
+__all__ = ['direct_product', 'pure_density', 'product_wave_function']
+
+from .density import direct_product, pure_density
+from .wave_function import product_wave_function

wavepacket/builder/density.py

@@ -0,0 +1,73 @@
+import numpy as np
+
+import wavepacket as wp
+from ..grid import State
+
+
+def pure_density(psi: State) -> State:
+    """
+    Given an input wave function, create the corresponding pure density operator.
+
+    This function only performs the direct product, it does not apply
+    further modifications like normalizations.
+
+    Parameters
+    ----------
+    psi : wp.grid.State
+         The input wave function
+
+    Returns
+    -------
+    wp.grid.State
+        The corresponding density operator.
+
+    Raises
+    ------
+    wp.BadStateError
+        If the input is not a valid wave function.
+
+    See also
+    --------
+    direct_product : This function is identical to `direct_product(psi, psi)`
+    """
+    return direct_product(psi, psi)
+
+
+def direct_product(ket: State, bra: State) -> State:
+    """
+    Returns the direct product of wave functions as a density operator.
+
+    Given two wave functions :math:`\psi, \phi`, this function returns the
+    density operator as :math:`| \psi \\rangle\langle \phi |`.
+    This operation can be useful to build up a
+    density operator piece by piece.
+
+    Parameters
+    ----------
+    ket : wp.grid.State
+         The ket state :math:`\psi`
+    bra : wp.grid.State
+         The bra state :math:`\phi`. Note that the function performs a
+         complex conjugation of this state prior to multiplication.
+
+    Returns
+    -------
+    wp.grid.State
+        The direct product of the two states.
+
+    Raises
+    ------
+    wp.BadStateError
+        If one of the input states is not a valid wave function.
+    wp.BadGridError
+        If the input states are defined on different grids.
+    """
+    if not ket.is_wave_function() or not bra.is_wave_function():
+        raise wp.BadStateError("Density operator can only be constructed from wave functions.")
+
+    if ket.grid != bra.grid:
+        raise wp.BadGridError("Grid for bra and ket states does not match")
+
+    rho_matrix = np.outer(ket.data, np.conj(bra.data))
+
+    return State(ket.grid, np.reshape(rho_matrix, ket.grid.operator_shape))

wavepacket/builder/wave_function.py

@@ -0,0 +1,63 @@
+from collections.abc import Sequence
+from typing import Iterable
+
+import numpy as np
+
+import wavepacket as wp
+import wavepacket.typing as wpt
+from ..grid import Grid, State
+
+
+def product_wave_function(grid: Grid,
+                          generators: wpt.Generator | Sequence[wpt.Generator],
+                          normalize: bool = True) -> State:
+    """
+    Builds a product wave function from a set of one-dimensional wave functions.
+
+    Parameters
+    ----------
+    grid : wp.grid.Grid
+           The grid on which the product wave function is assembled
+    generators : Sequence[wp.typing.Generator]
+                A list of callables that specifies the wave function
+                along each degree of freedom. The `generators` return
+                the one-dimensional functions in the DVR, i.e., raw function
+                values at the grid points.
+    normalize : bool, default=true
+               If the norm is non-zero and this value is set, the resulting
+               product wave function is normalized, otherwise the product
+               is returned directly.
+
+    Returns
+    -------
+    wp.grid.State
+        The product wave function in the Wavepacket-default weighted DVR.
+
+    Raises
+    ------
+    wp.InvalidValueError
+        If the number of generators does not match the grid dimensions.
+    """
+    generator_list = generators
+    if not isinstance(generator_list, Iterable):
+        generator_list = [generators]
+
+    if len(generator_list) != len(grid.dofs):
+        raise wp.InvalidValueError(
+            "To build a wave function, you need as many generators as degrees of freedoms."
+            f"Given {len(generator_list)} generators for {len(grid.dofs)} DOFs.")
+
+    result_data = np.ones(grid.shape, dtype=complex)
+    for dof_index, generator in enumerate(generator_list):
+        dof = grid.dofs[dof_index]
+        array = generator(dof.dvr_points)
+        array = dof.from_dvr(array, 0)
+
+        result_data *= grid.broadcast(array, dof_index)
+
+    result = State(grid, result_data)
+    norm = np.sqrt(wp.grid.trace(result))
+    if normalize and norm > 0:
+        return result / norm
+    else:
+        return result

wavepacket/exceptions.py

@@ -0,0 +1,48 @@
+class BadFunctionCall(Exception):
+    """
+    Signals that a function was called incorrectly.
+
+    A typical but rare use-case would be a function that was
+    called with incorrect parameters.
+    """
+    pass
+
+
+class BadGridError(Exception):
+    """
+    An invalid grid was supplied.
+
+    Most often, you attempt an operation between objects that must be defined on the
+    same grid. For example, the addition of two operators defined on different grids
+    is not a useful operation. The grid may also miss required properties, for example
+    an operation may expect a specific degree of freedom type along some index.
+    """
+    pass
+
+
+class BadStateError(Exception):
+    """
+    An invalid state was supplied.
+
+    Either the state is completely invalid (neither wave function nor density operator),
+    or you supplied the wrong type of state to a function, for example passing a
+    density operator where a wave function was required.
+    """
+    pass
+
+
+class ExecutionError(Exception):
+    """
+    An unrecoverable problem was encountered in foreign code.
+
+    The main example is the :py:class:`wavepacket.solver.odesolver` getting
+    an error back while integrating.
+    """
+    pass
+
+
+class InvalidValueError(Exception):
+    """
+    A function argument was incorrect, for example out of bounds.
+    """
+    pass

wavepacket/expression/__init__.py

@@ -0,0 +1,10 @@
+"""
+Classes that wrap operators into expressions for use in partial differential equations.
+"""
+
+__all__ = ['CommutatorLiouvillian',
+           'ExpressionBase', 'SchroedingerEquation']
+
+from .expressionbase import ExpressionBase
+from .liouvillian import CommutatorLiouvillian
+from .schroedingerequation import SchroedingerEquation

wavepacket/expression/expressionbase.py

@@ -0,0 +1,44 @@
+from abc import ABC, abstractmethod
+from typing import Optional
+
+from ..grid import State
+
+
+class ExpressionBase(ABC):
+    """
+    Base class for expressions.
+
+    By deriving from this class and implementing the method
+    :py:meth:`ExpressionBase.apply`, you can add custom expressions.
+
+    Notes
+    -----
+    All differential equations have the form
+    :math:`\dot \\rho = \mathcal{L}(\\rho)` (or equivalently
+    :math:`\dot \psi = \hat H \psi`), that is, the left-hand side
+    is just the time derivative. This matches the common convention for the
+    Liouville von-Neumann equation, but differs from the usual
+    form of the Schrödinger equation, where the imaginary factor
+    is on the left-hand side of the equation.
+    """
+
+    @abstractmethod
+    def apply(self, state: State, t: Optional[float] = None) -> State:
+        """
+        Applies the expression to the input state and returns the result.
+
+        Parameters
+        ----------
+        state : wp.grid.State
+            The state on which the expression is applied.
+        t : float, optional
+            The time at which the expression is evaluated. Default is None,
+            which will raise an exception if the contained expression is time-dependent.
+
+        Raises
+        ------
+        wp.BadStateError
+            If the state is invalid or has the wrong time. For example, a Schroedinger equaiton
+            makes little sense for a density operator.
+        """
+        pass

wavepacket/expression/liouvillian.py

@@ -0,0 +1,64 @@
+from typing import Optional
+
+import wavepacket as wp
+from .expressionbase import ExpressionBase
+from ..grid import State
+from ..operator import OperatorBase
+
+
+class CommutatorLiouvillian(ExpressionBase):
+    """
+    Represents a commutator expression in a Liouville von-Neumann equation.
+
+    Given an operator `H`, this commutator expression is given by
+    :math:`\mathcal{L}(\hat \\rho) = -\imath (\hat H \hat \\rho - \hat \\rho \hat H)`.
+
+    Parameters
+    ----------
+    op : wp.operator.OperatorBase
+        The operator to commute with the density operator.
+
+    Notes
+    -----
+    The extra factor of -i is added to ensure that the commutator can be directly
+    plugged into a Liouville von-Neumann equation. defined as
+    :math:`\\frac{\partial \hat \\rho}{\partial t} = \mathcal{L}(\hat \\rho)`.
+    If you need the raw commutator, you have to multiply the result with
+    the imaginary number.
+    """
+
+    def __init__(self, op: OperatorBase):
+        self._op = op
+
+    def apply(self, rho: State, t: Optional[float] = None) -> State:
+        """
+        Evaluates the commutator  for the given density operator and time.
+
+        Parameters
+        ----------
+        rho : wp.grid.State
+            The density operator to commute with
+        t : float, optional
+            The time at which the operator is evaluated. By default it is None,
+            which will throw if the contained operator is time-dependent.
+
+        Returns
+        -------
+        wp.grid.State
+            The result of the commutator.
+
+        Raises
+        ------
+        wp.BadGridError
+            If the grids of the density operator and the wrapped operator do not match.
+        wp.BadStateError
+            If the input state is not a valid density operator.
+        """
+        if rho.grid != self._op.grid:
+            raise wp.BadGridError("Input state is defined on the wrong grid.")
+
+        if not rho.is_density_operator():
+            raise wp.BadStateError("CommutatorLiouvillian requires a density operator.")
+
+        return State(rho.grid,
+                     -1j * (self._op.apply_from_left(rho.data, t) - self._op.apply_from_right(rho.data, t)))

wavepacket/expression/schroedingerequation.py

@@ -0,0 +1,65 @@
+from typing import Optional
+
+import wavepacket as wp
+from .expressionbase import ExpressionBase
+from ..grid import State
+from ..operator import OperatorBase
+
+
+class SchroedingerEquation(ExpressionBase):
+    """
+    Expression wrapper for a Schrödinger equation.
+
+    You should wrap a Hamiltonian in an object of this type
+    so that the solvers can subsequently solve the resulting
+    equation.
+
+    Parameters
+    ----------
+    hamiltonian : wp.operator.OperatorBase
+        The Hamiltonian that is wrapped by this class.
+
+    Notes
+    -----
+    The Schrödinger equation is given by
+    :math:`\dot \psi = -\imath \hat H \psi`,
+    so this class only multiplies the wave function with the
+    negative imaginary number, and wraps the input Hamiltonian
+    into an expression so that solvers can work with it.
+    """
+
+    def __init__(self, hamiltonian: OperatorBase):
+        self._hamiltonian = hamiltonian
+
+    def apply(self, psi: State, t: Optional[float] = None) -> State:
+        """
+        Evaluates the right-hand side of the Schrödinger equation
+        for the given input state and time.
+
+        Parameters
+        ----------
+        psi : wp.grid.State
+             The input state that is evaluated.
+        t : float, optional
+            The time at which the Schrödinger equation is evaluated.
+            The default `None` throws if the wrapped operator is time-dependent.
+
+        Returns
+        -------
+        wp.grid.State
+            The result of the evaluation.
+
+        Raises
+        ------
+        wp.grid.BadGridError
+            If the state's grid differs from that of the Hamiltonian.
+        wp.grid.BadStateError
+            If the input state is not a wave function.
+        """
+        if psi.grid != self._hamiltonian.grid:
+            raise wp.BadGridError("Input state has wrong grid.")
+
+        if not psi.is_wave_function():
+            raise wp.BadStateError("SchroedingerEquation requires a wave function.")
+
+        return State(psi.grid, -1j * self._hamiltonian.apply_to_wave_function(psi.data, t))

wavepacket/generators.py

@@ -0,0 +1,88 @@
+from typing import Optional
+
+import numpy as np
+
+import wavepacket as wp
+import wavepacket.typing as wpt
+
+
+class Gaussian(wpt.Generator):
+    """
+    Callable that defines a one-dimensional Gaussian function.
+
+    This callable can be supplied wherever a callable is required. An example
+    would be an initial wave function for :py:func:`wavepacket.builder.product_wave_function`,
+    or a potential wrapped in a :py:class:`wavepacket.operator.Potential1D`.
+
+    Parameters
+    ----------
+    x : float, default=0
+        The center of the Gaussian.
+    p : float, default=0
+        The momentum of the Gaussian.
+    rms, fwhm : float
+        You must specify the width of the Gaussian using exactly one of these values,
+        either through the root-mean-square width, or the full-width-at-half-maximum.
+
+    Raises
+    ------
+    wp.InvalidValueError
+        If the width of the Gaussian is not positive.
+    wp.BadFunctionCall
+        If both rms and fwhm have either been set or not supplied.
+
+    Notes
+    -----
+
+    Up to scaling, the functional form of the Gaussian is
+    :math:`f(x) = e^{-(x-x_0)^2 / 2 \sigma^2 + \imath p (x-x_0)}`.
+    Here, sigma is the rms width, which is connected to the FWHM by
+    :math:`\sigma = \mathrm{FWHM} / \sqrt{8 \ln 2}`.
+    """
+
+    def __init__(self, x: float = 0.0, p: float = 0.0,
+                 rms: Optional[float] = None, fwhm: Optional[float] = None):
+        if rms is not None and rms <= 0:
+            raise wp.InvalidValueError(f"RMS width of Gaussian is {rms}, but should be positive.")
+
+        if fwhm is not None and fwhm <= 0:
+            raise wp.InvalidValueError(f"FWHM of Gaussian is {rms}, but should be positive.")
+
+        if fwhm is not None and rms is not None:
+            raise wp.BadFunctionCall("Only one of RMS width or FWHM must be set, not both.")
+        if fwhm is None and rms is None:
+            raise wp.BadFunctionCall("One of RMS width or FWHM must be set.")
+
+        self._x = x
+        self._p = p
+        if rms:
+            self._rms = rms
+        else:
+            self._rms = fwhm / np.sqrt(8 * np.log(2))
+
+    def __call__(self, x: wpt.RealData) -> wpt.ComplexData:
+        shifted = x - self._x
+        arg = - shifted ** 2 / (2 * self._rms ** 2) + 1j * self._p * shifted
+        return np.exp(arg)
+
+
+class PlaneWave:
+    """
+    Callable that defines a plane wave.
+
+    You will typically use this callable for initial states. There are often
+    better options, especially if your FBR already defines a plane wave basis,
+    but sometimes you may just want to represent a reasonable plane wave and
+    not count indices to get the correct wave vector.
+
+    Parameters
+    ----------
+    k : float
+        The wave vector of the plane wave.
+    """
+
+    def __init__(self, k: float):
+        self._k = k
+
+    def __call__(self, x: wpt.RealData) -> wpt.ComplexData:
+        return np.exp(1j * self._k * x)

wavepacket/grid/__init__.py

@@ -0,0 +1,12 @@
+"""
+This module contains the classes to define a grid and represent states on it.
+"""
+
+__all__ = ['DofBase', 'Grid', 'PlaneWaveDof', 'State',
+           'dvr_density', 'trace']
+
+from .dofbase import DofBase
+from .grid import Grid
+from .planewavedof import PlaneWaveDof
+from .state import State
+from .state_utilities import dvr_density, trace