Skip to content

Solvers

Generated from the docstrings in the source code.

split_operator

Generic split-operator propagator for 1D, 2D and 3D grids.

One solver covers:

  • 2nd-order Strang splitting and 4th-order Yoshida (triple-jump) composition,
  • explicitly time-dependent potentials V(r, t),
  • the cubic Gross–Pitaevskii nonlinearity g|ψ|²ψ (BEC mean field, solitons),
  • imaginary-time propagation (ground states),
  • a rotating frame, H → H − ΩL_z (2D/3D, rotation about the z axis),
  • phenomenological damping, (i − γ)ℏ ∂ψ/∂t = Hψ, with the norm held fixed,
  • CPU (NumPy/SciPy) or CUDA GPU (CuPy) execution.

Time-dependent potentials are handled by treating t as an extra coordinate: the Strang step S(h) applies exp(−i V(t) h/2ℏ), exp(−i T h/ℏ), exp(−i V(t+h) h/2ℏ). This is a symmetric 2nd-order method, so the Yoshida composition S(w₁h) S(w₀h) S(w₁h) is 4th-order accurate.

The nonlinear sub-flow i ℏ ∂ψ/∂t = g|ψ|²ψ leaves |ψ| unchanged and is therefore integrated exactly by the phase factor exp(−i g|ψ|² h/2ℏ). In imaginary time the state is kept at unit norm and the nonlinear term uses the normalised density g|ψ|²/‖ψ‖² evaluated at the sub-step midpoint, which keeps the converged GP ground state 2nd-order accurate in dτ.

With rotation, −ΩL_z = Ω(y p_x − x p_y) is split together with the kinetic energy into H_x = p_x²/2m + Ω y p_x, diagonal in (k_x, y), and H_yz = (p_y² + p_z²)/2m − Ω x p_y, diagonal in (x, k_y, k_z). Each factor is applied exactly with 1D/2D FFTs, and the symmetric product e^{H_x h/2} e^{H_yz h} e^{H_x h/2} keeps the step 2nd order (and symmetric, so the Yoshida composition stays 4th order).

SplitOperatorSolver(grid, potential, dt: float, hbar: float = 1.0, m: float = 1.0, order: int = 2, nonlinearity: float = 0.0, imaginary_time: bool = False, absorber=None, workers: int = -1, t0: float = 0.0, rotation: float = 0.0, damping: float = 0.0, backend='numpy', device_id: int = 0)

Bases: Solver

Parameters:

Name Type Description Default
grid (Grid1D, Grid2D or Grid3D)
required
potential Potential or None

Static or time-dependent (time_dependent = True) potential. None means free space.

 required
dt float

Time step (imaginary-time step dτ when imaginary_time=True).

 required
hbar float

Reduced Planck constant and particle mass.

 1.0
m float

Reduced Planck constant and particle mass.

 1.0
order (2, 4)

Splitting order: 2 = Strang, 4 = Yoshida triple jump.

 2
nonlinearity float

Gross–Pitaevskii coupling g in V_eff = V + g|ψ|². With ψ normalised to 1, g includes the particle number (g = 4πℏ²aN/m in 3D).

 0.0
imaginary_time bool

Propagate in imaginary time t → −iτ and renormalise after every step. Converges to the lowest state not orthogonal to the initial guess.

 False
rotation float

Angular velocity Ω of the frame about the z axis (2D/3D grids). The frame rotates counter-clockwise for Ω > 0, so states with positive L_z are favoured. The potential is given in the rotating frame.

 0.0
damping float

Dissipation γ ≥ 0 in (i − γ)ℏ ∂ψ/∂t = Hψ (Choi et al. 1998, Tsubota et al. 2002). The norm is restored after every step, which plays the role of the chemical potential. Order 2 only.

 0.0
absorber AbsorbingBoundaryLayer

Mask applied after every step.

 None
workers int

Threads for scipy.fft (-1 = all cores); ignored on the GPU.

 -1
t0 float

Initial time (relevant for time-dependent potentials).

 0.0
backend (numpy, cupy, auto)

'cupy' runs on a CUDA GPU in double precision; 'auto' uses the GPU when CuPy and a CUDA device are available. evolve keeps ψ on the device between steps and only copies it back for callbacks and at the end; step copies to and from the device on every call. Time-dependent potentials are evaluated on the host and copied to the device at every sub-step.

 'numpy'
device_id int

CUDA device for the 'cupy' backend.

 0
Notes

Every factor is unitary, but for nonlinear problems on a uniform (or broad) background the split-step method has a parametric numerical instability when the kinetic phase of the highest grid wavenumber per (sub)step exceeds π (Weideman & Herbst, SIAM J. Numer. Anal. 23, 485 (1986)). Keep ℏ k_max² / (2m) · |w| dt < π with k_max = π/dx, where |w| ≤ 1.71 for order 4, otherwise round-off noise in the shortest wavelengths can grow exponentially. Localised states (solitons, trapped condensates) are much less sensitive.

Source code in quantumlab/solvers/split_operator.py 
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
def __init__(self, grid, potential, dt: float, hbar: float = 1.0, m: float = 1.0,
             order: int = 2, nonlinearity: float = 0.0, imaginary_time: bool = False,
             absorber=None, workers: int = -1, t0: float = 0.0,
             rotation: float = 0.0, damping: float = 0.0,
             backend='numpy', device_id: int = 0):
    super().__init__(grid, potential, dt, hbar, m)
    if order not in (2, 4):
        raise ValueError('order must be 2 or 4')
    self.time_dependent = bool(getattr(potential, 'time_dependent', False))
    if imaginary_time and order == 4:
        raise ValueError('order=4 uses a negative substep, which is unstable in imaginary time')
    if imaginary_time and self.time_dependent:
        raise ValueError('imaginary-time propagation requires a time-independent potential')
    if damping < 0:
        raise ValueError('damping must be non-negative')
    if damping > 0 and (order == 4 or imaginary_time):
        raise ValueError('damping requires order=2 real-time propagation')
    if rotation != 0.0 and ndim(grid) < 2:
        raise ValueError('rotation requires a 2D or 3D grid')

    self.order = order
    self.nonlinearity = float(nonlinearity)
    self.imaginary_time = imaginary_time
    self.rotation = float(rotation)
    self.damping = float(damping)
    # Steps that are not unitary are followed by renormalisation
    self._renormalise = imaginary_time or self.damping > 0
    self.absorber = absorber
    self.workers = workers
    self.t0 = t0
    self.steps_taken = 0
    self._be = get_backend(backend, workers=workers, device_id=device_id)
    self._xp = self._be.xp

    self._wf_class = wavefunction_class(grid)
    self._dV = volume_element(grid)
    # exp(c · E · h / ℏ): c = −i in real time, −1 in imaginary time and
    # −(i + γ)/(1 + γ²) with damping, from ∂ψ/∂t = Hψ / ((i − γ)ℏ)
    self._c = -1.0 if imaginary_time else -(1j + self.damping) / (1.0 + self.damping ** 2)
    self._T_k = self.hbar ** 2 * k_squared(grid) / (2.0 * self.m)
    self._substeps = (1.0,) if order == 2 else (YOSHIDA_W1, YOSHIDA_W0, YOSHIDA_W1)

    self._V_static = None
    if not self.time_dependent:
        self._V_static = (np.zeros(grid.shape) if potential is None
                          else np.asarray(potential.evaluate(grid), dtype=float))
    self._V_static_dev = None if self._V_static is None else self._be.to_device(self._V_static)
    self._mask_dev = None if absorber is None else self._be.to_device(np.asarray(absorber.mask))
    self.update_operators()

t: float property

Current simulation time.

update_operators()

(Re)build the cached kinetic and static-potential propagators.

Source code in quantumlab/solvers/split_operator.py 
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
def update_operators(self):
    """(Re)build the cached kinetic and static-potential propagators."""
    dev = self._be.to_device
    self._U_T = {}
    self._U_x_half = {}
    self._U_yz = {}
    self._U_V_half = {}
    for w in set(self._substeps):
        h = w * self.dt
        if self.rotation == 0.0:
            self._U_T[w] = dev(np.exp(self._c * self._T_k * h / self.hbar))
        else:
            H_x, H_yz = self._rotating_kinetic_parts()
            self._U_x_half[w] = dev(np.exp(self._c * H_x * h / (2.0 * self.hbar)))
            self._U_yz[w] = dev(np.exp(self._c * H_yz * h / self.hbar))
        if self._V_static is not None:
            self._U_V_half[w] = dev(np.exp(self._c * self._V_static * h / (2.0 * self.hbar)))

potential_at(t: float) -> np.ndarray

V(r, t) as a NumPy array on the host.

Source code in quantumlab/solvers/split_operator.py 
187
188
189
190
191
def potential_at(self, t: float) -> np.ndarray:
    """V(r, t) as a NumPy array on the host."""
    if self._V_static is not None:
        return self._V_static
    return np.asarray(self.potential.evaluate_at(self.grid, t), dtype=float)

evolve(wavefunction, n_steps: int, callback=None, callback_interval: int = 1)

Advance n_steps steps. If given, callback(step, t, wf) is called every callback_interval steps (step counted from 1 within this call). On the GPU, ψ stays on the device except for those callbacks.

Source code in quantumlab/solvers/split_operator.py 
273
274
275
276
277
278
279
280
281
282
283
284
def evolve(self, wavefunction, n_steps: int, callback=None, callback_interval: int = 1):
    """
    Advance ``n_steps`` steps. If given, ``callback(step, t, wf)`` is called
    every ``callback_interval`` steps (step counted from 1 within this call).
    On the GPU, ψ stays on the device except for those callbacks.
    """
    psi = self._upload(wavefunction)
    for n in range(1, n_steps + 1):
        psi = self._advance(psi)
        if callback is not None and n % callback_interval == 0:
            callback(n, self.t, self._download(psi))
    return self._download(psi)

multicomponent

Coupled Gross–Pitaevskii equations for multi-component condensates (e.g. two hyperfine states of ⁸⁷Rb, or two atomic species of equal mass):

iℏ ∂ψ_i/∂t = [ −ℏ²∇²/2m + V_i(r) + Σ_j g_ij |ψ_j|² − ΩL_z ] ψ_i

The kinetic (and rotation) factor is the one of SplitOperatorSolver and is applied to every component. The potential sub-flow is local and, in real time, leaves every |ψ_j| unchanged, so it is integrated exactly with the densities frozen. In imaginary time or with damping each component keeps its own norm N_i (the populations are fixed), and the coupled potential uses the densities rescaled to those norms at the sub-step midpoint, as for a single component, which keeps the converged states 2nd-order accurate in dτ.

CoupledGPESolver(grid, potentials: Sequence, couplings, dt: float, hbar: float = 1.0, m: float = 1.0, order: int = 2, imaginary_time: bool = False, rotation: float = 0.0, damping: float = 0.0, workers: int = -1, backend='numpy', device_id: int = 0, t0: float = 0.0)

Parameters:

Name Type Description Default
grid (Grid1D, Grid2D or Grid3D)
required
potentials sequence of Potential or None

One static potential per component (None means free space).

 required
couplings (n, n) array_like

Symmetric interaction matrix g_ij. The nonlinear term uses the actual densities |ψ_j|², so g_ij is per particle when ψ_j is normalised to its particle number.

 required
dt float

Time step (dτ in imaginary time).

 required
order int

As in SplitOperatorSolver; they apply to all components.

 2
imaginary_time int

As in SplitOperatorSolver; they apply to all components.

 2
rotation int

As in SplitOperatorSolver; they apply to all components.

 2
damping int

As in SplitOperatorSolver; they apply to all components.

 2
workers int

As in SplitOperatorSolver; they apply to all components.

 2
backend int

As in SplitOperatorSolver; they apply to all components.

 2
device_id int

As in SplitOperatorSolver; they apply to all components.

 2
t0 int

As in SplitOperatorSolver; they apply to all components.

 2
Source code in quantumlab/solvers/multicomponent.py 
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
def __init__(self, grid, potentials: Sequence, couplings, dt: float, hbar: float = 1.0,
             m: float = 1.0, order: int = 2, imaginary_time: bool = False, rotation: float = 0.0,
             damping: float = 0.0, workers: int = -1, backend='numpy', device_id: int = 0,
             t0: float = 0.0):
    self.n_components = len(potentials)
    g = np.asarray(couplings, dtype=float)
    if g.shape != (self.n_components, self.n_components):
        raise ValueError(f'couplings must be a {self.n_components}×{self.n_components} matrix')
    if not np.allclose(g, g.T):
        raise ValueError('couplings must be symmetric (g_ij = g_ji)')
    if any(getattr(p, 'time_dependent', False) for p in potentials):
        raise ValueError('CoupledGPESolver supports static potentials only')
    self.grid = grid
    self.potentials = list(potentials)
    self.couplings = g
    self.dt = dt
    self.hbar = hbar
    self.m = m
    self.t0 = t0
    self.steps_taken = 0
    # The single-component engine supplies the kinetic/rotation propagators,
    # the time-stepping constant c, the substep weights and the array backend.
    self._engine = SplitOperatorSolver(grid, None, dt, hbar, m, order=order, imaginary_time=imaginary_time,
                                       rotation=rotation, damping=damping, workers=workers,
                                       backend=backend, device_id=device_id, t0=t0)
    self._be = self._engine._be
    self._xp = self._be.xp
    self._c = self._engine._c
    self._renormalise = self._engine._renormalise
    self._substeps = self._engine._substeps
    self._dV = volume_element(grid)
    self._wf_class = wavefunction_class(grid)
    self._V = [self._be.to_device(np.zeros(grid.shape) if p is None else
                                  np.asarray(p.evaluate(grid), dtype=float)) for p in self.potentials]

evolve(states, n_steps: int, callback=None, callback_interval: int = 1)

Advance all components n_steps steps; callback(step, t, states) is called every callback_interval steps with a tuple of wavefunctions.

Source code in quantumlab/solvers/multicomponent.py 
134
135
136
137
138
139
140
141
142
143
144
def evolve(self, states, n_steps: int, callback=None, callback_interval: int = 1):
    """
    Advance all components ``n_steps`` steps; ``callback(step, t, states)``
    is called every ``callback_interval`` steps with a tuple of wavefunctions.
    """
    psis = self._upload(states)
    for n in range(1, n_steps + 1):
        psis = self._advance(psis)
        if callback is not None and n % callback_interval == 0:
            callback(n, self.t, self._download(psis))
    return self._download(psis)

backends

Array backends for the split-operator solver: NumPy/SciPy on the CPU or CuPy on a CUDA GPU. A backend bundles the array module, the FFTs and the host/device transfer functions, so the solver code is identical for both.

cupy_available() -> bool

True if CuPy is importable and at least one CUDA device is visible.

Source code in quantumlab/solvers/backends.py 
22
23
24
25
26
27
28
def cupy_available() -> bool:
    """True if CuPy is importable and at least one CUDA device is visible."""
    try:
        import cupy
        return cupy.cuda.runtime.getDeviceCount() > 0
    except Exception:
        return False

get_backend(backend='numpy', workers: int = -1, device_id: int = 0) -> ArrayBackend

Resolve a backend specification.

'numpy' (default) — SciPy FFTs on the CPU; 'cupy' — CUDA GPU, raises ImportError/RuntimeError if unavailable; 'auto' — CuPy when a CUDA device is present, otherwise NumPy; an ArrayBackend instance is returned unchanged.

Source code in quantumlab/solvers/backends.py 
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
def get_backend(backend='numpy', workers: int = -1, device_id: int = 0) -> ArrayBackend:
    """
    Resolve a backend specification.

    ``'numpy'`` (default) — SciPy FFTs on the CPU;
    ``'cupy'`` — CUDA GPU, raises ImportError/RuntimeError if unavailable;
    ``'auto'`` — CuPy when a CUDA device is present, otherwise NumPy;
    an ``ArrayBackend`` instance is returned unchanged.
    """
    if isinstance(backend, ArrayBackend):
        return backend
    if backend == 'numpy':
        return numpy_backend(workers)
    if backend == 'cupy':
        if not cupy_available():
            raise RuntimeError("backend='cupy' requested but CuPy or a CUDA device is not available")
        return cupy_backend(device_id)
    if backend == 'auto':
        return cupy_backend(device_id) if cupy_available() else numpy_backend(workers)
    raise ValueError("backend must be 'numpy', 'cupy', 'auto' or an ArrayBackend")

split_step_1d

split_step_2d

SplitStep2DSolver(grid, potential, dt: float, hbar: float = 1.0, m: float = 1.0)

Bases: Solver

2D Split-Step Fourier Method (SSFM) solver for the time-dependent Schrödinger equation on a 2D spatial grid.

The symmetric splitting operator is applied as: ψ(t+dt) = U_V · FFT2⁻¹[ U_T · FFT2[ U_V · ψ(t) ] ]

where U_V = exp(-i·V·dt/(2ℏ)) and U_T = exp(-i·T_k·dt/ℏ).

Source code in quantumlab/solvers/split_step_2d.py 
17
18
19
20
21
22
23
def __init__(self, grid, potential, dt: float, hbar: float = 1.0, m: float = 1.0):
    super().__init__(grid, potential, dt, hbar, m)
    self.V = self.potential.evaluate(self.grid)
    self.T_k = self.hbar**2 * (self.grid.K_x**2 + self.grid.K_y**2) / (2.0 * self.m)
    self.U_V = None
    self.U_T = None
    self.update_operators()

update_operators()

Pre-compute the position and momentum space propagators.

Source code in quantumlab/solvers/split_step_2d.py 
25
26
27
28
def update_operators(self):
    """Pre-compute the position and momentum space propagators."""
    self.U_V = np.exp(-1j * self.V * self.dt / (2.0 * self.hbar))
    self.U_T = np.exp(-1j * self.T_k * self.dt / self.hbar)

step(wavefunction: WaveFunction2D) -> WaveFunction2D

Advance the wavefunction by one time step using SSFM.

Source code in quantumlab/solvers/split_step_2d.py 
30
31
32
33
34
35
36
37
38
39
40
41
def step(self, wavefunction: WaveFunction2D) -> WaveFunction2D:
    """Advance the wavefunction by one time step using SSFM."""
    psi = wavefunction.psi.copy()
    # Half-step in position space
    psi *= self.U_V
    # Full step in momentum space
    psi_k = fft2(psi)
    psi_k *= self.U_T
    psi = ifft2(psi_k)
    # Half-step in position space
    psi *= self.U_V
    return WaveFunction2D(self.grid, psi)

split_step_3d

SplitStep3DSolver(grid, potential, dt: float, hbar: float = 1.0, m: float = 1.0, workers: int = -1)

Bases: Solver

3D Split-Step Fourier Method (SSFM) solver for the time-dependent Schrödinger equation on a 3D spatial grid.

The symmetric Lie-Trotter-Suzuki splitting is applied as: ψ(t+dt) = U_V/2 · IFFT3[ U_T · FFT3[ U_V/2 · ψ(t) ] ]

where U_V = exp(-i·V·dt/(2ℏ)) and U_T = exp(-i·T_k·dt/ℏ), with T_k = ℏ²(K_x² + K_y² + K_z²) / (2m).

This scheme is 2nd-order accurate in time (global error O(dt²)) and exactly unitary when no absorbing boundary is applied.

Parameters:

Name Type Description Default
grid Grid3D

The 3D spatial grid.

 required
potential Potential

The potential energy operator V(x, y, z).

 required
dt float

Time step.

 required
hbar float

Reduced Planck constant (natural units).

 1.0
m float

Particle mass (natural units).

 1.0
workers int

Number of CPU threads for scipy FFT. -1 uses all available cores.

 -1
Source code in quantumlab/solvers/split_step_3d.py 
37
38
39
40
41
42
43
44
45
46
47
def __init__(self, grid, potential, dt: float, hbar: float = 1.0,
             m: float = 1.0, workers: int = -1):
    super().__init__(grid, potential, dt, hbar, m)
    self.workers = workers
    self.V = self.potential.evaluate(self.grid)
    self.T_k = (self.hbar ** 2 *
                (self.grid.K_x ** 2 + self.grid.K_y ** 2 + self.grid.K_z ** 2) /
                (2.0 * self.m))
    self.U_V = None
    self.U_T = None
    self.update_operators()

update_operators()

Pre-compute the position-space and momentum-space propagators.

Source code in quantumlab/solvers/split_step_3d.py 
49
50
51
52
def update_operators(self):
    """Pre-compute the position-space and momentum-space propagators."""
    self.U_V = np.exp(-1j * self.V * self.dt / (2.0 * self.hbar))
    self.U_T = np.exp(-1j * self.T_k * self.dt / self.hbar)

step(wavefunction: WaveFunction3D) -> WaveFunction3D

Advance the wavefunction by one time step using symmetric SSFM.

Source code in quantumlab/solvers/split_step_3d.py 
54
55
56
57
58
59
60
61
62
63
64
65
def step(self, wavefunction: WaveFunction3D) -> WaveFunction3D:
    """Advance the wavefunction by one time step using symmetric SSFM."""
    psi = wavefunction.psi.copy()
    # Half-step in position space
    psi *= self.U_V
    # Full step in momentum space
    psi_k = fftn(psi, workers=self.workers)
    psi_k *= self.U_T
    psi = ifftn(psi_k, workers=self.workers)
    # Half-step in position space
    psi *= self.U_V
    return WaveFunction3D(self.grid, psi)

split_step_3d_gpu

3D Split-Step Fourier solver with CuPy GPU acceleration & zero PCI-E memory bottleneck. Automatic fallback to CPU (SciPy) if CuPy is unavailable. Deterministic: same grid, same potential, same dt -> bit-identical evolution.

SplitStep3DGPU(grid, potential, dt: float, hbar: float = 1.0, m: float = 1.0, device_id: int = 0, renormalize_interval: int = 0)

Bases: Solver

GPU-accelerated 3D SSFM with double-precision complex FP64 support. Grid operators are transferred to GPU once at init. Supports in-VRAM batch stepping via step_sequence() or step_gpu() to eliminate PCI-E transfer bottlenecks.

Source code in quantumlab/solvers/split_step_3d_gpu.py 
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
def __init__(
    self,
    grid,
    potential,
    dt: float,
    hbar: float = 1.0,
    m: float = 1.0,
    device_id: int = 0,
    renormalize_interval: int = 0,
):
    super().__init__(grid, potential, dt, hbar, m)
    self.device_id = device_id
    self.renormalize_interval = renormalize_interval
    self._step_count = 0
    self._gpu = CUPY_AVAILABLE

    if self._gpu:
        try:
            num_devices = cp.cuda.runtime.getDeviceCount()
            if device_id < 0 or device_id >= num_devices:
                raise ValueError(
                    f"Invalid GPU device_id {device_id}. Available CUDA devices: {num_devices}"
                )
            cp.cuda.Device(device_id).use()
        except Exception as err:
            print(f"[WARN] Failed to select CUDA device {device_id}: {err}. Falling back to CPU.")
            self._gpu = False

    if self._gpu:
        self.V = cp.asarray(self.potential.evaluate(self.grid), dtype=cp.float64)
        self.T_k = cp.asarray(
            (self.hbar ** 2 *
             (self.grid.K_x ** 2 + self.grid.K_y ** 2 + self.grid.K_z ** 2) /
             (2.0 * self.m)),
            dtype=cp.float64
        )
    else:
        self.V = np.asarray(self.potential.evaluate(self.grid), dtype=np.float64)
        self.T_k = np.asarray(
            (self.hbar ** 2 *
             (self.grid.K_x ** 2 + self.grid.K_y ** 2 + self.grid.K_z ** 2) /
             (2.0 * self.m)),
            dtype=np.float64
        )

    self.update_operators()

update_operators() -> None

Re-compute potential and kinetic evolution operators.

Source code in quantumlab/solvers/split_step_3d_gpu.py 
75
76
77
78
79
80
81
82
def update_operators(self) -> None:
    """Re-compute potential and kinetic evolution operators."""
    if self._gpu:
        self.U_V = cp.exp(-1j * self.V * self.dt / (2.0 * self.hbar))
        self.U_T = cp.exp(-1j * self.T_k * self.dt / self.hbar)
    else:
        self.U_V = np.exp(-1j * self.V * self.dt / (2.0 * self.hbar))
        self.U_T = np.exp(-1j * self.T_k * self.dt / self.hbar)

step(wavefunction: WaveFunction3D) -> WaveFunction3D

Single-step propagation returning CPU WaveFunction3D object.

Source code in quantumlab/solvers/split_step_3d_gpu.py 
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
def step(self, wavefunction: WaveFunction3D) -> WaveFunction3D:
    """Single-step propagation returning CPU WaveFunction3D object."""
    self._step_count += 1
    if self._gpu:
        psi_gpu = cp.asarray(wavefunction.psi, dtype=cp.complex128)
        psi_gpu = self._step_gpu_internal(psi_gpu)
        return WaveFunction3D(self.grid, cp.asnumpy(psi_gpu))
    else:
        psi = wavefunction.psi.copy()
        psi *= self.U_V
        psi_k = fftn(psi)
        psi_k *= self.U_T
        psi = ifftn(psi_k)
        psi *= self.U_V
        if self.renormalize_interval > 0 and self._step_count % self.renormalize_interval == 0:
            norm = np.sqrt(np.sum(np.abs(psi)**2) * self.grid.dx * self.grid.dy * self.grid.dz)
            if norm > 0:
                psi /= norm
        return WaveFunction3D(self.grid, psi)

step_sequence(initial_wf: WaveFunction3D, steps: int) -> WaveFunction3D

In-VRAM multi-step propagation eliminating PCI-E transfer bottlenecks. Transfers initial psi to GPU ONCE, steps N times on GPU, and downloads final state.

Source code in quantumlab/solvers/split_step_3d_gpu.py 
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
def step_sequence(self, initial_wf: WaveFunction3D, steps: int) -> WaveFunction3D:
    """
    In-VRAM multi-step propagation eliminating PCI-E transfer bottlenecks.
    Transfers initial psi to GPU ONCE, steps N times on GPU, and downloads final state.
    """
    if not self._gpu:
        wf = initial_wf
        for _ in range(steps):
            wf = self.step(wf)
        return wf

    psi_gpu = cp.asarray(initial_wf.psi, dtype=cp.complex128)
    for _ in range(steps):
        self._step_count += 1
        psi_gpu = self._step_gpu_internal(psi_gpu)

    return WaveFunction3D(self.grid, cp.asnumpy(psi_gpu))

base