Solving Schrodinger-Poisson system using RK-Spectral method

schrodinger_poisson_rk_spectral.m

% Solving Schrodinger-Poisson system
%
% Reference:
%   P. K. Shukla and B. Eliasson. Formation and Dynamics of Dark Solitons
%   and Vortices in Quantum Electron Plasmas. Physical Review Letters, 96,
%   245001, 2006. https://doi.org/10.1103/physrevlett.96.245001
%   (The charge states on page 3:
%        The sign of n_i seems wrong. Should be n1 = -1 , n2 = 1, n3 = 1, n4 = -1)
%   

% ========= 参数设置 =========
Nx = 256; Ny = 256;      % 网格点数
Lx = 40; Ly = 40;        % 空间区域 [-L/2, L/2]
x = linspace(-Lx/2, Lx/2, Nx+1); x(end) = [];
y = linspace(-Ly/2, Ly/2, Ny+1); y(end) = [];
[X, Y] = meshgrid(x, y);

% FFT 波数向量
kx = (2*pi/Lx) * [0:Nx/2 -Nx/2+1:-1];
ky = (2*pi/Ly) * [0:Ny/2 -Ny/2+1:-1];
[KX, KY] = meshgrid(kx, ky);
K2 = KX.^2 + KY.^2;

% ===== 量子参数与时间推进设置 =====
A = 5;
dt = 1/256/8;
T = 10;
Nt = round(T / dt);

% ===== 初始条件：四个涡旋 =====
P = 10;
n = 1;
psi = vortex(-4, P, -n, X, Y) .* ...
      vortex( 2, P, n, X, Y) .* ...
      vortex(-2,-P, n, X, Y) .* ...
      vortex( 4,-P, -n, X, Y);

% P = 0;
% n = -2;
% psi = vortex(-3, P, n, X, Y) .* ...
%       vortex( 3, P, n, X, Y);

[EX, EY] = electron(psi, KX, KY);

% 创建 VideoWriter 对象
v = VideoWriter(sprintf('sp-rks-n=%d-1overdt=%d.avi', n, round(1/dt)), 'Motion JPEG AVI');
v.FrameRate = 8;  % 每秒帧数，可根据需要调整
open(v);  % 打开视频写入对象


fig = figure;
ax = axes('Position', [0.1 0.15 0.35 0.7]);
ax2 = axes('InnerPosition', [0.55, 0.15, 0.28, 0.7]);

plot_density(ax, x, y, psi, 0);
plot_electron_current(ax2, X, Y, EX, EY, 0);

frame = getframe(gcf);
writeVideo(v, frame);

M = 64;
save(sprintf("rk-data/psi-%d-%05d.mat", M, 0), "psi")

sigma = 0;
% ===== 时间推进：RK4 =====
for nt = 1:Nt
    psi = rk4_step(psi, K2, A, dt, sigma);
    [EX, EY] = electron(psi, KX, KY);

    % 可视化
    if mod(nt, M) == 0
	    fprintf("%d/%d\n", round(nt/M), round(Nt/M))
        save(sprintf("rk-data/psi-%d-%05d.mat", M, round(nt/M)), "psi")

        plot_density(ax, x, y, psi, nt*dt)
        plot_electron_current(ax2, X, Y, EX, EY, nt*dt);

        frame = getframe(gcf);
        writeVideo(v, frame);
    end
end
close(v);

% ===== subfunctions =====
function [EX, EY] = electron(psi, KX, KY)
    psi_hat = fft2(psi);
    EX = 2*imag(conj(psi).*ifft2(1i*KX.*psi_hat));
    EY = 2*imag(conj(psi).*ifft2(1i*KY.*psi_hat));
end

function v = vortex(x0, y0, n, X, Y)
    R2 = (X - x0).^2 + (Y - y0).^2;
    theta = atan2(Y - y0, X - x0);
    v = tanh(sqrt(R2)) .* exp(1i * n * theta);
end

function psi_new = rk4_step(psi, K2, A, dt, sigma)
    k1 = rhs(psi, K2, A, sigma);
    k2 = rhs(psi + 0.5*dt*k1, K2, A, sigma);
    k3 = rhs(psi + 0.5*dt*k2, K2, A, sigma);
    k4 = rhs(psi + dt*k3, K2, A, sigma);
    psi_new = psi + dt/6 * (k1 + 2*k2 + 2*k3 + k4);
end

function dpsi = rhs(psi, K2, A, sigma)
    rho = abs(psi).^2;

    % 解 Poisson 方程 ∇²φ = ρ - 1，用 FFT（周期边界）
    rho_hat = fft2(rho - 1);
    K2_safe = K2; K2_safe(K2==0) = 1;         % 避免除以 0
    phi_hat = rho_hat ./ (-K2_safe); 
    phi_hat(K2==0) = 0;                       % 设置 φ 平均值为 0
    phi = real(ifft2(phi_hat));

    % 计算拉普拉斯 ∇²ψ
    psi_hat = fft2(psi);
    lap_psi = ifft2(-K2 .* psi_hat);

    % % Schrödinger 方程右端
    dpsi = 1i * (A * lap_psi  + phi .* psi - rho .* psi)...
        - sigma .* psi;
end

function plot_density(ax, x, y, psi, t)
    imagesc(ax, x, y, abs(psi).^2);
    axis(ax, 'equal', 'tight');
    set(ax, 'YDir', 'normal');
    title(ax, sprintf('|\\psi|^2 at t = %.3f', t));
    colormap(ax, parula);
    colorbar(ax);
end

function plot_electron_current(ax2, X, Y, EX, EY, t)
    quiver_sparse(ax2, X, Y, EX, EY, 4, 1)
    axis(ax2, 'equal', 'tight');
    title(sprintf('Electron current at t = %.3f', t));
end

function quiver_sparse(ax, x, y, u, v, step, scale, varargin)
%QUIVER_SPARSE 稀疏绘制矢量场箭头
%   quiver_sparse(ax, x, y, u, v, step)
%   quiver_sparse(ax, x, y, u, v, step, scale)
%   quiver_sparse(ax, x, y, u, v, step, scale, 'LineWidth', 2, ...)

% 参数说明：
% x, y    —— 网格坐标矩阵
% u, v    —— 矢量场
% step    —— 抽样间隔（整数）
% scale   —— 箭头缩放系数（可选，默认1）
% varargin—— 传给 quiver 的额外绘图参数（例如颜色、线宽等）

    if nargin < 7 || isempty(scale)
        scale = 1;
    end

    % 下采样
    x_sparse = x(1:step:end, 1:step:end);
    y_sparse = y(1:step:end, 1:step:end);
    u_sparse = u(1:step:end, 1:step:end);
    v_sparse = v(1:step:end, 1:step:end);

    % 绘图
    quiver(ax, x_sparse, y_sparse, u_sparse, v_sparse, scale, varargin{:});
end