Stationary states

Generated from the docstrings in the source code.

eigensolver

Direct eigensolver for H = −ℏ²/2m ∇² + V(r) on 1D/2D/3D grids.

The Hamiltonian is never stored: it is applied matrix-free with FFTs and handed to LOBPCG, preconditioned by the inverse kinetic operator in k-space. For a real potential the spectral Hamiltonian is real-symmetric, so the eigenvectors are real.

eigenstates(grid, potential, n_states: int = 6, hbar: float = 1.0, m: float = 1.0, tol: float = None, max_iter: int = 1000, seed: int = 0, method: str = 'auto', workers: int = -1) -> Tuple[np.ndarray, List]

Lowest n_states eigenpairs of H.

Parameters:

Name	Type	Description	Default
method	(auto, lobpcg, lanczos)	'lanczos' uses ARPACK (robust, best for small grids); 'lobpcg' uses a k-space preconditioned block solver (fast for large 2D/3D grids). 'auto' picks Lanczos below 20 000 grid points.	'auto'
tol	float	Lanczos: relative eigenvalue accuracy (default 1e-12). LOBPCG: residual-norm tolerance (default 1e-6); the eigenvalue error scales as the residual squared.	None
seed	int	Seed of the deterministic random starting block (LOBPCG).	0

Returns:

Name	Type	Description
energies	(ndarray, shape(n_states))	Ascending eigenvalues.
states	list of WaveFunction	Normalised eigenfunctions (∫|ψ|² dV = 1). Degenerate levels are returned as an arbitrary orthonormal basis of their subspace.

Source code in quantumlab/stationary/eigensolver.py

def eigenstates(grid, potential, n_states: int = 6, hbar: float = 1.0, m: float = 1.0,
                tol: float = None, max_iter: int = 1000, seed: int = 0,
                method: str = 'auto', workers: int = -1) -> Tuple[np.ndarray, List]:
    """
    Lowest ``n_states`` eigenpairs of H.

    Parameters
    ----------
    method : {'auto', 'lobpcg', 'lanczos'}
        'lanczos' uses ARPACK (robust, best for small grids); 'lobpcg' uses a
        k-space preconditioned block solver (fast for large 2D/3D grids).
        'auto' picks Lanczos below 20 000 grid points.
    tol : float, optional
        Lanczos: relative eigenvalue accuracy (default 1e-12).
        LOBPCG: residual-norm tolerance (default 1e-6); the eigenvalue error
        scales as the residual squared.
    seed : int
        Seed of the deterministic random starting block (LOBPCG).

    Returns
    -------
    energies : ndarray, shape (n_states,)
        Ascending eigenvalues.
    states : list of WaveFunction
        Normalised eigenfunctions (∫|ψ|² dV = 1). Degenerate levels are
        returned as an arbitrary orthonormal basis of their subspace.
    """
    shape = grid.shape
    axes = tuple(range(len(shape)))
    size = int(np.prod(shape))
    V = np.asarray(potential.evaluate(grid), dtype=float)
    T_k = hbar ** 2 * k_squared(grid) / (2.0 * m)

    def _as_block(v):
        v = np.asarray(v)
        return v.reshape(shape + (-1,)), v.ndim == 1

    def matvec(v):
        block, squeeze = _as_block(v)
        k_block = fftn(block, axes=axes, workers=workers) * T_k[..., None]
        out = (ifftn(k_block, axes=axes, workers=workers).real + V[..., None] * block).reshape(size, -1)
        return out[:, 0] if squeeze else out

    H = LinearOperator((size, size), matvec=matvec, matmat=matvec, dtype=float)

    if method == 'auto':
        method = 'lanczos' if size < 20000 else 'lobpcg'

    if method == 'lanczos':
        energies, vecs = eigsh(H, k=n_states, which='SA', tol=1e-12 if tol is None else tol)
    elif method == 'lobpcg':
        shift = max(1.0, float(np.abs(V).mean()))
        inv_T = 1.0 / (T_k + shift)

        def precondition(v):
            block, squeeze = _as_block(v)
            k_block = fftn(block, axes=axes, workers=workers) * inv_T[..., None]
            out = ifftn(k_block, axes=axes, workers=workers).real.reshape(size, -1)
            return out[:, 0] if squeeze else out

        M = LinearOperator((size, size), matvec=precondition, matmat=precondition, dtype=float)
        rng = np.random.default_rng(seed)
        block_size = n_states + max(2, n_states // 2)
        X0 = rng.standard_normal((size, block_size))
        res_tol = 1e-6 if tol is None else tol
        with warnings.catch_warnings():
            # LOBPCG warns about the extra padding vectors too; the requested
            # states are checked explicitly below.
            warnings.simplefilter('ignore', UserWarning)
            energies, vecs = lobpcg(H, X0, M=M, tol=res_tol, maxiter=max_iter, largest=False)
        energies, vecs = energies[:n_states], vecs[:, :n_states]
        residuals = np.linalg.norm(H.matmat(vecs) - vecs * energies, axis=0) / np.linalg.norm(vecs, axis=0)
        if np.any(residuals > 10 * res_tol):
            warnings.warn(f'LOBPCG did not converge: max residual {residuals.max():.2e}', RuntimeWarning)
    else:
        raise ValueError("method must be 'auto', 'lobpcg' or 'lanczos'")

    order = np.argsort(energies)
    energies = energies[order]
    vecs = vecs[:, order]

    dV = volume_element(grid)
    wf_class = wavefunction_class(grid)
    states = []
    for j in range(n_states):
        psi = vecs[:, j].reshape(shape) / np.sqrt(dV)
        psi /= np.sqrt(np.sum(psi ** 2) * dV)
        # Fix the global sign so the largest lobe is positive (reproducible output)
        idx = np.unravel_index(np.argmax(np.abs(psi)), shape)
        if psi[idx] < 0:
            psi = -psi
        states.append(wf_class(grid, psi.astype(complex)))
    return energies, states

imaginary_time

Imaginary-time propagation (ψ ← e^{−Hτ}ψ / ‖·‖) for ground and excited states, including the nonlinear Gross–Pitaevskii ground state.

The converged state of 2nd-order Strang imaginary-time stepping differs from the exact eigenstate by O(dτ²); the reported energy is evaluated with the exact spectral Hamiltonian, so its error is O(dτ⁴) for the ground state.

ground_state(grid, potential, dtau: float = 0.01, nonlinearity: float = 0.0, initial=None, tol: float = 1e-11, max_steps: int = 50000, check_interval: int = 20, hbar: float = 1.0, m: float = 1.0, orthogonal_to: Optional[List] = None, workers: int = -1, rotation: float = 0.0, backend='numpy')

Relax to the lowest-energy state by imaginary-time propagation.

Parameters:

Name	Type	Description	Default
dtau	float or sequence of float	Imaginary time step. A decreasing sequence runs a continuation: the large steps relax slow modes (e.g. of vortex lattices) in few steps, the small ones lower the O(dτ²) error floor. max_steps and tol apply to each stage, and info describes the last one (steps is the total).	0.01
nonlinearity	float	Gross–Pitaevskii coupling g (ψ normalised to 1). g = 0 gives the linear Schrödinger ground state.	0.0
initial	WaveFunction or ndarray	Starting guess. A good guess (e.g. Thomas–Fermi) speeds convergence.	None
tol	float	Stop when the energy changes by less than tol between checks.	1e-11
orthogonal_to	list of WaveFunction	States projected out after each step (Gram–Schmidt), used to obtain excited states of the linear problem.	None
rotation	float	Angular velocity Ω of a frame rotating about z (2D/3D). Energies and μ are then those of H − ΩL_z. Above the critical Ω the GP ground state contains vortices; imaginary time only reaches it if the guess breaks rotational symmetry (e.g. an imprinted vortex or a lattice seed).	0.0
backend	(numpy, cupy, auto)	Where the propagation runs (see SplitOperatorSolver). Energies are evaluated on the host every check_interval steps.	'numpy'

Returns:

Name	Type	Description
wf	WaveFunction
info	ConvergenceInfo	energy is the GP energy functional E[ψ] (equal to the eigenvalue for g = 0), chemical_potential is μ = ⟨ψ|H_GP|ψ⟩ and residual is ‖H_GP ψ − μψ‖.

Source code in quantumlab/stationary/imaginary_time.py

def ground_state(grid, potential, dtau: float = 0.01, nonlinearity: float = 0.0,
                 initial=None, tol: float = 1e-11, max_steps: int = 50000,
                 check_interval: int = 20, hbar: float = 1.0, m: float = 1.0,
                 orthogonal_to: Optional[List] = None, workers: int = -1,
                 rotation: float = 0.0, backend='numpy'):
    """
    Relax to the lowest-energy state by imaginary-time propagation.

    Parameters
    ----------
    dtau : float or sequence of float
        Imaginary time step. A decreasing sequence runs a continuation: the
        large steps relax slow modes (e.g. of vortex lattices) in few steps,
        the small ones lower the O(dτ²) error floor. ``max_steps`` and
        ``tol`` apply to each stage, and ``info`` describes the last one
        (``steps`` is the total).
    nonlinearity : float
        Gross–Pitaevskii coupling g (ψ normalised to 1). g = 0 gives the
        linear Schrödinger ground state.
    initial : WaveFunction or ndarray, optional
        Starting guess. A good guess (e.g. Thomas–Fermi) speeds convergence.
    tol : float
        Stop when the energy changes by less than ``tol`` between checks.
    orthogonal_to : list of WaveFunction, optional
        States projected out after each step (Gram–Schmidt), used to obtain
        excited states of the linear problem.
    rotation : float
        Angular velocity Ω of a frame rotating about z (2D/3D). Energies and
        μ are then those of H − ΩL_z. Above the critical Ω the GP ground state
        contains vortices; imaginary time only reaches it if the guess breaks
        rotational symmetry (e.g. an imprinted vortex or a lattice seed).
    backend : {'numpy', 'cupy', 'auto'}
        Where the propagation runs (see ``SplitOperatorSolver``). Energies are
        evaluated on the host every ``check_interval`` steps.

    Returns
    -------
    wf : WaveFunction
    info : ConvergenceInfo
        ``energy`` is the GP energy functional E[ψ] (equal to the eigenvalue
        for g = 0), ``chemical_potential`` is μ = ⟨ψ|H_GP|ψ⟩ and ``residual``
        is ‖H_GP ψ − μψ‖.
    """
    if np.ndim(dtau) > 0:
        wf, total = initial, 0
        for step_size in dtau:
            wf, info = ground_state(grid, potential, float(step_size), nonlinearity, wf, tol, max_steps,
                                    check_interval, hbar, m, orthogonal_to, workers, rotation, backend)
            total += info.steps
        info.steps = total
        return wf, info
    if orthogonal_to and nonlinearity != 0.0:
        raise ValueError('orthogonal_to is only meaningful for the linear problem')
    solver = SplitOperatorSolver(grid, potential, dtau, hbar, m, order=2,
                                 nonlinearity=nonlinearity, imaginary_time=True, workers=workers,
                                 rotation=rotation, backend=backend)
    dV = volume_element(grid)
    wf_class = wavefunction_class(grid)
    V = solver.potential_at(0.0)
    lower = [np.asarray(s.psi) for s in (orthogonal_to or [])]

    if initial is None:
        psi = _default_guess(grid)
    else:
        psi = np.array(getattr(initial, 'psi', initial), dtype=complex)

    def project(p):
        for q in lower:
            p = p - _inner(q, p, dV) * q
        return p / np.sqrt(np.real(_inner(p, p, dV)))

    psi = project(psi)
    wf = wf_class(grid, psi)
    energy_prev = np.inf
    converged = False
    steps = 0
    while steps < max_steps:
        if lower:
            for _ in range(check_interval):
                wf = solver.step(wf)
                wf.psi = project(wf.psi)
        else:
            # Keeps ψ on the device (GPU backend) for the whole interval
            wf = solver.evolve(wf, check_interval)
        steps += check_interval
        energy, _, _ = _energy_and_mu(wf.psi, grid, V, hbar, m, nonlinearity, dV, rotation)
        if abs(energy - energy_prev) < tol:
            converged = True
            break
        energy_prev = energy

    energy, mu, residual = _energy_and_mu(wf.psi, grid, V, hbar, m, nonlinearity, dV, rotation)
    return wf, ConvergenceInfo(converged, steps, energy, mu, residual)

excited_states(grid, potential, n_states: int, dtau: float = 0.01, tol: float = 1e-11, max_steps: int = 50000, seed: int = 0, hbar: float = 1.0, m: float = 1.0, workers: int = -1)

Lowest n_states eigenstates of the linear Hamiltonian, found one by one with imaginary time + Gram–Schmidt against the states already converged. Starting guesses are seeded random fields (deterministic for a fixed seed).

Returns:

Name	Type	Description
energies	ndarray
states	list of WaveFunction
infos	list of ConvergenceInfo

Source code in quantumlab/stationary/imaginary_time.py

def excited_states(grid, potential, n_states: int, dtau: float = 0.01, tol: float = 1e-11,
                   max_steps: int = 50000, seed: int = 0, hbar: float = 1.0, m: float = 1.0,
                   workers: int = -1):
    """
    Lowest ``n_states`` eigenstates of the linear Hamiltonian, found one by one
    with imaginary time + Gram–Schmidt against the states already converged.
    Starting guesses are seeded random fields (deterministic for a fixed seed).

    Returns
    -------
    energies : ndarray
    states : list of WaveFunction
    infos : list of ConvergenceInfo
    """
    rng = np.random.default_rng(seed)
    envelope = np.abs(_default_guess(grid))
    states, energies, infos = [], [], []
    for _ in range(n_states):
        guess = envelope * rng.standard_normal(grid.shape)
        wf, info = ground_state(grid, potential, dtau=dtau, initial=guess, tol=tol,
                                max_steps=max_steps, hbar=hbar, m=m,
                                orthogonal_to=states, workers=workers)
        states.append(wf)
        energies.append(info.energy)
        infos.append(info)
    return np.array(energies), states, infos

coupled_ground_state(grid, potentials, couplings, initial, dtau: float = 0.005, tol: float = 1e-11, max_steps: int = 50000, check_interval: int = 20, hbar: float = 1.0, m: float = 1.0, rotation: float = 0.0, backend='numpy', workers: int = -1)

Lowest-energy state of coupled Gross–Pitaevskii equations by imaginary time, at fixed populations N_i given by the norms of initial.

Parameters:

Name	Type	Description	Default
initial	sequence of WaveFunction or ndarray	One starting state per component. For immiscible mixtures the guess should break the symmetry between the components (e.g. displaced clouds); a symmetric guess can stay on a symmetric saddle point.	required

Returns:

Name	Type	Description
states	tuple of WaveFunction
info	CoupledConvergenceInfo	energy is the total energy functional (of H − ΩL_z when rotating), chemical_potentials and residuals (‖H_iψ_i − μ_iψ_i‖/√N_i) are per component.

Source code in quantumlab/stationary/imaginary_time.py

def coupled_ground_state(grid, potentials, couplings, initial, dtau: float = 0.005, tol: float = 1e-11,
                         max_steps: int = 50000, check_interval: int = 20, hbar: float = 1.0,
                         m: float = 1.0, rotation: float = 0.0, backend='numpy', workers: int = -1):
    """
    Lowest-energy state of coupled Gross–Pitaevskii equations by imaginary
    time, at fixed populations N_i given by the norms of ``initial``.

    Parameters
    ----------
    initial : sequence of WaveFunction or ndarray
        One starting state per component. For immiscible mixtures the guess
        should break the symmetry between the components (e.g. displaced
        clouds); a symmetric guess can stay on a symmetric saddle point.

    Returns
    -------
    states : tuple of WaveFunction
    info : CoupledConvergenceInfo
        ``energy`` is the total energy functional (of H − ΩL_z when rotating),
        ``chemical_potentials`` and ``residuals`` (‖H_iψ_i − μ_iψ_i‖/√N_i) are
        per component.
    """
    from quantumlab.solvers.multicomponent import CoupledGPESolver
    g = np.asarray(couplings, dtype=float)
    solver = CoupledGPESolver(grid, potentials, g, dtau, hbar, m, order=2, imaginary_time=True,
                              rotation=rotation, backend=backend, workers=workers)
    dV = volume_element(grid)
    wf_class = wavefunction_class(grid)
    Vs = [np.zeros(grid.shape) if p is None else np.asarray(p.evaluate(grid), dtype=float) for p in potentials]
    states = tuple(wf_class(grid, np.array(getattr(s, 'psi', s), dtype=complex)) for s in initial)
    energy_prev = np.inf
    converged = False
    steps = 0
    while steps < max_steps:
        states = solver.evolve(states, check_interval)
        steps += check_interval
        energy, _, _ = _coupled_energy_mu([s.psi for s in states], grid, Vs, g, hbar, m, rotation, dV)
        if abs(energy - energy_prev) < tol:
            converged = True
            break
        energy_prev = energy
    energy, mus, residuals = _coupled_energy_mu([s.psi for s in states], grid, Vs, g, hbar, m, rotation, dV)
    return states, CoupledConvergenceInfo(converged, steps, energy, mus, residuals)

conjugate_gradient

Ground states by preconditioned nonlinear conjugate gradient on the sphere ‖ψ‖² = N (Antoine, Levitt & Tang, J. Comput. Phys. 343, 92 (2017)).

Imaginary-time propagation is a first-order method whose fixed point carries an O(dτ²) error, and soft modes (vortex lattices, elongated condensates) relax in it as exp(−Δτ) with a tiny gap Δ. Minimising the energy directly removes both problems: the fixed point is an exact stationary state, so the residual ‖Hψ − μψ‖ can be driven to small values, and conjugate directions converge on soft modes much faster than steepest descent.

The preconditioner is P = P_T^½ P_V P_T^½ with P_T = (α + ℏ²k²/2m)⁻¹ in Fourier space and P_V = (α + V + g|ψ|²)⁻¹ in real space, which tames both the kinetic and the potential stiffness.

ground_state_pcg(grid, potential, initial=None, nonlinearity: float = 0.0, rotation: float = 0.0, tol: float = 1e-08, max_iter: int = 10000, hbar: float = 1.0, m: float = 1.0, callback=None, callback_interval: int = 100, workers: int = -1)

Minimise E[ψ] = ⟨ψ|−ℏ²∇²/2m + V − ΩL_z|ψ⟩ + (g/2)∫|ψ|⁴ at fixed norm.

Parameters:

Name	Type	Description	Default
initial	WaveFunction or ndarray	Starting state; its norm N is kept. For vortex states start from a guess with vortices (e.g. from damped dynamics or imaginary time).	None
tol	float	Stop when the residual ‖Hψ − μψ‖/√N falls below tol.	1e-08
callback	callable	callback(iteration, residual, energy) every callback_interval iterations.	None

Returns:

Name	Type	Description
wf	WaveFunction
info	ConvergenceInfo	steps is the number of iterations, converged whether tol was reached, residual = ‖Hψ − μψ‖/√N.

Source code in quantumlab/stationary/conjugate_gradient.py

def ground_state_pcg(grid, potential, initial=None, nonlinearity: float = 0.0, rotation: float = 0.0,
                     tol: float = 1e-8, max_iter: int = 10000, hbar: float = 1.0, m: float = 1.0,
                     callback=None, callback_interval: int = 100, workers: int = -1):
    """
    Minimise E[ψ] = ⟨ψ|−ℏ²∇²/2m + V − ΩL_z|ψ⟩ + (g/2)∫|ψ|⁴ at fixed norm.

    Parameters
    ----------
    initial : WaveFunction or ndarray, optional
        Starting state; its norm N is kept. For vortex states start from a
        guess with vortices (e.g. from damped dynamics or imaginary time).
    tol : float
        Stop when the residual ‖Hψ − μψ‖/√N falls below ``tol``.
    callback : callable, optional
        ``callback(iteration, residual, energy)`` every ``callback_interval``
        iterations.

    Returns
    -------
    wf : WaveFunction
    info : ConvergenceInfo
        ``steps`` is the number of iterations, ``converged`` whether ``tol``
        was reached, ``residual`` = ‖Hψ − μψ‖/√N.
    """
    if getattr(potential, 'time_dependent', False):
        raise ValueError('ground_state_pcg requires a time-independent potential')
    dV = volume_element(grid)
    V = np.zeros(grid.shape) if potential is None else np.asarray(potential.evaluate(grid), dtype=float)
    T_k = hbar ** 2 * k_squared(grid) / (2.0 * m)
    g = float(nonlinearity)
    psi = _default_guess(grid) if initial is None else np.array(getattr(initial, 'psi', initial), dtype=complex)

    def inner(a, b):
        return float(np.real(np.sum(np.conj(a) * b)) * dV)

    if initial is None:
        psi /= np.sqrt(inner(psi, psi))                        # default guess at unit norm

    def h_psi(p):
        return apply_hamiltonian(p, grid, V, hbar, m, nonlinearity=g, workers=workers, rotation=rotation)

    def quartic(p):
        return float(np.sum((p.real ** 2 + p.imag ** 2) ** 2) * dV)

    def energy_from(p, H_p):
        # ⟨ψ|Hψ⟩ with H including g|ψ|² counts the interaction twice: E = ⟨ψ|Hψ⟩ − (g/2)∫|ψ|⁴
        return inner(p, H_p) - 0.5 * g * quartic(p)

    # Each iteration applies H twice (trial point and accepted point); the accepted
    # point's Hψ is reused for the next residual and energy.
    N = inner(psi, psi)
    Hpsi = h_psi(psi)
    d_prev = r_prev = Pr_prev = None
    converged = False
    for it in range(max_iter + 1):
        mu = inner(psi, Hpsi) / N
        r = Hpsi - mu * psi
        residual = np.sqrt(inner(r, r) / N)
        E0 = energy_from(psi, Hpsi)
        if callback is not None and it % callback_interval == 0:
            callback(it, residual, E0)
        if residual < tol:
            converged = True
            break
        if it == max_iter:
            break
        # The shift keeps both preconditioner factors positive, also for V < 0
        alpha = max(1.0, mu - float(V.min()))
        PT_half = 1.0 / np.sqrt(alpha + T_k)
        PV = 1.0 / (alpha + V + g * np.abs(psi) ** 2)
        Pr = ifftn(PT_half * fftn(PV * ifftn(PT_half * fftn(r, workers=workers), workers=workers),
                                  workers=workers), workers=workers)
        Pr -= inner(psi, Pr) / N * psi                        # tangent space of the sphere
        beta = 0.0 if r_prev is None else max(0.0, inner(r - r_prev, Pr) / inner(r_prev, Pr_prev))
        d = -Pr if d_prev is None else -Pr + beta * d_prev
        d -= inner(psi, d) / N * psi
        if inner(r, d) >= 0.0:                                 # not a descent direction: restart
            d = -Pr
        u = d * np.sqrt(N / inner(d, d))                       # tangent direction with ‖u‖² = N
        # Line search along the great circle ψ(θ) = ψ cos θ + u sin θ (norm preserved).
        # While energy differences are resolvable, θ comes from a quadratic fit of E(θ);
        # near convergence they scale as residual² and drown in round-off, so a secant
        # step on the directional derivative E'(θ) = 2 Re⟨Hψ(θ), ψ'(θ)⟩ is used instead.
        slope0 = 2.0 * inner(Hpsi, u)
        trial = min(0.5, 0.5 * np.sqrt(inner(d, d) / N))
        psi_t = psi * np.cos(trial) + u * np.sin(trial)
        H_t = h_psi(psi_t)
        E_t = energy_from(psi_t, H_t)
        if abs(E_t - E0) > 1e-10 * max(1.0, abs(E0)):
            curvature = (E_t - E0 - slope0 * trial) / trial ** 2
            theta = -slope0 / (2.0 * curvature) if curvature > 0 else trial
        else:
            slope_t = 2.0 * inner(H_t, -psi * np.sin(trial) + u * np.cos(trial))
            theta = trial * slope0 / (slope0 - slope_t) if slope_t > slope0 else trial
        for _ in range(30):
            if theta == trial:
                candidate, H_c, E_c = psi_t, H_t, E_t
            else:
                candidate = psi * np.cos(theta) + u * np.sin(theta)
                H_c = h_psi(candidate)
                E_c = energy_from(candidate, H_c)
            # accept unless the energy rises by more than round-off
            if E_c <= E0 + 1e-12 * max(1.0, abs(E0)):
                break
            theta *= 0.5
        psi, Hpsi = candidate, H_c
        d_prev, r_prev, Pr_prev = d, r, Pr

    return wavefunction_class(grid)(grid, psi), ConvergenceInfo(converged, it, E0, mu, residual)

hamiltonian

apply_hamiltonian(psi: np.ndarray, grid, V: np.ndarray, hbar: float = 1.0, m: float = 1.0, nonlinearity: float = 0.0, workers: int = -1, rotation: float = 0.0) -> np.ndarray

Hψ = −ℏ²/2m ∇²ψ + (V + g|ψ|²)ψ − Ω L_z ψ with all derivatives evaluated spectrally (exact for band-limited ψ on the periodic grid). L_z is taken about the z axis, matching SplitOperatorSolver(rotation=Ω).

Source code in quantumlab/stationary/hamiltonian.py

def apply_hamiltonian(psi: np.ndarray, grid, V: np.ndarray, hbar: float = 1.0, m: float = 1.0,
                      nonlinearity: float = 0.0, workers: int = -1, rotation: float = 0.0) -> np.ndarray:
    """
    Hψ = −ℏ²/2m ∇²ψ + (V + g|ψ|²)ψ − Ω L_z ψ with all derivatives evaluated
    spectrally (exact for band-limited ψ on the periodic grid). L_z is taken
    about the z axis, matching ``SplitOperatorSolver(rotation=Ω)``.
    """
    T_k = hbar ** 2 * k_squared(grid) / (2.0 * m)
    psi_k = fftn(psi, workers=workers)
    H_psi = ifftn(T_k * psi_k, workers=workers)
    V_eff = V if nonlinearity == 0.0 else V + nonlinearity * np.abs(psi) ** 2
    H_psi = H_psi + V_eff * psi
    if rotation != 0.0:
        d_dx = ifftn(1j * grid.K_x * psi_k, workers=workers)
        d_dy = ifftn(1j * grid.K_y * psi_k, workers=workers)
        L_z_psi = -1j * hbar * (grid.X * d_dy - grid.Y * d_dx)
        H_psi = H_psi - rotation * L_z_psi
    return H_psi