1. The time-evolution Schrödinger equation: $$i\hbar\frac{\partial\psi}{\partial t} = \hat{H}\psi$$ where ($\psi$) represents the wave function, and ($\hat{H}$) is the Hamiltonian operator.

2. The wave function's evolution is performed using a split-operator method, dividing it into the evolution of kinetic energy $K(p, t)$ and potential energy $V(x, t)$. The Hamiltonian operator can be expressed as: $$\hat{H} = \hat{K}(p, t) + \hat{V}(x, t)$$

3. The action of the kinetic energy operator $\hat{K}(p, t)$ and potential energy operator $\hat{V}(x, t)$ on the wave function is carried out using the following equations:

   • Action of the potential energy: $\psi(x, t) \rightarrow \psi(x, t) \cdot (-1)^k \cdot \exp(-0.5j\cdot\Delta t\cdot V(x, t))$
   • Action of the kinetic energy: $\psi(p, t) \rightarrow \psi(p, t) \cdot \exp(-1j\cdot\Delta t\cdot K(p, t))$
   • Here, $\Delta t$ represents the time step, and $k$ is the index of discrete grid points.

4. The expectation value of the Hamiltonian operator (the expectation value of the Hamiltonian) can be expressed as: $$\langle H \rangle = \langle K \rangle + \langle V \rangle$$

5. Formulas for calculating the expectation values:

   • Coordinate expectation value: $\langle x \rangle = \sum_{i} x_i \cdot |\psi(x_i, t)|^2$
   • Momentum expectation value: $\langle p \rangle = \sum_{i} p_i \cdot |\psi(p_i, t)|^2$
   • Hamiltonian expectation value: $\langle H \rangle = \sum_{i} H(x_i, t) \cdot |\psi(x_i, t)|^2$
   • Time derivative of momentum expectation value: $\frac{d\langle p \rangle}{dt} = -\sum_{i} \frac{dV}{dx}(x_i, t) \cdot |\psi(x_i, t)|^2$
   • Time derivative of coordinate expectation value: $\frac{d\langle x \rangle}{dt} = \sum_{i} \frac{dp}{dt}(p_i, t) \cdot |\psi(p_i, t)|^2$

These equations describe the core mathematical operations and calculations used in the code for simulating the one-dimensional Schrödinger equation, employing a split-operator method and verifying Ehrenfest's theorems.

import numpy as np
from scipy import fftpack # Tools for fourier transform
from scipy import linalg # Linear algebra for dense matrix
from numba import njit


class SplitOpSchrodinger1D(object):
    """
    The second-order split-operator propagator of the 1D Schrodinger equation
    in the coordinate representation
    with the time-dependent Hamiltonian H = K(p, t) + V(x, t).
    """
    def __init__(self, *, x_grid_dim, x_amplitude, v, k, dt, diff_k=None, diff_v=None, t=0, abs_boundary=1., **kwargs):
        """
        :param x_grid_dim: the grid size
        :param x_amplitude: the maximum value of the coordinates
        :param v: the potential energy (as a function)
        :param k: the kinetic energy (as a function)
        :param diff_k: the derivative of the potential energy for the Ehrenfest theorem calculations
        :param diff_v: the derivative of the kinetic energy for the Ehrenfest theorem calculations
        :param t: initial value of time
        :param dt: time increment
        :param abs_boundary: absorbing boundary
        :param kwargs: ignored
        """

        # saving the properties
        self.x_grid_dim = x_grid_dim
        self.x_amplitude = x_amplitude
        self.v = v
        self.k = k
        self.diff_v = diff_v
        self.diff_k = diff_k
        self.t = t
        self.dt = dt
        self.abs_boundary = abs_boundary

        # Check that all attributes were specified
        # make sure self.x_amplitude has a value of power of 2
        assert 2 ** int(np.log2(self.x_grid_dim)) == self.x_grid_dim, \
            "A value of the grid size (x_grid_dim) must be a power of 2"

        # get coordinate step size
        self.dx = 2. * self.x_amplitude / self.x_grid_dim

        # generate coordinate range
        x = self.x = (np.arange(self.x_grid_dim) - self.x_grid_dim / 2) * self.dx
        # The same as
        # self.x = np.linspace(-self.x_amplitude, self.x_amplitude - self.dx , self.x_grid_dim)

        # generate momentum range as it corresponds to FFT frequencies
        p = self.p = (np.arange(self.x_grid_dim) - self.x_grid_dim / 2) * (np.pi / self.x_amplitude)

        # allocate the array for wavefunction
        self.wavefunction = np.zeros(self.x.size, dtype=np.complex)

        ####################################################################################################
        #
        # Codes for efficient evaluation
        #
        ####################################################################################################

        if isinstance(abs_boundary, (float, int, np.ndarray)):
            @njit
            def expV(wavefunction, t):
                """
                function to efficiently evaluate
                    wavefunction *= (-1) ** k * exp(-0.5j * dt * v)
                """
                wavefunction *= (-1) ** np.arange(wavefunction.size) * abs_boundary * np.exp(-0.5j * dt * v(x, t))
        else:
            try:
                abs_boundary(x)
            except TypeError:
                raise ValueError("abs_boundary must be a numba function or a numerical constant or a numpy array")

            @njit
            def expV(wavefunction, t):
                """
                function to efficiently evaluate
                    wavefunction *= (-1) ** k * exp(-0.5j * dt * v)
                """
                wavefunction *= (-1) ** np.arange(wavefunction.size) * abs_boundary(x) * np.exp(-0.5j * dt * v(x, t))

        self.expV = expV

        @njit
        def expK(wavefunction, t):
            """
            function to efficiently evaluate
                wavefunction *= exp(-1j * dt * k)
            """
            wavefunction *= np.exp(-1j * dt * k(p, t))

        self.expK = expK

        # Check whether the necessary terms are specified to calculate the first-order Ehrenfest theorems
        if diff_k and diff_v:

            # Get codes for efficiently calculating the Ehrenfest relations

            @njit
            def get_p_average_rhs(density, t):
                return np.sum(density * diff_v(x, t))

            self.get_p_average_rhs = get_p_average_rhs

            # The code above is equivalent to
            # self.get_p_average_rhs = njit(lambda density, t: np.sum(density * diff_v(x, t)))

            @njit
            def get_v_average(density, t):
                return np.sum(v(x, t) * density)

            self.get_v_average = get_v_average

            @njit
            def get_x_average(density):
                return np.sum(x * density)

            self.get_x_average = get_x_average

            @njit
            def get_x_average_rhs(density, t):
                # diff_k = momentum
                return np.sum(diff_k(p, t) * density)

            self.get_x_average_rhs = get_x_average_rhs

            @njit
            def get_k_average(density, t):
                return np.sum(k(p, t) * density)

            self.get_k_average = get_k_average

            @njit
            def get_p_average(density):
                return np.sum(p * density)

            self.get_p_average = get_p_average

            # Lists where the expectation values of x and p
            self.x_average = []
            self.p_average = []

            # Lists where the right hand sides of the Ehrenfest theorems for x and p
            self.x_average_rhs = []
            self.p_average_rhs = []

            # List where the expectation value of the Hamiltonian will be calculated
            self.hamiltonian_average = []

            # List where the expectation value of the kinetic energy will be stored
            self.k_average = []

            # List where the expectation value of the potential energy will be stored
            self.v_average = []

            # Allocate array for storing coordinate or momentum density of the wavefunction
            self.density = np.zeros(self.wavefunction.shape, dtype=np.float)

            # sequence of alternating signs for getting the wavefunction in the momentum representation
            self.minus = (-1) ** np.arange(self.x_grid_dim)

            # Flag requesting tha the Ehrenfest theorem calculations
            self.is_ehrenfest = True
        else:
            # Since diff_v and diff_k are not specified, we are not going to evaluate the Ehrenfest relations
            self.is_ehrenfest = False

    def propagate(self, time_steps=1):
        """
        Time propagate the wave function saved in self.wavefunction
        :param time_steps: number of self.dt time increments to make
        :return: self.wavefunction
        """
        for _ in range(time_steps):

            # advance the wavefunction by dt
            self.single_step_propagation()

            # calculate the Ehrenfest theorems
            self.get_ehrenfest()

        return self.wavefunction

    def single_step_propagation(self):
        """
        Perform a single step propagation of the wavefunction. The wavefunction is normalized.
        :return: self.wavefunction
        """
        # make a half step in time
        self.t += 0.5 * self.dt

        # efficiently evaluate
        #   wavefunction *= (-1) ** k * exp(-0.5j * dt * v)
        self.expV(self.wavefunction, self.t)

        # going to the momentum representation
        self.wavefunction = fftpack.fft(self.wavefunction, overwrite_x=True)

        # efficiently evaluate
        #   wavefunction *= exp(-1j * dt * k)
        self.expK(self.wavefunction, self.t)

        # going back to the coordinate representation
        self.wavefunction = fftpack.ifft(self.wavefunction, overwrite_x=True)

        # efficiently evaluate
        #   wavefunction *= (-1) ** k * exp(-0.5j * dt * v)
        self.expV(self.wavefunction, self.t)

        # make a half step in time
        self.t += 0.5 * self.dt

        # normalize
        # this line is equivalent to
        # self.wavefunction /= np.sqrt(np.sum(np.abs(self.wavefunction) ** 2 ) * self.dx)
        self.wavefunction /= linalg.norm(self.wavefunction) * np.sqrt(self.dx)

        return self.wavefunction

    def get_ehrenfest(self):
        """
        Calculate observables entering the Ehrenfest theorems
        """
        if self.is_ehrenfest:
            # evaluate the coordinate density
            np.abs(self.wavefunction, out=self.density)
            self.density *= self.density
            # normalize
            self.density /= self.density.sum()

            # save the current value of <x>
            self.x_average.append(self.get_x_average(self.density))

            self.p_average_rhs.append(-self.get_p_average_rhs(self.density, self.t))

            # save the potential energy
            self.hamiltonian_average.append(self.get_v_average(self.density, self.t))

            # calculate density in the momentum representation
            wavefunction_p = fftpack.fft(self.minus * self.wavefunction, overwrite_x=True)

            # get the density in the momentum space
            np.abs(wavefunction_p, out=self.density)
            self.density *= self.density
            # normalize
            self.density /= self.density.sum()

            # save the current value of <p>
            self.p_average.append(self.get_p_average(self.density))

            self.x_average_rhs.append(self.get_x_average_rhs(self.density, self.t))

            # save the kinetic energy
            self.k_average.append(self.get_k_average(self.density, self.t))

            # save the potential energy
            self.v_average.append(self.get_v_average(self.density, 0))

            # add the kinetic energy to get the hamiltonian
            self.hamiltonian_average[-1] += self.k_average[-1]

    def set_wavefunction(self, wavefunc):
        """
        Set the initial wave function
        :param wavefunc: 1D numpy array or function specifying the wave function
        :return: self
        """
        if isinstance(wavefunc, np.ndarray):
            # wavefunction is supplied as an array

            # perform the consistency checks
            assert wavefunc.shape == self.wavefunction.shape,\
                "The grid size does not match with the wave function"

            # make sure the wavefunction is stored as a complex array
            np.copyto(self.wavefunction, wavefunc.astype(np.complex))

        else:
            try:
                self.wavefunction[:] = wavefunc(self.x)
            except TypeError:
                raise ValueError("wavefunc must be either function or numpy.array")

        # normalize
        self.wavefunction /= linalg.norm(self.wavefunction) * np.sqrt(self.dx)

        return self

This code snippet appears to be setting up some parameters and initializing an array for a wavefunction, likely for some quantum physics or signal processing simulation. Let's break down what the code is doing step by step and represent it in LaTeX equations:

1. The code first checks that the x_grid_dim attribute is a power of 2. This is done by ensuring that 2^int(np.log2(self.x_grid_dim)) equals self.x_grid_dim. In LaTeX, this check can be represented as:

   $$2^{\lfloor\log_2(\text{self.x\_grid\_dim})\rfloor} = \text{self.x\_grid\_dim}$$

         # Check that all attributes were specified
         # make sure self.x_amplitude has a value of power of 2
         assert 2 ** int(np.log2(self.x_grid_dim)) == self.x_grid_dim, \
            "A value of the grid size (x_grid_dim) must be a power of 2"
   

2. The code calculates the coordinate step size dx as half of the range x_amplitude divided by the grid dimension x_grid_dim. This is represented as:

   $$\text{dx} = \frac{2 \cdot \text{x\_amplitude}}{\text{x\_grid\_dim}}$$

         # get coordinate step size
         self.dx = 2. * self.x_amplitude / self.x_grid_dim
   

3. It generates a coordinate range x using NumPy's arange function, centered around 0 and spaced by dx. This is equivalent to creating an array of evenly spaced values from -x_amplitude to x_amplitude - dx with a total of x_grid_dim points. In LaTeX:

   $$\text{x} = \left(\text{np.arange}(\text{x\_grid\_dim}) - \frac{\text{x\_grid\_dim}}{2}\right) \cdot \text{dx}$$

         # generate coordinate range
         x = self.x = (np.arange(self.x_grid_dim) - self.x_grid_dim / 2) * self.dx
   

4. Similarly, it generates a momentum range p that corresponds to FFT (Fast Fourier Transform) frequencies. p is generated using a similar formula to x, but it uses the mathematical constant π. In LaTeX:

   $$\text{p} = \left(\text{np.arange}(\text{x\_grid\_dim}) - \frac{\text{x\_grid\_dim}}{2}\right) \cdot \frac{\pi}{\text{x\_amplitude}}$$

         # generate momentum range as it corresponds to FFT frequencies
         p = self.p = (np.arange(self.x_grid_dim) - self.x_grid_dim / 2) * (np.pi / self.x_amplitude)
   

5. Finally, it initializes an array called wavefunction of complex numbers with the same size as the x array. The array is filled with zeros. In LaTeX:

   $$\text{self.wavefunction} = \text{np.zeros}(\text{x.size}, \text{dtype}=\text{np.complex})$$

         # allocate the array for wavefunction
         self.wavefunction = np.zeros(self.x.size, dtype=np.complex)
   

In summary, this code prepares the necessary parameters and initializes a complex-valued array for further calculations, likely in the context of a quantum mechanics or signal processing simulation.

The provided code snippet is part of the implementation of the Ehrenfest theorems within the context of simulating the time evolution of a quantum system described by the Schrödinger equation. Let's break down the code and explain how it relates to the Ehrenfest theorems:

1. Ehrenfest Theorems Overview: The Ehrenfest theorems are fundamental principles in quantum mechanics that provide a connection between the classical and quantum descriptions of a physical system. There are two main Ehrenfest theorems:

   • First Ehrenfest Theorem: It relates the time derivatives of the quantum expectation values of position $x$ and momentum $p$ to the corresponding classical quantities in a classical Hamiltonian system. Mathematically, it states:

     • $\frac{d\langle x \rangle}{dt} = \frac{1}{m}\langle p \rangle$
     • $\frac{d\langle p \rangle}{dt} = -\frac{d\langle V \rangle}{dx}$
     • Here, $\langle \cdot \rangle$ represents the quantum expectation value, $m$ is the particle mass, and $V$ is the potential energy.

   • Second Ehrenfest Theorem: It relates the time derivative of the quantum expectation value of the Hamiltonian $H$ to the classical Hamiltonian function. Mathematically, it states: $\frac{d\langle H \rangle}{dt} = \langle \frac{\partial H}{\partial t} \rangle$

2. Code Explanation: The provided code implements these Ehrenfest theorems as follows:

   • The method get_ehrenfest calculates various observables to verify the Ehrenfest theorems. It checks if the simulation has been configured for Ehrenfest theorem calculations (self.is_ehrenfest).

   • It begins by evaluating the coordinate density of the wavefunction $\psi$. The absolute value of the wavefunction is computed and squared to obtain the probability density $|\psi|^2$, which represents the probability of finding the particle at a given position.

              # evaluate the coordinate density
              np.abs(self.wavefunction, out=self.density)
              self.density *= self.density
   

   • The code then normalizes the density to ensure that the probabilities sum up to 1.

              # normalize
              self.density /= self.density.sum()
   

   • The current value of the quantum expectation value of position $\langle x \rangle$ is computed using the probability density. This corresponds to the first part of the first Ehrenfest theorem.

              # save the current value of <x>
              self.x_average.append(self.get_x_average(self.density))
   

   • The code calculates $\frac{d\langle p \rangle}{dt}$, which corresponds to the second part of the first Ehrenfest theorem. This quantity is stored in self.p_average_rhs and is related to the time derivative of the momentum.

              self.p_average_rhs.append(-self.get_p_average_rhs(self.density, self.t))
   

   • The potential energy $V$ at the current time is calculated and stored in self.hamiltonian_average. This is part of the second Ehrenfest theorem.

               # save the potential energy
               self.hamiltonian_average.append(self.get_v_average(self.density, self.t))
   

   • The code then transforms the wavefunction into momentum space using a Fast Fourier Transform (FFT). This allows the calculation of momentum-related quantities.

               # calculate density in the momentum representation
               wavefunction_p = fftpack.fft(self.minus * self.wavefunction, overwrite_x=True)
   

   • Similar to the position space, it computes the momentum density, normalizes it, and calculates the quantum expectation value of momentum $\langle p \rangle$, which is stored in self.p_average.

               # get the density in the momentum space
               np.abs(wavefunction_p, out=self.density)
               self.density *= self.density
   
               # normalize
               self.density /= self.density.sum()
   
               # save the current value of <p>
               self.p_average.append(self.get_p_average(self.density))
   

   • It computes $\frac{d\langle x \rangle}{dt}$ using the momentum space quantities and stores it in self.x_average_rhs. This corresponds to the time derivative of the position and is part of the first Ehrenfest theorem.

                self.x_average_rhs.append(self.get_x_average_rhs(self.density, self.t))
   

   • The kinetic energy $K$ at the current time is calculated and stored in self.k_average. This is also part of the second Ehrenfest theorem.

                # save the kinetic energy
                self.k_average.append(self.get_k_average(self.density, self.t))
   

   • Finally, the code calculates the potential energy at time $t=0$ and stores it in self.v_average.

                # save the potential energy
                self.v_average.append(self.get_v_average(self.density, 0))
   

   • The kinetic energy is added to the potential energy to calculate the total Hamiltonian $H$ at the current time, and this value is stored in self.hamiltonian_average.

                # add the kinetic energy to get the hamiltonian
                self.hamiltonian_average[-1] += self.k_average[-1]