function visualizeGSWavefunction2D(run_index)
    
    set(0,'defaulttextInterpreter','latex')
    set(groot, 'defaultAxesTickLabelInterpreter','latex'); set(groot, 'defaultLegendInterpreter','latex');
    
    folder_path = './Data';

    % Load data
    Data = load(sprintf(horzcat(folder_path, '/Run_%03i/psi_gs.mat'),run_index),'psi','Params','Transf','Observ');
    
    Params = Data.Params;
    Transf = Data.Transf; 
    Observ = Data.Observ; 
    if isgpuarray(Data.psi)
        psi = gather(Data.psi);
    else
        psi = Data.psi;
    end
    
    if isgpuarray(Data.Observ.residual)
        Observ.residual = gather(Data.Observ.residual);
    else
        Observ.residual = Data.Observ.residual;
    end

    % Set long format for output
    format long
    
    % Coordinates in micrometers
    x = Transf.x * Params.l0 * 1e6; 
    y = Transf.y * Params.l0 * 1e6; 
    
    % Plotting
    figure('Position', [100, 100, 1600, 900]);
    clf
    
    % Compute probability density |psi|^2
    n = abs(psi).^2;
    
    % Plot |psi(x,y)|^2 (density in x and y plane)
    subplot('Position', [0.05, 0.55, 0.28, 0.4])
    pcolor(x, y, n');
    shading interp;
    colorbar;
    xlabel('$x$ [$\mu$m]', 'FontSize', 14); ylabel('$y$ [$\mu$m]', 'FontSize', 14);
    title('$|\Psi_{xy}|^2$', 'FontSize', 14);
    
    % Plot real part of psi in the x-y plane
    subplot('Position', [0.36, 0.55, 0.28, 0.4])
    pcolor(x, y, real(psi)');
    shading interp;
    colorbar;
    xlabel('$x$ [$\mu$m]', 'FontSize', 14); ylabel('$y$ [$\mu$m]', 'FontSize', 14);
    title('Re$\{\Psi_{xy}\}$', 'FontSize', 14);

    % Plot imaginary part of psi in the x-y plane
    subplot('Position', [0.67, 0.55, 0.28, 0.4])
    pcolor(x, y, imag(psi)');
    shading interp;
    colorbar;
    xlabel('$x$ [$\mu$m]', 'FontSize', 14); ylabel('$y$ [$\mu$m]', 'FontSize', 14);
    title('Im$\{\Psi_{xy}\}$', 'FontSize', 14);

    % Plot residual (time steps vs -log10(residual))
    subplot('Position', [0.05, 0.05, 0.26, 0.4]);
    plot(-log10(Observ.residual), '-b')
    ylabel('$-\mathrm{log}_{10}(r)$', 'FontSize', 14); xlabel('Time steps', 'FontSize', 14);
    title('Residual', 'FontSize', 14);

    % Plot total energy over time
    subplot('Position', [0.36, 0.05, 0.26, 0.4]);
    plot(Observ.EVec, '-b')
    ylabel('$E$', 'FontSize', 14); xlabel('Time steps', 'FontSize', 14);
    title('Total Energy', 'FontSize', 14);

    % Plot chemical potential over time
    subplot('Position', [0.67, 0.05, 0.26, 0.4]);
    plot(Observ.mucVec, '-b')
    ylabel('$\mu$', 'FontSize', 14); xlabel('Time steps', 'FontSize', 14);
    title('Chemical Potential', 'FontSize', 14);
end