KDK

The KDK module implements the split-step Fourier method using a "kick–drift–kick" (KDK) algorithm to evolve the Schrödinger–Poisson equation. The scheme is a second-order Strang splitting that decouples the kinetic and potential parts of the Hamiltonian.

The full mathematical derivation is given in Algorithms §2; this page documents the two low-level entry points that the main loop calls.

Core Functions

apply_kick_step!

apply_kick_step!(
    device_ψ,
    ψ,
    V::Function,
    rho,
    rho_max,
    rho_max_id::CartesianIndex{3},
    grid::SimulationGrid,
    config_gravity::GravityConfig,
    config_tidal::TidalFieldConfig,
    config_device::DeviceConfig,
    dt::Real,
    i::Int64,
    t::Vector{<:Real}
) -> Tuple{Any, Any, Any}

Each call performs step (c), then step (a) of Algorithm 1 in the paper: it computes the gravitational potential from the current wavefunction (including the tidal field, if enabled), applies the half-timestep potential operator, and returns the Fourier transform of the kicked wavefunction ready for the drift step. The closing half kick, step (d), is applied by the next iteration's kick call; with KDK_flag = true the kick length is dt/2, so consecutive kick–drift–kick steps form the symmetric Strang sequence.

The implementation handles:

• device_ψ is consumed by this function on the GPU path (the underlying buffer is freed before returning). The caller must not reuse the binding.
• On the GPU, intermediate buffers (device_Φ_all, device_Φ_WaveDM, potential_grav_static, device_V, potential_all, nonlinear_term, device_ψ) are released through CUDA.unsafe_free! so peak GPU memory stays close to a single wavefunction's worth.
• The tidal field (if MW_tidal_field = true) is added to $\Phi_{\rm total}$ at this point.

apply_drift_step!

apply_drift_step!(
    spec,
    linear_phase,
    border,
    gpu::Bool,
    DeviceArray
) -> Any

The drift step fuses step (b) of Algorithm 1 with the absorbing boundary: it multiplies the Fourier-space wavefunction by the linear phase factor $e^{i\,\nabla^2\,\Delta t/2}$, transforms back to real space, and applies the absorbing boundary mask $\mathcal B(\mathbf r)$.

The spec and linear_phase arguments are consumed by this function; their bindings should not be reused after the call.

Algorithm Description

Write the SPE evolution operator as the sum of two pieces:

\[\partial_t \psi \;=\; \bigl(\hat D + \hat N\bigr)\psi,\]

with the kinetic dispersion operator

\[\hat D = \tfrac{i}{2}\nabla^2,\]

and the potential operator

\[\hat N = -i\,\Phi_{\rm total}.\]

The second-order Strang split applied over a single timestep $\Delta t$ is

\[\psi(\mathbf r, t+\Delta t) \;\approx\; e^{\frac{\Delta t}{2}\hat N}\,e^{\Delta t \hat D}\,e^{\frac{\Delta t}{2}\hat N}\,\psi(\mathbf r, t).\]

In Fourier space the drift is diagonal: $e^{\Delta t \hat D} \to e^{-\frac{i}{2}\Delta t\,k^2}$.

Steps (c) + (a) — apply_kick_step!

1. Compute the gravitational potential from the current density:

   \[\rho = |\psi|^2, \qquad \Phi_{\rm DM}(\mathbf k) = -\frac{4\pi}{k^2}\bigl[\rho(\mathbf k) - \langle\rho\rangle\,\delta(\mathbf k)\bigr].\]

2. If MW_tidal_field = true, add the time-dependent MW (+LMC) tidal field to $\Phi_{\rm total}$.

3. Evaluate the user-supplied potential $V(\mathbf r, t, \psi)$ on the grid and assemble $\Phi_{\rm total}$.

4. Apply the half-timestep potential kick to the wavefunction:

   \[\psi \;\leftarrow\; e^{\frac{\Delta t}{2}\hat N}\,\psi \;=\; e^{-i\,\Phi_{\rm total}\,\Delta t/2}\,\psi.\]

5. Transform to Fourier space and return the spectrum for the drift step.

Step (b) — apply_drift_step!

1. Multiply the Fourier-space wavefunction by the linear phase:

   \[\tilde\psi \;\leftarrow\; e^{i\,\nabla^2\,\Delta t/2}\,\tilde\psi \;=\; e^{-i\,k^2\,\Delta t/2}\,\tilde\psi.\]

2. Inverse Fourier transform back to real space: $\psi \leftarrow \mathcal F^{-1}\{\tilde\psi\}$.

3. Apply the absorbing boundary mask $\mathcal B(\mathbf r)$:

   \[\psi \;\leftarrow\; \psi \odot \mathcal B(\mathbf r).\]

The boundary mask is set up by setup_absorption_boundary(Xmax, Ymax, Zmax, x, y, z, absorb_coeff, dt). It multiplies $\psi$ by

\[\mathcal B(\mathbf r) = \exp\!\left[-\alpha\left(6-\sum_{\text{faces}}\tanh\frac{x_i\pm L/2}{L/50}\right)\Delta t\right],\]

so the damping is strongest at the box faces and decays smoothly into the interior over a transition layer of width $\sim L/50$. The factor is not strictly zero at the boundary; with the default $\alpha=10$ it is $\exp(-\alpha\,\Delta t)$ at a face for a field that is uniform in the other two directions.

Sign Convention

The drift implementation in setup_fft_operators uses the convention

\[\nabla^2 \;\longrightarrow\; -(2\pi k_x)^2 - (2\pi k_y)^2 - (2\pi k_z)^2,\]

so linear_phase = fftshift(exp.(im * Laplacian * dt/2)) evaluates to the same diagonal factor as the mathematical form above.

Memory Management

The internal helper _release!(buf, gpu) wraps the conditional in a single function so that call sites remain readable.