Split-Step Fourier Method

[Abdelhakim95Ben]

I hope everyone is doing well.

Could you please confirm if this code is correct for solving the Gross-Pitaevskii equation numerically for the case of a harmonic oscillator?

Thank you.

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

Constants

hbar = 1.0 # Reduced Planck constant
m = 1.0 # Mass of the particle
g = 1.0 # Interaction strength
L = 10.0 # Length of the spatial domain
N = 256 # Number of spatial grid points
dx = L / N # Spatial grid step
dt = 0.01 # Time step for evolution
t_max = 5.0 # Maximum time for evolution

Spatial and momentum grids

x = np.linspace(-L / 2, L / 2, N, endpoint=False)
k = 2 * np.pi * np.fft.fftfreq(N, d=dx) # Momentum space grid

Potential: Harmonic oscillator (V = 0.5 * x^2)

V = 0.5 * x**2

Initial wavefunction (Gaussian)

sigma_init = 1.0
psi0 = (1 / (np.pi0.25 * sigma_init0.5)) * np.exp(-x2 / (2 * sigma_init2)).astype(np.complex128)

Normalize initial wavefunction

norm_psi0 = np.sqrt(np.sum(np.abs(psi0)**2) * dx)
psi0 /= norm_psi0

Split-step Fourier method function

def split_step_fourier(psi, V, g, k, dt, steps):
"""
Solves the GPE using the Split-Step Fourier Method.

Parameters:
    psi (numpy.ndarray): Initial wavefunction.
    V (numpy.ndarray): Potential array.
    g (float): Nonlinear interaction strength.
    k (numpy.ndarray): Momentum space array.
    dt (float): Time step.
    steps (int): Number of time steps.
    
Returns:
    numpy.ndarray: Final wavefunction.
"""
psi_k = np.fft.fft(psi)  # Initial wavefunction in momentum space

for step in range(steps):
    # Nonlinear and potential evolution in real space
    psi = psi * np.exp(-1j * (V + g * np.abs(psi)**2) * dt / 2)
    
    # Kinetic evolution in momentum space
    psi_k = np.fft.fft(psi)
    psi_k *= np.exp(-1j * (hbar * k**2 / (2 * m)) * dt)
    psi = np.fft.ifft(psi_k)
    
    # Nonlinear and potential evolution in real space
    psi = psi * np.exp(-1j * (V + g * np.abs(psi)**2) * dt / 2)

return psi

Number of time steps

steps = int(t_max / dt)

Start timing

start_time = time.time()

Solve GPE using SSFM

psi_numerical = split_step_fourier(psi0, V, g, k, dt, steps)

End timing

end_time = time.time()

Normalize numerical solution

norm_numerical = np.sqrt(np.sum(np.abs(psi_numerical)**2) * dx)
psi_numerical /= norm_numerical

Plotting results

plt.figure(figsize=(12, 6))

Plot the initial and final wavefunction densities

plt.subplot(1, 2, 1)
plt.plot(x, np.abs(psi0)**2, label="Initial |ψ(x)|²", lw=2)
plt.plot(x, np.abs(psi_numerical)**2, label="Numerical |ψ(x)|²", linestyle="--", lw=2)
plt.title("Wavefunction: Initial vs Numerical Solution")
plt.xlabel("Position (x)")
plt.ylabel("Density |ψ(x)|²")
plt.legend()
plt.grid()

Plot the real part of the final wavefunction

plt.subplot(1, 2, 2)
plt.plot(x, np.real(psi_numerical), label="Re(ψ(x))", color="blue", lw=2)
plt.title("Real Part of the Final Wavefunction")
plt.xlabel("Position (x)")
plt.ylabel("Re(ψ(x))")
plt.legend()
plt.grid()

Show timing

plt.figtext(0.15, 0.85, f"Execution Time: {end_time - start_time:.2f} seconds", fontsize=12)

plt.tight_layout()
plt.show()