Quantum Dynamics in a Magnetic Field: Ground State, Currents, and Time Evolution

This notebook models a two-component (spin-1/2) Bose-Einstein Condensate or quantum particle in a 2D space subjected to an external magnetic field, an electric field, and a confining potential well.

The Hamiltonian of the system includes the kinetic term with minimal coupling, a non-linear Gross-Pitaevskii (GP) interaction, Zeeman spin splitting, and a confining potential:

$$H = \frac{1}{2m}(\vec{p} - q\vec{A})^2 + V_{trap}(\vec{r}) + g|\psi|^2 - \frac{\hbar q B}{2m}\sigma_z + q A_t$$

We will calculate the ground state using Imaginary Time Evolution, analyze the Probability and Spin Currents, and perform a Real-Time Evolution using a Runge-Kutta 4th order (RK4) solver.

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation

# --- GLOBAL PLOT SETTINGS (Publication Quality) ---
plt.rcParams.update({
    "text.usetex": True,
    "font.family": "serif",
    "font.serif": ["Computer Modern"],
    "font.size": 11,
    "axes.labelsize": 11,
    "xtick.labelsize": 10,
    "ytick.labelsize": 10,
    "legend.fontsize": 10,
    "figure.titlesize": 12,
})

1. System Setup and Grid Discretization

We define the 2D spatial grid and the corresponding momentum grid using Fast Fourier Transforms (FFT). We also initialize a trial wave function (a Gaussian wave packet in the spin-up component).

# Spatial Grid definition
Lx, Ly = 5, 5
N = 128
x = np.linspace(-Lx, Lx, N, endpoint=False)
y = np.linspace(-Ly, Ly, N, endpoint=False)
X, Y = np.meshgrid(x, y)
dx = x[1] - x[0]
dy = y[1] - y[0]

# Momentum Grid definition (for FFT derivatives)
kx = np.fft.fftfreq(N, d=dx) * 2 * np.pi
ky = np.fft.fftfreq(N, d=dy) * 2 * np.pi
KX, KY = np.meshgrid(kx, ky)

def normalize(psi):
    """Normalizes the 2-component spinor wave function."""
    norm = np.sqrt(np.sum(np.abs(psi)**2) * dx * dy)
    return psi / norm

# Initial trial wave function (Gaussian in spin-up)
sigma = 2.5
psi = np.zeros((2, N, N), dtype=complex)
psi[0, :, :] = np.exp(-(X**2) / (2 * sigma**2)) # Spin UP
psi[1, :, :] = np.zeros_like(X)                 # Spin DOWN
psi = normalize(psi)

2. Gauge Fields and Potentials

We apply a uniform magnetic field $B$ using the Landau gauge $\vec{A} = (0, Bx)$ and an electric field via the scalar potential $A_t = Ex$. We also define a confining well potential in the $x$-direction to prevent boundary issues with the periodic conditions assumed by the FFT.

# Physical Parameters
g = 0     # Gross-Pitaevskii interaction strength
B = 1     # Magnetic field
E = 0.5   # Electric field
q = 1     # Charge
m = 1     # Mass
hbar = 1  # Reduced Planck constant

# Gauge formulation: A = (A_x, A_y), Scalar potential = A_t
A_x = np.zeros_like(psi)
A_y = np.zeros_like(psi)
A_y[0, :, :] = B * X 
A_y[1, :, :] = B * X 

A_t = np.zeros_like(psi)
A_t[0, :, :] = E * X 
A_t[1, :, :] = E * X 

# Confining potential well
omega_B = q * B / m
V_0 = 9 * hbar * omega_B
x_max = Lx / 1.5        
potential_well = V_0 * (1 + np.tanh((np.abs(X) - 0.75 * x_max) * 9))

3. Hamiltonian Operators

Using spectral derivatives via FFT, we construct the operators for the Hamiltonian. FFT is highly accurate for smooth periodic functions and avoids the numerical dissipation of finite differences.

def calc_gp_potential(psi, g):
    """Non-linear Gross-Pitaevskii potential term."""
    density = np.sum(np.abs(psi)**2, axis=0)
    return g * density

def kinetic_hamiltonian(psi, A_x, A_y, A_t, KX, KY, hbar=1, m=1, q=1):
    """Applies minimal coupling kinetic Hamiltonian using spectral derivatives."""
    psi_k = np.fft.fft2(psi, axes=(1, 2))

    # Spectral derivatives
    d_psi_x = np.fft.ifft2(1j * KX * psi_k, axes=(1, 2))
    d_psi_y = np.fft.ifft2(1j * KY * psi_k, axes=(1, 2))
    laplacian = np.fft.ifft2(-(KX**2 + KY**2) * psi_k, axes=(1, 2))
    
    A_dot_grad = A_x * d_psi_x + A_y * d_psi_y
    A_sq = A_x**2 + A_y**2

    term1 = -(hbar**2 / (2 * m)) * laplacian
    term2 = 1j * (hbar * q / m) * A_dot_grad
    term3 = (q**2 / (2 * m)) * A_sq * psi
    term4 = q * A_t * psi

    return term1 + term2 + term3 + term4

def spin_hamiltonian(psi, B, hbar=1, m=1, q=1):
    """Zeeman spin splitting."""
    omega_B = q * B / m
    E_spin = np.zeros_like(psi)
    E_spin[0, :, :] = -hbar * omega_B / 2 * psi[0, :, :]
    E_spin[1, :, :] = hbar * omega_B / 2 * psi[1, :, :]
    return E_spin

4. Ground State via Imaginary Time Evolution

We find the ground state by evolving the system in imaginary time $\tau = it$. High energy modes decay exponentially faster, ensuring convergence to the ground state.

$$\psi(\tau + \Delta\tau) \approx (1 - H\Delta\tau)\psi(\tau)$$

dtau = 0.001
max_steps = 20000
energy_tolerance = 1e-10
energies = []

print("Starting imaginary time evolution...")

for step in range(max_steps):
    H_kin = kinetic_hamiltonian(psi, A_x, A_y, A_t, KX, KY)
    H_gp  = calc_gp_potential(psi, g) * psi
    H_sp  = spin_hamiltonian(psi, B)
    H_trap = potential_well * psi
    
    H_total = H_kin + H_gp + H_sp + H_trap
    
    # Forward Euler step in imaginary time
    psi_new = psi - H_total * dtau
    psi = normalize(psi_new)
    
    # Check convergence periodically
    if step % 100 == 0:
        expected_energy = np.sum(np.conj(psi) * H_total).real * dx * dy
        energies.append(expected_energy)
        
        if len(energies) > 1:
            delta_E = np.abs(energies[-1] - energies[-2])
            if delta_E < energy_tolerance:
                print(f"Convergence reached at step {step}! Energy = {expected_energy:.6f}")
                break

# Plot energy convergence and final density
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 4))

ax1.plot(energies, color='purple')
ax1.set_title("Energy Convergence")
ax1.set_xlabel("Steps (x100)")
ax1.set_ylabel("Total Energy")
ax1.grid(True, alpha=0.3)

im = ax2.imshow(np.abs(psi[0])**2, cmap='inferno', origin='lower', extent=[-Lx, Lx, -Ly, Ly])
ax2.set_xlabel(r'$x/l_B$')
ax2.set_ylabel(r'$y/l_B$')
ax2.set_title(r'Density $|\psi_\uparrow|^2$')
cbar = plt.colorbar(im, ax=ax2)
cbar.set_label(r'$|\psi|^2$', rotation=270, labelpad=15)

plt.tight_layout()
plt.show()

# Optional: Save ground state data
# np.savez('ground_state_data.npz', psi=psi, x=x, y=y, Lx=Lx, Ly=Ly, N=N, g=g, B=B)

Starting imaginary time evolution...
Convergence reached at step 8400! Energy = -0.112857

5. Probability and Spin Currents

Once the ground state is reached, we evaluate the orbital and spin currents.

The orbital current is given by: $$\vec{J}_{orb} = \frac{1}{2m} \left( \psi^\dagger(\hat{\vec{p}} - q\vec{A})\psi + h.c. \right)$$

The spin current is defined via the spin texture: $$\vec{J}_{spin} = \frac{\hbar}{2m} \nabla \times (\psi^\dagger \vec{\sigma} \psi)$$

def calculate_currents(psi, A_x, A_y, KX, KY, q=1, hbar=1, m=1):
    psi_k = np.fft.fft2(psi, axes=(1, 2))
    
    kx_psi = np.fft.ifft2(KX * psi_k, axes=(1, 2))
    ky_psi = np.fft.ifft2(KY * psi_k, axes=(1, 2))
    
    psi_up, psi_down = psi[0], psi[1]
    psi_up_conj, psi_down_conj = np.conj(psi_up), np.conj(psi_down)
    
    # Orbital Current Calculation
    J_orb_x = (psi_up_conj*(kx_psi[0] - q*A_x[0]*psi_up) + psi_down_conj*(kx_psi[1] - q*A_x[1]*psi_down)).real / m
    J_orb_y = (psi_up_conj*(ky_psi[0] - q*A_y[0]*psi_up) + psi_down_conj*(ky_psi[1] - q*A_y[1]*psi_down)).real / m
    
    # Spin Current Calculation (Spin Texture components)
    spin_texture_xy = psi_up_conj * psi_down
    C_x = 2 * spin_texture_xy.real
    C_y = 2 * spin_texture_xy.imag
    C_z = np.abs(psi_up)**2 - np.abs(psi_down)**2
    
    C_z_k = np.fft.fft2(C_z)
    
    J_spin_x = (hbar / (2*m)) * np.fft.ifft2(1j * KY * C_z_k).real
    J_spin_y = (hbar / (2*m)) * np.fft.ifft2(-1j * KX * C_z_k).real
    
    return J_orb_x, J_orb_y, J_spin_x, J_spin_y

J_orb_x, J_orb_y, J_spin_x, J_spin_y = calculate_currents(psi, A_x, A_y, KX, KY)
J_total_x = J_orb_x + J_spin_x
J_total_y = J_orb_y + J_spin_y
J_max = np.max(np.sqrt(J_orb_x**2 + J_orb_y**2))

# --- PLOTTING CURRENTS ---
fig, axes = plt.subplots(1, 3, figsize=(12, 3.8), constrained_layout=True)
density = np.abs(psi[0])**2 + np.abs(psi[1])**2

def plot_stream_panel(ax, Jx, Jy, title, colormap_color):
    im = ax.imshow(density, extent=[-Lx, Lx, -Ly, Ly], origin='lower', cmap='inferno')
    
    J_mag = np.sqrt(Jx**2 + Jy**2)
    lw = 2.5 * (J_mag / J_max)
    threshold = 0.03 * J_max
    
    Jx_plot = np.where(J_mag > threshold, Jx, np.nan)
    Jy_plot = np.where(J_mag > threshold, Jy, np.nan)
    
    ax.streamplot(X, Y, Jx_plot, Jy_plot, color=colormap_color, linewidth=lw, density=1.0, arrowsize=1.2)
    ax.set_title(title)
    ax.set_xlabel(r'$x/l_B$')
    ax.set_ylabel(r'$y/l_B$')
    ax.set_xlim(-Lx, Lx)
    ax.set_ylim(-Ly, Ly)
    return im

plot_stream_panel(axes[0], J_orb_x, J_orb_y, r'Orbital Current $\vec{J}_{orb}$', 'cyan')
plot_stream_panel(axes[1], J_spin_x, J_spin_y, r'Spin Current $\vec{J}_{spin}$', 'white')
im = plot_stream_panel(axes[2], J_total_x, J_total_y, r'Total Current $\vec{J}_{tot}$', 'lime')

cbar = fig.colorbar(im, ax=axes, shrink=0.8, pad=0.02)
cbar.set_label(r'$|\psi|^2$', rotation=270, labelpad=15)
plt.show()

# 1D Cross-section plot
idx_y = N // 2
fig1, ax1d = plt.subplots(figsize=(5, 3.5))
ax1d.plot(x, J_spin_y[idx_y, :] / J_max, label=r"$J_{\mathrm{spin}}$")
ax1d.plot(x, J_orb_y[idx_y, :] / J_max, label=r"$J_{\mathrm{orb}}$")
ax1d.plot(x, J_total_y[idx_y, :] / J_max, label=r"$J_{\mathrm{tot}}$", linestyle='--')
ax1d.set_xlabel(r'$x/l_B$')
ax1d.set_ylabel(r'$J_y / J_{max}$')
ax1d.set_xlim(-Lx/1.5, Lx/1.5)
ax1d.legend(loc="upper right", frameon=False)
plt.show()

6. Real-Time Dynamics

We now evolve the system in real time using the Schrödinger equation. We employ a 4th-order Runge-Kutta (RK4) method to ensure numerical stability.

$$i\hbar \frac{\partial \psi}{\partial t} = H\psi$$

def compute_time_derivative(psi_state, hbar=1):
    H_gp = calc_gp_potential(psi_state, g) * psi_state
    
    # FIX 1: Removed 'dx, dy' from the arguments so they don't overwrite 'hbar' and 'm'
    H_kin = kinetic_hamiltonian(psi_state, A_x, A_y, A_t, KX, KY)
    
    H_sp = spin_hamiltonian(psi_state, B)
    H_trap = potential_well * psi_state

    H_total = H_gp + H_kin + H_sp + H_trap
    return -1j * H_total / hbar

# FIX 2: Sudden change in the electric field to induce motion. 
# If the ground state was calculated with E=0.5, setting E=0.0 will make it oscillate.
# For a much stronger kick, you can try E_new = 2.0
E_new = 0.0 
A_t[0, :, :] = E_new * X 
A_t[1, :, :] = E_new * X 

def rk4_step(psi_state, dt):
    """Runge-Kutta 4th order solver step."""
    k1 = compute_time_derivative(psi_state)
    k2 = compute_time_derivative(psi_state + 0.5 * dt * k1)
    k3 = compute_time_derivative(psi_state + 0.5 * dt * k2)
    k4 = compute_time_derivative(psi_state + dt * k3)

    return psi_state + (k1 + 2 * k2 + 2 * k3 + k4) * dt / 6

# FIX 3: Adjusted time parameters for visible dynamics
dt = 0.002            # Larger time steps
max_time_steps = 2000 # Total simulated time = 4.0 (Enough to see movement)
frames = []

def get_total_density(p):
    return np.abs(p[0])**2 + np.abs(p[1])**2

frames.append(get_total_density(psi))
print("Starting real-time evolution...")

for step in range(max_time_steps):
    psi = rk4_step(psi, dt)
    
    # Save frames more frequently
    if step % 20 == 0:
        frames.append(get_total_density(psi))

print("Evolution complete. Generating animation...")

# Setup Animation
fig_anim, ax_anim = plt.subplots(figsize=(5, 4))
ax_anim.set_xlabel(r'$x / l_B$')
ax_anim.set_ylabel(r'$y / l_B$')
ax_anim.set_title(r'Time Evolution $|\psi|^2$')

# Calculate global maximum so the color scale remains fixed across frames
vmax_global = np.max([np.max(f) for f in frames])

im_anim = ax_anim.imshow(frames[0], cmap='inferno', origin='lower', 
                         extent=[-Lx, Lx, -Ly, Ly], vmin=0, vmax=vmax_global)
fig_anim.colorbar(im_anim, ax=ax_anim, label=r'$|\psi|^2$')

def update_frame(frame_idx):
    im_anim.set_data(frames[frame_idx])
    return [im_anim]

ani = animation.FuncAnimation(fig_anim, update_frame, frames=len(frames), blit=True)
plt.close() # Prevent static plot from showing, only the animation below

# To display the animation natively in Jupyter Notebook:
from IPython.display import HTML
HTML(ani.to_jshtml())

Starting real-time evolution...
Evolution complete. Generating animation...

Once Loop Reflect