wavepacket 0.1.0__py3-none-any.whl

wavepacket/operator/fbroperators.py

@@ -0,0 +1,94 @@
+import numpy as np
+
+import wavepacket as wp
+import wavepacket.typing as wpt
+from .operatorbase import OperatorBase
+from ..grid import Grid
+
+
+class PlaneWaveFbrOperator(OperatorBase):
+    """
+    Base class for operators that are diagonal in a plane wave basis.
+
+    The :py:class:`wavepacket.grid.PlaneWaveDof` describes an expansion of
+    the wave function in plane waves. In the corresponding FBR, derivatives
+    are diagonal and essentially transform into a multiplication with the
+    wave vector / FBR grid.
+
+    A key difference to a generic FBR operator is that this operator
+    does not apply the shift after the FFT, which increases performance.
+
+    Parameters
+    ----------
+    grid : wp.grid.Grid
+        The grid on which the operator is defined.
+    dof_index : int
+        The degree of freedom along which the operator is defined.
+    generator : wpt.Generator
+        A callable that gives the operator value for each FBR point.
+
+    Raises
+    ------
+    wp.InvalidValueError
+        If the supplied degree of freedom is not a plane wave expansion.
+    """
+
+    def __init__(self, grid: Grid, dof_index: int, generator: wpt.Generator):
+        if not isinstance(grid.dofs[dof_index], wp.grid.PlaneWaveDof):
+            raise wp.InvalidValueError(
+                f"PlaneWaveFbrOperator requires a PlaneWaveDof, but got {grid.dofs[dof_index].__class__}")
+
+        self._wf_index = dof_index
+        self._ket_index = grid.normalize_index(dof_index)
+        self._bra_index = self._ket_index + len(grid.dofs)
+
+        # shifting the data here allows us to skip the fftshift() on the input data in apply*()
+        data = generator(grid.dofs[dof_index].fbr_points)
+        shifted_data = np.fft.ifftshift(data)
+        self._wf_data = grid.broadcast(shifted_data, dof_index)
+        self._ket_data = grid.operator_broadcast(shifted_data, dof_index)
+        self._bra_data = grid.operator_broadcast(shifted_data, dof_index, is_ket=False)
+
+        super().__init__(grid)
+
+    def apply_to_wave_function(self, psi: wpt.ComplexData, t: float) -> wpt.ComplexData:
+        psi_fft = np.fft.fft(psi, axis=self._wf_index)
+        return np.fft.ifft(psi_fft * self._wf_data, axis=self._wf_index)
+
+    def apply_from_left(self, rho: wpt.ComplexData, t: float) -> wpt.ComplexData:
+        rho_fft = np.fft.fft(rho, axis=self._ket_index)
+        return np.fft.ifft(rho_fft * self._ket_data, axis=self._ket_index)
+
+    def apply_from_right(self, rho: wpt.ComplexData, t: float) -> wpt.ComplexData:
+        rho_fft = np.fft.ifft(rho, axis=self._bra_index)
+        return np.fft.fft(rho_fft * self._bra_data, axis=self._bra_index)
+
+
+class CartesianKineticEnergy(PlaneWaveFbrOperator):
+    """
+    Convenience class that implements the common Cartesian kinetic energy operator.
+
+    The form of the operator is :math:`-\\frac{1}{2m} \ \\frac{\partial^2}{\partial x^2}`.
+    It requires the degree of freedom to be a :py:class:`wavepacket.grid.PlaneWaveDof`.
+
+    Parameters
+    ----------
+    grid : wp.grid.Grid
+        The grid on which the operator is defined.
+    dof_index : inst
+        Degree of freedom along which the operator acts
+    mass : float
+        The mass of the particle.
+
+    Raises
+    ------
+    wp.InvalidValueError
+        If the mass is not positive, or if the degree of freedom does not describe a
+        plane wave expansion.
+    """
+
+    def __init__(self, grid: Grid, dof_index: int, mass: float):
+        if mass <= 0:
+            raise wp.InvalidValueError(f"Particle mass must be positive, but is {mass}")
+
+        super().__init__(grid, dof_index, lambda fbr_points: fbr_points ** 2 / (2 * mass))

wavepacket/operator/operatorbase.py

@@ -0,0 +1,204 @@
+from abc import ABC, abstractmethod
+from collections.abc import Sequence
+import typing
+
+import numpy as np
+
+import wavepacket as wp
+import wavepacket.typing as wpt
+from ..grid import Grid, State
+
+
+class OperatorBase(ABC):
+    """
+    Base class of an operator.
+
+    An operator can be applied to a wave function or from the left or right to a
+    density operator. It is defined on a grid, and can only operate on states on that
+    grid.
+
+    Parameters
+    ----------
+    grid : wp.grid.Grid
+        The grid on which the operator is defined.
+        Particular operators may require additional parameters.
+
+    Attributes
+    ----------
+    grid
+    """
+
+    def __init__(self, grid: Grid):
+        self._grid = grid
+
+    @property
+    def grid(self):
+        """
+        Returns the grid on which the operator is defined.
+        """
+        return self._grid
+
+    def apply(self, state: State, t: typing.Optional[float] = None) -> State:
+        """
+        Applies the operator onto a wave function or a density operator from the left.
+
+        This is a convenience function if you just want to apply the operator without
+        detailed knowledge of the state.
+
+        Parameters
+        ----------
+        state : wp.grid.State
+            The state that the operator is applied on.
+        t : float, optional
+            The time at which the operator is applied.
+             The default value `None` raises an exception for time-dependent operators.
+
+        Returns
+        -------
+        wp.grid.State
+            The result of applying the operator on the state.
+
+        Raises
+        ------
+        wp.BadGridError
+            If the state's grid does not match the grid of the operator.
+        wp.BadStateError
+            If the state is neither a wave function nor a density operator.
+
+        """
+        if state.grid != self._grid:
+            raise wp.BadGridError("Grid of state does not match grid of operator.")
+
+        if state.is_wave_function():
+            return State(state.grid, self.apply_to_wave_function(state.data, t))
+        elif state.is_density_operator():
+            return State(state.grid, self.apply_from_left(state.data, t))
+        else:
+            raise wp.BadStateError("Cannot apply the operator to an invalid state.")
+
+    def __add__(self, other: typing.Self) -> typing.Self:
+        """
+        Adds two operators and returns the result as a :py:class:`wavepacket.operator.OperatorSum`.
+        """
+        return OperatorSum([self, other])
+
+    @abstractmethod
+    def apply_to_wave_function(self, psi: wpt.ComplexData, t: float) -> wpt.ComplexData:
+        """
+        Applies the operator on a wave function.
+
+        This function is mainly for Wavepacket-internal use. It ignores most error
+        handling (usually done by the wrapping :py:class:`wavepacket.expression.ExpressionBase`),
+        and operates directly on coefficients to avoid the creation of temporary states.
+
+        Parameters
+        ----------
+        psi : wpt.ComplexData
+            The coefficients describing the wave function on which the operator acts.
+        t : float
+            The time at which the operator should be evaluated.
+
+        Returns
+        -------
+        wpt.ComplexData
+            The coefficients of the resulting wave function.
+        """
+        pass
+
+    @abstractmethod
+    def apply_from_left(self, rho: wpt.ComplexData, t: float) -> wpt.ComplexData:
+        """
+        Applies the operator on a density operator from the left.
+
+        This function is mainly for Wavepacket-internal use. It ignores most error
+        handling (usually done by the wrapping :py:class:`wavepacket.expression.ExpressionBase`),
+        and operates directly on coefficients to avoid the creation of temporary states.
+
+        Parameters
+        ----------
+        rho : wpt.ComplexData
+            The coefficients describing the density operator on which the operator acts.
+        t : float
+            The time at which the operator should be evaluated.
+
+        Returns
+        -------
+        wpt.ComplexData
+            The coefficients of the resulting density operator.
+        """
+        pass
+
+    @abstractmethod
+    def apply_from_right(self, rho: wpt.ComplexData, t: float) -> wpt.ComplexData:
+        """
+        Applies the operator on a density operator from the right.
+
+        This function is mainly for Wavepacket-internal use. It ignores most error
+        handling (usually done by the wrapping :py:class:`wavepacket.expression.ExpressionBase`),
+        and operates directly on coefficients to avoid the creation of temporary states.
+
+        Parameters
+        ----------
+        rho : wpt.ComplexData
+            The coefficients describing the density operator on which the operator acts.
+        t : float
+            The time at which the operator should be evaluated.
+
+        Returns
+        -------
+        wpt.ComplexData
+            The coefficients of the resulting density operator.
+        """
+        pass
+
+
+class OperatorSum(OperatorBase):
+    """
+    An operator that represents the sum of multiple other operators.
+
+    You do not normally construct this operator directly. It is the result of
+    adding two or more operators together. All functionality is simply forwarded
+    to the individual operators.
+
+    Parameters
+    ----------
+    ops : Sequence[wp.operator.OperatorBase]
+        The operators that should be summed up.
+
+    Raises
+    ------
+    wp.BadGridError
+        If the operators are defined on different grids.
+    """
+
+    def __init__(self, ops: Sequence[OperatorBase]):
+        if not ops:
+            raise wp.InvalidValueError("OperatorSum needs at least one operator to sum.")
+        for op in ops:
+            if op.grid != ops[0].grid:
+                raise wp.BadGridError("All grids in a sum operator must be equal.")
+
+        self._ops = ops
+        grid = ops[0].grid
+        super().__init__(grid)
+
+    def apply_to_wave_function(self, psi: wpt.ComplexData, t: float) -> wpt.ComplexData:
+        result = np.zeros(self.grid.shape, dtype=np.complex128)
+        for op in self._ops:
+            result += op.apply_to_wave_function(psi, t)
+
+        return result
+
+    def apply_from_left(self, rho: wpt.ComplexData, t: float) -> wpt.ComplexData:
+        result = np.zeros(self.grid.operator_shape, dtype=np.complex128)
+        for op in self._ops:
+            result += op.apply_from_left(rho, t)
+
+        return result
+
+    def apply_from_right(self, rho: wpt.ComplexData, t: float) -> wpt.ComplexData:
+        result = np.zeros(self.grid.operator_shape, dtype=np.complex128)
+        for op in self._ops:
+            result += op.apply_from_right(rho, t)
+
+        return result

wavepacket/operator/operatorutils.py

@@ -0,0 +1,30 @@
+from typing import Optional
+
+import numpy as np
+
+import wavepacket as wp
+from .operatorbase import OperatorBase
+from ..grid import State
+
+
+def expectation_value(op: OperatorBase, state: State, t: Optional[float] = None) -> complex:
+    """
+    Calculates the expectation value of an operator for a given state.
+
+    Parameters
+    ----------
+    op : wp.operator.OperatorBase
+        The operator whose expectation value is calculated.
+    state : wp.grid.State
+        The wave function or density operator that is used for the calculation.
+    t : float, optional
+        The time at which the operator should be evaluated.
+         Only required for time-dependent operators, where the default value `None` raises an exception.
+    """
+    new = op.apply(state, t)
+
+    if state.is_wave_function():
+        return np.vdot(state.data, new.data)
+    else:
+        matrix_data = np.reshape(new.data, [new.grid.size, new.grid.size])
+        return np.trace(matrix_data)

wavepacket/operator/potentials.py

@@ -0,0 +1,43 @@
+import wavepacket.typing as wpt
+from .operatorbase import OperatorBase
+from ..grid import Grid
+
+
+class Potential1D(OperatorBase):
+    """
+    Operator that represents a real-valued one-dimensional potential.
+
+    Within the DVR approximation [1]_, potential energy operators are
+    diagonal in the DVR. That is, they are completely described by a value
+    for each grid point.
+
+    Parameters
+    ----------
+    grid : wp.grid.Grid
+        The grid on which the operator is defined
+    dof_index : int
+        the index of the degree of freedom along which the potential is defined.
+    generator : wpt.RealGenerator
+        A callable that generates a potential energy value for each grid point of the respective DOF.
+
+    References
+    ----------
+    .. [1] https://sourceforge.net/p/wavepacket/wiki/Numerics.DVR>
+    """
+
+    def __init__(self, grid: Grid, dof_index: int, generator: wpt.RealGenerator):
+        data = generator(grid.dofs[dof_index].dvr_points)
+        self._wf_data = grid.broadcast(data, dof_index)
+        self._ket_data = grid.operator_broadcast(data, dof_index)
+        self._bra_data = grid.operator_broadcast(data, dof_index, False)
+
+        super().__init__(grid)
+
+    def apply_to_wave_function(self, psi: wpt.ComplexData, t: float) -> wpt.ComplexData:
+        return self._wf_data * psi
+
+    def apply_from_left(self, rho: wpt.ComplexData, t: float) -> wpt.ComplexData:
+        return self._ket_data * rho
+
+    def apply_from_right(self, rho: wpt.ComplexData, t: float) -> wpt.ComplexData:
+        return self._bra_data * rho

wavepacket/solver/__init__.py

@@ -0,0 +1,8 @@
+"""
+This module contains classes to solve the equations of motion set up by the
+classes in the :py:mod:`wavepacket.expression` module.
+"""
+__all__ = ['OdeSolver', 'SolverBase']
+
+from .odesolver import OdeSolver
+from .solverbase import SolverBase

wavepacket/solver/odesolver.py

@@ -0,0 +1,70 @@
+import numpy as np
+from scipy.integrate import solve_ivp
+
+import wavepacket as wp
+import wavepacket.typing as wpt
+from .solverbase import SolverBase
+from ..grid import Grid, State
+from ..expression import ExpressionBase
+
+
+def _inner_solver(t: float, y: wpt.ComplexData,
+                  eq: ExpressionBase, grid: Grid) -> wpt.ComplexData:
+    # 1. Reconstruct a state from the data array y
+    if len(y) == grid.size:
+        state = State(grid, np.reshape(y, grid.shape))
+    else:
+        state = State(grid, np.reshape(y, grid.operator_shape))
+
+    # 2. Insert into the equation expression
+    dot_state = eq.apply(state, t)
+
+    # 3. Extract the result as a vector again and return
+    return dot_state.data.ravel()
+
+
+class OdeSolver(SolverBase):
+    """
+    A solver that uses scipy's ODE integrators as backend.
+
+    This solver merely wraps `scipy.integrate.solve_ivp()` in a Wavepacket
+    solver. ODESolvers tend to be workhorses that can be applied to most
+    systems without a second thought at the expense of performance.
+    See [1]_ for more details on the convergence of solvers.
+
+    Parameters
+    ----------
+    expr : wp.expression.ExpressionBase
+        The right-hand side of the differential equation that encapsulates the quantum system.
+    dt : float
+        The elementary time step for the time evolution.
+    **kwargs
+        Additional keyword parameters can be supplied that are directly forwarded to scipy.
+        The most important parameters are "method" (defaults to "RK45") and "rtol", "atol"
+        for the relative and absolute error tolerances (booth default to 1e-6).
+
+    References
+    ----------
+    .. [1] https://sourceforge.net/p/wavepacket/cpp/blog/2021/04/convergence-2
+    """
+
+    def __init__(self, expr: ExpressionBase, dt: float, **kwargs):
+        super().__init__(dt)
+
+        self._expression = expr
+        self._kwargs = kwargs
+
+        # The default error tolerance is typically too generous for our use cases. Override
+        self._kwargs.setdefault('rtol', 1e-6)
+
+    def step(self, state: State, t: float) -> State:
+        t_span = (t, t + self.dt)
+        y0 = state.data.ravel()
+        args = (self._expression, state.grid)
+
+        solution = solve_ivp(_inner_solver, t_span, y0, t_eval=[t + self._dt], args=args, **self._kwargs)
+        if solution.status != 0:
+            raise wp.ExecutionError("Bad return value from integrating"
+                                    f"Status: {solution.status} != 0; Message: {solution.msg}")
+
+        return State(state.grid, np.reshape(solution.y, state.data.shape))

wavepacket/solver/solverbase.py

@@ -0,0 +1,119 @@
+from abc import ABC, abstractmethod
+
+import wavepacket as wp
+from ..grid import State
+
+
+class SolverBase(ABC):
+    """
+    Abstract base class for all differential equation solvers.
+
+    Each solver must take in the constructor the size of an
+    elementary time step, and it must implement a method :py:meth:`step`
+    to evolve a state by one elementary time step forward in time.
+
+    Usually, a solver does not need to know what step it evolves or
+    how a differential equation looks in detail. The form of the
+    differential equation is encapsulated in a
+    :py:class:`wavepacket.expression.ExpressionBase` instance that
+    is usually supplied to the specific implementations.
+
+    A solver may, however, implicitly assume properties of the
+    differential equation. It may also take additional arguments,
+    these details are documented in the specific implementations.
+
+    Parameters
+    ----------
+    dt : float
+        The size of an elementary time step.
+
+
+    Attributes
+    ----------
+    dt
+
+    Raises
+    ------
+    wp.InvalidValueError
+        If the timestep is not positive.
+    """
+
+    def __init__(self, dt: float):
+        if dt <= 0:
+            raise wp.InvalidValueError(f"Require positive timestep, got {dt}")
+
+        self._dt = dt
+
+    @property
+    def dt(self) -> float:
+        """
+        Returns the size of the elementary time step.
+        """
+        return self._dt
+
+    @abstractmethod
+    def step(self, state: State, t: float) -> State:
+        """
+        Evolves the given state for one elementary time step.
+
+        Given a wave function or density operator at time t,
+        this function should return the propagated wave function
+        or density operator at the new time t+dt, where dt is
+        the elementary time step.
+
+        Parameters
+        ----------
+        state : wp.grid.State
+            The state to be evolved in time
+        t : float
+            The time at which the time evolution starts
+
+        Returns
+        -------
+        wp.grid.State
+            The state at the new time t+dt.
+        """
+        pass
+
+    def propagate(self, state0: State, t0: float,
+                  num_steps: int, include_first: bool = True):
+        """
+        Generator function that yields the propagated wave function at multiple time steps.
+
+        This function allows you to propagate in one go with a for loop.
+        It repeatedly calls :py:meth:`step` and returns the wave function and current time.
+
+        Parameters
+        ----------
+        state0 : wp.grid.State
+            The initial state to be propagated in time.
+        t0 : float
+            The initial time at which the state is given
+        num_steps : int
+            For how many elementary time steps the state should be propagated.
+        include_first : bool
+            If true, the function will start by yielding the initial state.
+
+        Yields
+        ------
+        float
+            The time at which the state is yielded. Starts optionally with t0 and progresses in units of dt.
+        wp.grid.State
+            The propagated state at the given time.
+
+        Raises
+        ------
+        wp.InvalidValueError
+            If num_steps is negative.
+        """
+        if num_steps < 0:
+            raise wp.InvalidValueError("Cannot propagate for negative number of steps.")
+
+        if include_first:
+            yield t0, state0
+
+        state = state0
+        for step in range(num_steps):
+            t = t0 + step * self._dt
+            state = self.step(state, t)
+            yield t + self._dt, state

wavepacket/testing/__init__.py

@@ -0,0 +1,11 @@
+"""
+Utility classes and functions for writing Wavepacket tests.
+"""
+
+__all__ = ['DummyDof', 'DummyOperator',
+           'assert_close', 'random_state']
+
+from .assertions import assert_close
+from .dummydof import DummyDof
+from .dummyoperator import DummyOperator
+from .random import random_state

wavepacket/testing/assertions.py

@@ -0,0 +1,20 @@
+import numpy.testing
+
+import wavepacket as wp
+
+
+def assert_close(actual: wp.grid.State, expected: wp.grid.State, diff: float = 0) -> None:
+    """
+    Assertion helper: Verify that two states are on the same grid and similar to each other.
+
+    Parameters
+    ----------
+    actual : wp.grid.State
+        The expected state
+    expected : wp.grid.State
+        The state to be tested.
+    diff : float
+        The maximum absolute difference between any coefficients of the two states.
+    """
+    assert actual.grid == expected.grid
+    numpy.testing.assert_allclose(actual.data, expected.data, rtol=0, atol=diff)

wavepacket/testing/dummydof.py

@@ -0,0 +1,22 @@
+import wavepacket.typing as wpt
+from ..grid import DofBase
+
+
+class DummyDof(DofBase):
+    """
+    Empty DOF without transformations.
+    """
+    def __init__(self, dvr_array: wpt.RealData, fbr_array: wpt.RealData):
+        super().__init__(dvr_array, fbr_array)
+
+    def from_fbr(self, data: wpt.ComplexData, index: int, is_ket: bool = True) -> wpt.ComplexData:
+        return data
+
+    def to_dvr(self, data: wpt.ComplexData, index: int) -> wpt.ComplexData:
+        return data
+
+    def from_dvr(self, data: wpt.ComplexData, index: int) -> wpt.ComplexData:
+        return data
+
+    def to_fbr(self, data: wpt.ComplexData, index: int, is_ket: bool = True) -> wpt.ComplexData:
+        return data

wavepacket/testing/dummyoperator.py

@@ -0,0 +1,26 @@
+import wavepacket as wp
+import wavepacket.typing as wpt
+from ..grid import Grid
+from ..operator import OperatorBase
+
+
+class DummyOperator(OperatorBase):
+    """
+    Empty operator that throws when it is applied.
+
+    Used for testing where we need an operator, but do not get
+    as far as actually doing something.
+    """
+    def __init__(self, grid: Grid):
+        super().__init__(grid)
+
+    def apply_to_wave_function(self, psi: wpt.ComplexData, t: float) -> wpt.ComplexData:
+        raise wp.BadFunctionCall("Should be patched.")
+
+    def apply_from_left(self, rho: wpt.ComplexData, t: float) -> wpt.ComplexData:
+        raise wp.BadFunctionCall("Should be patched.")
+
+    def apply_from_right(self, rho: wpt.ComplexData, t: float) -> wpt.ComplexData:
+        raise wp.BadFunctionCall("Should be patched.")
+
+

wavepacket/testing/random.py

@@ -0,0 +1,16 @@
+import numpy as np
+
+import wavepacket as wp
+
+
+def random_state(grid, seed) -> wp.grid.State:
+    """
+    Creates a random wave function.
+
+    Note that this function does not employ high-quality randomization,
+    it is only meant to create states without accidental symmetries or lots of code.
+    """
+    rng = np.random.default_rng(seed)
+    data = rng.random(grid.shape) + 1j * rng.random(grid.shape)
+
+    return wp.grid.State(grid, data)

wavepacket/typing/__init__.py

@@ -0,0 +1,7 @@
+"""
+Several aliases and definitions for type checking.
+"""
+
+__all__ = ['ComplexData', 'RealData', 'Generator', 'RealGenerator']
+
+from .data_types import ComplexData, RealData, Generator, RealGenerator