wavepacket 0.1.0__py3-none-any.whl

wavepacket/grid/dofbase.py

@@ -0,0 +1,191 @@
+"""
+Definition of an abstract Degree of Freedom (DOF).
+"""
+
+from abc import abstractmethod, ABC
+
+import wavepacket as wp
+import wavepacket.typing as wpt
+
+
+class DofBase(ABC):
+    """
+    Abstract base class of a one-dimensional basis expansion.
+
+    We assemble a multidimensional grid as direct product of one-dimensional grids.
+    To distinguish these two types of grids, we call the one-dimensional grids DOF
+    (Degree of Freedom). They are defined by a basis expansion (the FBR), its corresponding
+    grid definition in real space (the DVR, see :doc:`/representations`),
+    the transformation between the two representations, and
+    some descriptive quantum numbers for the FBR.
+
+    Note that the transformation functions are highly flexible, but awkward and error-prone
+    to use, so they should generally be avoided outside of Wavepacket-internal code.
+    In general, you should use convenience functions instead that perform the required
+    transformation behind the scenes, such as :py:func:`wavepacket.grid.dvr_density`.
+
+    Parameters
+    ----------
+    dvr_points : real-valued array_like
+        The grid points in real space.
+    fbr_points : real-valued array_like
+        This is typically used to assign useful numbers to the underlying basis.
+        For example, in a plane-wave expansion, we would use the wave vectors as fbr_grid.
+        The size must match that of the `dvr_points`.
+
+    Attributes
+    ----------
+    dvr_points
+    fbr_points
+    size
+
+    Raises
+    ------
+    wp.InvalidValueError
+        If the input grids are empty, multidimensional, or if the FBR and DVR grid do not match in size.
+
+    See Also
+    --------
+    wavepacket.grid.Grid : The multidimensional grid formed from one or more degrees of freedom.
+    wavepacket.grid.PlaneWaveDof: An implementation of a degree of freedom
+                             based on a plane wave expansion.
+    """
+
+    def __init__(self, dvr_points: wpt.RealData, fbr_points: wpt.RealData):
+        if len(dvr_points) == 0 or len(fbr_points) == 0:
+            raise wp.InvalidValueError("Degrees of freedom may not be empty.")
+
+        if dvr_points.ndim != 1 or fbr_points.ndim != 1:
+            raise wp.InvalidValueError("A degree of freedom represents only one-dimensional data.")
+
+        if dvr_points.size != fbr_points.size:
+            raise wp.InvalidValueError("The DVR and FBR grids must have the same size.")
+
+        self._dvr_points = dvr_points
+        self._fbr_points = fbr_points
+
+    @property
+    def dvr_points(self) -> wpt.RealData:
+        """
+        Numpy array that gives the grid points in real space as supplied in the constructor.
+        """
+        return self._dvr_points
+
+    @property
+    def fbr_points(self) -> wpt.RealData:
+        """
+        Numpy array that gives the grid points of the underlying basis as supplied in the constructor.
+        """
+        return self._fbr_points
+
+    @property
+    def size(self) -> int:
+        """
+        The size of the grid (number of elements of `dvr_points`/`fbr_points`)
+        """
+        return self._dvr_points.size
+
+    @abstractmethod
+    def to_fbr(self, data: wpt.ComplexData, index: int, is_ket: bool = True) -> wpt.ComplexData:
+        """
+        Translates a dimension of the input coefficients from the Wavepacket-default "weighted DVR"
+        into the FBR.
+
+        This function is not meant for public use. It does not handle errors explicitly,
+        and is just awkward to use; you need to reach through the state abstraction and transform
+        each index correctly.
+
+        Parameters
+        ----------
+        data : wp.typing.ComplexData
+            The input coefficients of the state to transform.
+        index : int
+            The index of the coefficient array that should be transformed.
+        is_ket : bool, default=True
+            If the index is the coefficient for a ket state (True) or a bra state.
+
+        Returns
+        -------
+        wpt.ComplexData
+            The appropriately transformed coefficients. You will generally not wrap the results
+            into a :py:class:`wavepacket.grid.State`, because that class implicitly assumes a
+            weighted DVR transformation.
+        """
+        pass
+
+    @abstractmethod
+    def from_fbr(self, data: wpt.ComplexData, index: int, is_ket: bool = True) -> wpt.ComplexData:
+        """
+        Translates a dimension of the input coefficients from the FBR  into the Wavepacket-default
+        "weighted DVR"
+
+        This function is not meant for public use. It does not handle errors explicitly,
+        and is just awkward to use; you need to reach through the state abstraction and transform
+        each index correctly.
+
+        Parameters
+        ----------
+        data : wp.typing.ComplexData
+            The input coefficients of the state to transform.
+        index : int
+            The index of the coefficient array that should be transformed.
+        is_ket : bool, default=True
+            If the index is the coefficient for a ket state (True) or a bra state.
+
+        Returns
+        -------
+        wpt.ComplexData
+            The appropriately transformed coefficients. You generally need to wrap the result in
+            a :py:class:`wavepacket.grid.State` before further use.
+        """
+        pass
+
+    @abstractmethod
+    def to_dvr(self, data: wpt.ComplexData, index: int) -> wpt.ComplexData:
+        """
+        Translates a dimension of the input coefficients from the Wavepacket-default "weighted DVR"
+        into the DVR.
+
+        This function is not meant for public use. It does not handle errors explicitly,
+        and is just awkward to use; you need to reach through the state abstraction and transform
+        each index correctly.
+
+        Parameters
+        ----------
+        data : wp.typing.ComplexData
+            The input coefficients of the state to transform.
+        index : int
+            The index of the coefficient array that should be transformed.
+
+        Returns
+        -------
+        wpt.ComplexData
+            The appropriately transformed coefficients. You will generally not wrap the results
+            into a :py:class:`wavepacket.grid.State`, because that class implicitly assumes a
+            weighted DVR transformation.
+        """
+        pass
+
+    @abstractmethod
+    def from_dvr(self, data: wpt.ComplexData, index: int) -> wpt.ComplexData:
+        """
+        Translates a dimension of the input coefficients from the DVR into the  Wavepacket-default "weighted DVR".
+
+        This function is not meant for public use. It does not handle errors explicitly,
+        and is just awkward to use; you need to reach through the state abstraction and transform
+        each index correctly.
+
+        Parameters
+        ----------
+        data : wp.typing.ComplexData
+            The input coefficients of the state to transform.
+        index : int
+            The index of the coefficient array that should be transformed.
+
+        Returns
+        -------
+        wpt.ComplexData
+            The appropriately transformed coefficients. They need to be wrapped in a
+            :py:class:`wavepacket.grid.State` before further use.
+        """
+        pass

wavepacket/grid/grid.py

@@ -0,0 +1,145 @@
+from collections.abc import Sequence
+import math
+
+import numpy as np
+
+import wavepacket as wp
+import wavepacket.typing as wpt
+from .dofbase import DofBase
+
+
+class Grid:
+    """
+    Definition of a one- or multidimensional grid.
+
+    This class collects multiple :py:class:`wavepacket.grid.DofBase`-derived objects,
+    each corresponding to a one-dimensional basis expansion, and represents the resulting
+    one- or multidimensional grid that you can operate with.
+
+    Parameters
+    ----------
+    dofs : Sequence[wp.grid.DofBase] | DofBase
+        The degree(s) of freedom that make up the grid. The order determines the order of
+        the coefficient array indices.
+
+    Attributes
+    ----------
+    shape
+    operator_shape
+    dofs
+    size
+
+    Raises
+    ------
+    wp.InvalidValueError
+        If no degrees of freedom are supplied.
+    """
+
+    def __init__(self, dofs: Sequence[DofBase] | DofBase):
+        if dofs is None:
+            raise wp.InvalidValueError("A grid needs at least one Degree of freedom defined.")
+
+        if isinstance(dofs, DofBase):
+            self._dofs = [dofs]
+        else:
+            self._dofs = list(dofs)
+
+        self._shape = tuple(dof.size for dof in self._dofs)
+
+    @property
+    def shape(self) -> tuple[int, ...]:
+        """
+        The dimensions of the grid, for use in e.g. array creation.
+        """
+        # note: dvr shape
+        return self._shape
+
+    @property
+    def operator_shape(self) -> tuple[int, ...]:
+        """
+        The dimensions of an operator matrix on this grid.
+
+        An example of such a matrix would be the coefficient array for a density operator.
+        The operator dimensions are the dimensions of the grid concatenated with themselves.
+        For example, a grid with dimensions (5, 4) has operator dimensions (5, 4, 5, 4).
+        """
+        return self._shape + self._shape
+
+    @property
+    def dofs(self) -> Sequence[DofBase]:
+        """
+        The vector of degrees of freedom that make up the grid.
+        """
+        return self._dofs
+
+    @property
+    def size(self) -> int:
+        """
+        The total number of grid points.
+        """
+        return math.prod(dof.size for dof in self._dofs)
+
+    def normalize_index(self, index: int) -> int:
+        """
+        Maps indices on the interval [0, rank_of_grid].
+
+        Meant for Wavepacket-internal use.
+
+        For operator matrices, we want to address bra and ket indices separately,
+        and with normal Python convention, i.e., -n meaning the n-th entry from the end.
+        However, an operator matrix has 2N dimensions, with the first N denoting bra indices,
+        and the last N denoting ket indices. Then, the arithmetic becomes cumbersome unless we
+        first map the input index onto the range [0,N].
+        """
+        if index < -len(self._dofs) or index >= len(self._dofs):
+            raise IndexError("Index of degree of freedom out of bounds.")
+
+        if index < 0:
+            return index + len(self._dofs)
+        else:
+            return index
+
+    def broadcast(self, data: wpt.ComplexData, index: int) -> wpt.ComplexData:
+        """
+        Transforms a 1D array into a more suitable form for scaling.
+
+        Meant for Wavepacket-internal use. Note that this function is rather slow,
+        and should not be used in tight loops.
+
+        This function has a very specific purpose. Imagine, you have a grid
+        with shape (5, 4, 3). On this grid, you have a wave function specified by a coefficient
+        array "a" of the same shape. Now you define a potential along the second degree of freedom only.
+        The potential is given by a one-dimensional array "V" of size 4.
+        If you want to apply this potential to the wave function within the DVR approximation,
+        the new coefficients are given as :math:`b_{ijk} = V_j a_{ijk}`.
+
+        Unfortunately, Numpy does not offer a function for this scaling operation.
+        What we can do instead is to blow up the array V into a 3D array
+        of shape (1, 4, 1). then you can map the above multiplication onto Numpy's
+        broadcasting rules. This reshaping is done by this function.
+        """
+        # Note: rather slow, only use for precomputation
+        new_shape = len(self._dofs) * [1]
+        new_shape[index] = self._dofs[index].size
+        return np.reshape(data, new_shape)
+
+    def operator_broadcast(self, data: wpt.ComplexData, dof_index: int, is_ket: bool = True) -> wpt.ComplexData:
+        """
+        Similar to broadcast, but blows up the array into a form suitable for multiplication with operators.
+
+        Meant for Wavepacket-internal use. Note that this function is rather slow,
+        and should not be used in tight loops.
+
+        See :py:meth:`broadcast` for a description of the problem. For the (5, 4, 3)
+        grid shape, this function would blow up the potential array into a shape
+        (1, 4, 1, 1, 1, 1) or (1, 1, 1, 1, 4, 1). The is_ket parameter switches between the
+        two variants, we call the first three indices "ket" and the latter three "bra" indices.
+        """
+        new_shape = (2 * len(self._dofs)) * [1]
+
+        shape_index = self.normalize_index(dof_index)
+        if not is_ket:
+            shape_index += len(self._dofs)
+
+        new_shape[shape_index] = self._dofs[dof_index].size
+        return np.reshape(data, new_shape)

wavepacket/grid/planewavedof.py

@@ -0,0 +1,102 @@
+import math
+
+import numpy as np
+
+import wavepacket as wp
+import wavepacket.typing as wpt
+from .dofbase import DofBase
+
+
+def _broadcast(data: wpt.ComplexData, ndim, index: int) -> wpt.ComplexData:
+    shape = np.ones(ndim, dtype=int)
+    shape[index] = len(data)
+
+    return np.reshape(data, shape)
+
+
+class PlaneWaveDof(DofBase):
+    """
+    One-dimensional plane wave basis expansion.
+
+    This expansion is a good base choice for non-rotational degrees of freedom.
+    The DVR grid consists of equally-spaced grid points from xmin to (xmax - dx),
+    the FBR grid are the wave vectors of the plane waves centered around 0.
+    In the FBR, derivatives become simple multiplications wih the wave vectors, so
+    this grid allows a simple representation of kinetic energy operators.
+
+    The transformation between DVR and FBR uses a FFT, so the performance is ok, even
+    though you might need more grid points than with more suitable expansions.
+
+    Be aware that this degree of freedom implicitly uses periodic boundary conditions
+    in real space. That is, if the wave function leaves the grid on one side, it reenters
+    the grid on the other side. This problem can only be mitigated with negative imaginary
+    potentials. Periodic boundary conditions also hold in the FBR, causing aliasing.
+    Nevertheless, the monitoring of convergence is rather simple, see [1]_.
+
+    Parameters
+    ----------
+    xmin : float
+        The start of the grid.
+    xmax : float
+        The end of the grid. Note that the last grid point is at xmax - dx.
+    n : int
+        The number of grid points.
+
+    Raises
+    ------
+    wp.InvalidValueError
+        If the length of the grid is negative or the number of grid points is non-positive.
+
+    References
+    ----------
+    .. [1] https://sourceforge.net/p/wavepacket/cpp/blog/2020/11/convergence-1-equally-space-grids
+    """
+
+    def __init__(self, xmin: float, xmax: float, n: int):
+        if xmin > xmax:
+            raise wp.InvalidValueError("Range should be positive")
+
+        if n <= 0:
+            raise wp.InvalidValueError(f"Number of grid points must be positive, but is {n}")
+
+        dx = (xmax - xmin) / n
+
+        dvr = np.linspace(xmin, xmax, n, endpoint=False)
+        if n % 2 == 0:
+            fbr = np.linspace(-math.pi / dx, math.pi / dx, n, endpoint=False)
+        else:
+            fbr = np.linspace(-math.pi / dx * (n - 1) / n, math.pi / dx * (n - 1) / n, n)
+
+        super().__init__(dvr, fbr)
+
+        self._sqrt_weights: wpt.RealData = math.sqrt(dx) * np.ones(n)
+        self._phase: wpt.ComplexData = np.exp(-1j * fbr * xmin) / np.sqrt(n)
+        self._conj_phase: wpt.ComplexData = n * np.conj(self._phase)
+
+    def to_fbr(self, data: wpt.ComplexData, index: int, is_ket: bool = True) -> wpt.ComplexData:
+        if is_ket:
+            phase = _broadcast(self._phase, data.ndim, index)
+            transformed = np.fft.fftshift(np.fft.fft(data, axis=index), axes=index)
+        else:
+            phase = _broadcast(self._conj_phase, data.ndim, index)
+            transformed = np.fft.fftshift(np.fft.ifft(data, axis=index), axes=index)
+
+        return phase * transformed
+
+    def from_fbr(self, data: wpt.ComplexData, index: int, is_ket: bool = True) -> wpt.ComplexData:
+        if is_ket:
+            phase = _broadcast(self._conj_phase, data.ndim, index)
+            untransformed = phase * data
+            return np.fft.ifft(np.fft.ifftshift(untransformed, axes=index), axis=index)
+        else:
+            phase = _broadcast(self._phase, data.ndim, index)
+            untransformed = phase * data
+            return np.fft.fft(np.fft.ifftshift(untransformed, axes=index), axis=index)
+
+    def to_dvr(self, data: wpt.ComplexData, index: int) -> wpt.ComplexData:
+        conversion_factor = _broadcast(self._sqrt_weights, data.ndim, index)
+        return data / conversion_factor
+
+    def from_dvr(self, data: wpt.ComplexData, index: int) -> wpt.ComplexData:
+        conversion_factor = _broadcast(self._sqrt_weights, data.ndim, index)
+        return data * conversion_factor

wavepacket/grid/state.py

@@ -0,0 +1,93 @@
+import numbers
+from dataclasses import dataclass
+import typing
+
+import wavepacket as wp
+import wavepacket.typing as wpt
+from .grid import Grid
+
+
+@dataclass(frozen=True)
+class State:
+    """
+    This class holds the definition of a specific quantum state.
+
+    A state can be a wave function or a density operator.
+    While invalid states can be constructed, they have no use.
+    A state is comprised of three parts:
+
+    1. The grid on which the state is defined.
+    2. The expansion coefficients.
+    3. The representation / basis for the expansion.
+       This is always weighted DVR, see :doc:`/representations`.
+
+    Once constructed, a State is meant to be immutable. It supports elementary
+    arithmetic operations, however, to allow easy composition.
+
+    The technical reason for a state class is to store grids together with the coefficients,
+    which greatly simplifies the Wavepacket API.
+
+    Attributes
+    ----------
+    grid
+        The grid on which the state is defined.
+    data
+        The coefficients of the state.
+    """
+    grid: Grid
+    data: wpt.ComplexData
+
+    def is_wave_function(self) -> bool:
+        """
+        Returns whether the state represents a wave function.
+        """
+        return self.data.shape == self.grid.shape
+
+    def is_density_operator(self) -> bool:
+        """
+        Returns whether the state represents a density operator.
+        """
+        return self.data.shape == self.grid.operator_shape
+
+    def __add__(self, other: typing.Self | numbers.Number) -> typing.Self:
+        if isinstance(other, State):
+            self._check_states(other)
+            return State(self.grid, self.data + other.data)
+        else:
+            return State(self.grid, self.data + other)
+
+    def __radd__(self, other: numbers.Number) -> typing.Self:
+        return self + other
+
+    def __sub__(self, other: typing.Self | numbers.Number) -> typing.Self:
+        if isinstance(other, State):
+            self._check_states(other)
+            return State(self.grid, self.data - other.data)
+        else:
+            return State(self.grid, self.data - other)
+
+    def __rsub__(self, other: numbers.Number) -> typing.Self:
+        return State(self.grid, other - self.data)
+
+    def __mul__(self, other: numbers.Number) -> typing.Self:
+        return State(self.grid, self.data * other)
+
+    def __rmul__(self, other: numbers.Number) -> typing.Self:
+        return self * other
+
+    def __truediv__(self, other: numbers.Number) -> typing.Self:
+        if other == 0.0:
+            raise ZeroDivisionError("State cannot be divided by zero.")
+
+        return State(self.grid, self.data / other)
+
+    def __neg__(self) -> typing.Self:
+        return State(self.grid, -self.data)
+
+    def _check_states(self, other: typing.Self) -> None:
+        if self.grid != other.grid:
+            raise wp.BadGridError("Binary operations with states on different grids are not supported.")
+
+        if self.data.shape != other.data.shape:
+            raise wp.BadStateError(
+                "Binary operations can only be performed for two wave_functions or density operators.")

wavepacket/grid/state_utilities.py

@@ -0,0 +1,89 @@
+import numpy as np
+
+import wavepacket as wp
+import wavepacket.typing as wpt
+from .grid import Grid
+from .state import State
+
+
+def _take_diagonal(data: wpt.ComplexData, grid: Grid) -> wpt.RealData:
+    # note: the diagonal should be real and positive but for numerical issues
+    matrix_form = np.reshape(data, (grid.size, grid.size))
+    diagonal = np.abs(np.diag(matrix_form))
+    return np.reshape(diagonal, grid.shape)
+
+
+def dvr_density(state: State) -> wpt.RealData:
+    """
+    Returns the density of the input state at the DVR grid points.
+
+    The density is returned as a real-valued coefficient array
+    with the same shape as the underlying grid. The density is mostly a
+    dead end: useful for plotting and inspection, but not for
+    further computations.
+
+    Parameters
+    ----------
+    state : wp.grid.State
+        The state (wave function or density operator) whose density should be computed.
+
+    Returns
+    -------
+    wpt.RealData
+        The coefficient array of the density values at the DVR grid points.
+
+    Raises
+    ------
+    wp.BadStateError
+        If the supplied state is neither a wave function nor a density operator.
+    """
+    if state.is_wave_function():
+        data = state.data
+        for index, dof in enumerate(state.grid.dofs):
+            data = dof.to_dvr(data, index)
+
+        return np.abs(data * data)
+    elif state.is_density_operator():
+        data = state.data
+        grid = state.grid
+
+        # both, the bra and the ket indices are converted
+        for index, dof in enumerate(grid.dofs):
+            data = dof.to_dvr(data, index)
+            data = dof.to_dvr(data, index + len(grid.dofs))
+
+        return _take_diagonal(data, grid)
+    else:
+        raise wp.BadStateError("Input is not a valid state.")
+
+
+def trace(state: State) -> float:
+    """
+    Returns the trace of the supplied input state.
+
+    For density operators, this is the usual trace norm, i.e., sum over the
+    diagonal elements. For wave functions, it is the square of the usual 
+    L2 norm.
+
+    Parameters
+    ----------
+    state : wp.grid.State
+        The state (wave function or density operator) for which to calculate the trace.
+
+    Returns
+    -------
+    float
+        The trace of the input state.
+
+    Raises
+    ------
+    wp.BadStateError
+        If the supplied state is neither a wave function nor a density operator.
+    """
+    if state.is_wave_function():
+        return np.abs(state.data ** 2).sum()
+    elif state.is_density_operator():
+        diagonal = _take_diagonal(state.data, state.grid)
+        return diagonal.sum()
+    else:
+        raise wp.BadStateError("Input is not a valid state.")

wavepacket/logging.py

@@ -0,0 +1,34 @@
+import math
+import numbers
+
+from .grid import State, trace
+from .operator import Potential1D, expectation_value
+
+
+def log(t: numbers.Real, state: State, precision: int = 6) -> None:
+    """
+    Prints some data about the state for inspection.
+
+    The idea is that you call this function during every solver step and get
+    a log with the most important values about the propagation, for example
+    the state trace (if it deviates from one, this may be caused by poor convergence).
+
+    Parameters
+    ----------
+    t : float
+        The time at which you log.
+    state : wp.grid.State
+        The state to log.
+    precision : int, default=6
+        How many decimal places should be printed.
+    """
+    print(f"\n\nt = {float(t):.{precision}},     trace = {trace(state):.{precision}}\n")
+
+    for index, dof in enumerate(state.grid.dofs):
+        x = Potential1D(state.grid, 0, lambda dvr_grid: dvr_grid)
+        x2 = Potential1D(state.grid, 0, lambda dvr_grid: dvr_grid ** 2)
+
+        x_expect = expectation_value(x, state).real
+        x_expect2 = expectation_value(x2, state).real
+
+        print(f"<x_{index}> = {x_expect:.{precision}}  =/- {math.sqrt(x_expect2 - x_expect ** 2):.{precision}}")

wavepacket/operator/__init__.py

@@ -0,0 +1,12 @@
+"""
+This module contains classes that define operators on a given grid.
+"""
+
+__all__ = ['CartesianKineticEnergy', 'OperatorBase',
+           'PlaneWaveFbrOperator', 'Potential1D',
+           'expectation_value']
+
+from .operatorbase import OperatorBase
+from .operatorutils import expectation_value
+from .fbroperators import CartesianKineticEnergy, PlaneWaveFbrOperator
+from .potentials import Potential1D