Split-step algorithm

Split-step algorithm#

propagate(psi: Array, *, dt: float | complex, mass: float) Array[source]#

Freely propagates the given wavefunction in time without an external potential:

\[\Psi (x,t+dt) = e^{-\frac{i}{\hbar} T dt} \Psi (x,t).\]
Parameters:
  • psi (fa.Array) – The initial wavefunction \(\Psi(x,t)\).

  • dt (Union[float, complex]) – The timestep \(dt\).

  • mass (float) – The wavefunction’s mass.

Returns:

The freely propagated wavefunction \(\Psi(x,t+dt)\).

Return type:

fa.Array

split_step(psi: Array, *, dt: float, mass: float, V: Array) Array[source]#

Split-step is a pseudo-spectral Fourier method to solve the time-dependent Schrödinger equation. The time evolution of a wavefunction under a constant Hamiltonian $H$ is given by:

\[\Psi (x,t+dt) = e^{-\frac{i}{\hbar}H dt} \Psi (x,t).\]

The time evolution operator can be approximated by [1] [2]

\[e^{-\frac{i}{\hbar}H dt} = e^{-\frac{i}{2\hbar} V dt} e^{-\frac{i}{\hbar} T dt} e^{-\frac{i}{2\hbar} V dt} + \mathcal O (dt^3).\]

Note that the kinetic energy operator \(T\) is diagonal in frequency space and that \(V\) is diagonal in position space. The split-step method utilizes this as follows.

  1. Apply \(e^{-\frac{i}{2\hbar}V dt}\) in position space.

  2. Perform an FFT to get the wavefunction in frequency space.

  3. Apply \(e^{-\frac{i}{\hbar}T dt}\) in frequency space.

  4. Perform an inverse FFT to get the wavefunction in position space.

  5. Apply \(e^{-\frac{i}{2\hbar}V dt}\) in position space.

By this, the computation of the wavefunction’s time evolution is significantly faster. Note that the timestep \(dt\) should be chosen small to reduce the overall error.

Parameters:
  • psi (fa.Array) – The initial wavefunction \(\Psi(x,t)\).

  • dt (float) – The timestep \(dt\).

  • mass (float) – The wavefunction’s mass.

  • V (fa.Array) – The potential in position space.

Returns:

The wavefunction evolved in time: \(\Psi(x,t+dt)\).

Return type:

fa.Array

See also

matterwave.propagate

Used to freely propagate the wavefunction.

jax.lax.scan

Useful function to speed up an iteration.

References

Example

This example shows how to perform a split-step application.

>>> from matterwave import split_step, get_ground_state_ho
>>> import fftarray as fa
>>> from matterwave.constants import Rubidium87 as Rb87
>>> from scipy.constants import pi
>>> # Initialize constants
>>> mass = Rb87.mass # kg
>>> omega_x_init = 2.*pi*0.1 # Hz
>>> omega_x = 2.*pi # Hz
>>> dt = 1e-4 # s
>>> # Define the Dimension
>>> dim = fa.dim_from_constraints("x", pos_min = -200e-6, pos_max = 200e-6, freq_middle = 0., n = 2048)
>>> # Get the coordinates as an Array
>>> x = fa.coords_from_dim(dim, "pos")
>>> # Define the potential
>>> V = 0.5 * mass * omega_x**2. * x**2.
>>> # Initialize the wavefunction
>>> psi_init = get_ground_state_ho(dim, omega=omega_x_init, mass=mass)
>>> # Perform the split-step
>>> psi_final = split_step(psi_init, dt=dt, mass=mass, V=V)
split_step_imag_time(psi: Array, *, dt: float, mass: float, V: Array) Array[source]#

Imaginary time evolution: \(dt \rightarrow -i dt\) using split-step. Normalizes the wave function after each step. For more details see split_step.

Parameters:
  • psi (fa.Array) – The initial wavefunction \(\Psi(x,t)\).

  • dt (float) – The timestep \(dt\).

  • mass (float) – The wavefunction’s mass.

  • V (fa.Array) – The potential in position space.

Returns:

The wavefunction after one imaginary time evolution step.

Return type:

fa.Array

See also

matterwave.split_step

Used to propagate the wavefunction.

matterwave.normalize

Used to normalize the wavefunction.

Example

This example shows how to perform a split-step application with imaginary time.

>>> from matterwave import split_step, get_ground_state_ho
>>> import fftarray as fa
>>> from matterwave.constants import Rubidium87 as Rb87
>>> from scipy.constants import pi
>>> # Initialize constants
>>> mass = Rb87.mass # kg
>>> omega_x_init = 2.*pi*0.1 # Hz
>>> omega_x = 2.*pi # Hz
>>> dt = 1e-4 # s
>>> # Define the Dimension
>>> dim = fa.dim_from_constraints("x", pos_min = -200e-6, pos_max = 200e-6, freq_middle = 0., n = 2048)
>>> # Get the coordinates as an Array
>>> x = fa.coords_from_dim(dim, "pos")
>>> # Define the potential
>>> V = 0.5 * mass * omega_x**2. * x**2.
>>> # Initialize the wavefunction
>>> psi_init = get_ground_state_ho(dim, omega=omega_x_init, mass=mass)
>>> # Perform the split-step
>>> psi_final = split_step_imag_time(psi_init, dt=dt, mass=mass, V=V)