Calculations/Dipolar-Gas-Simulator/+Scripts/extractThomasFermiRadii.m

function [Rx, Ry] = extractThomasFermiRadii(folder_path, run_index, SuppressPlotFlag)    
    set(0,'defaulttextInterpreter','latex')
    set(groot, 'defaultAxesTickLabelInterpreter','latex'); set(groot, 'defaultLegendInterpreter','latex');
    
    format long
    
    % 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
    
    % Axes scaling and coordinates in micrometers
    x                       = Transf.x * Params.l0 * 1e6; 
    y                       = Transf.y * Params.l0 * 1e6;
    z                       = Transf.z * Params.l0 * 1e6;

    dx = x(2)-x(1); dy = y(2)-y(1); dz = z(2)-z(1);
    
    % Compute probability density |psi|^2
    n                       = abs(psi).^2;
    nxz                     = squeeze(trapz(n*dy,2));
    nyz                     = squeeze(trapz(n*dx,1));
    nxy                     = squeeze(trapz(n*dz,3));
    
    state                   = nxy';

    % Fit
    fitResult               = fitThomasFermiProfile2D(state, x, y);
    Rx                      = fitResult.Rx;
    Ry                      = fitResult.Ry;

    if ~SuppressPlotFlag
        fig = figure(1); 
        clf
        set(gcf,'Position', [100, 100, 1800, 500])
        t = tiledlayout(1, 3, 'TileSpacing', 'compact', 'Padding', 'compact');

        % Original density plot
        nexttile;
        plot1 = pcolor(x, y, state);
        set(plot1, 'EdgeColor', 'none');
        colormap(gca, Helper.Colormaps.plasma())
        colorbar;
        xlabel('$x$ ($\mu$m)', 'Interpreter', 'latex', 'FontSize', 14)
        ylabel('$y$ ($\mu$m)', 'Interpreter', 'latex', 'FontSize', 14)
        title('Original Density', 'Interpreter', 'latex', 'FontSize', 16)
        axis square;
        set(gca, 'FontSize', 14);
    
        % Fitted TF profile plot
        nexttile;
        plot2 = pcolor(x, y, fitResult.fittedDensity);
        set(plot2, 'EdgeColor', 'none');
        colormap(gca, Helper.Colormaps.plasma())
        colorbar;
        xlabel('$x$ ($\mu$m)', 'Interpreter', 'latex', 'FontSize', 14)
        ylabel('$y$ ($\mu$m)', 'Interpreter', 'latex', 'FontSize', 14)
        title(sprintf('TF Fit: $R_x = %.2f\\,\\mu\\mathrm{m}$, $R_y = %.2f\\,\\mu\\mathrm{m}$', Rx, Ry), ...
                    'Interpreter', 'latex', 'FontSize', 16)
        axis square;
        set(gca, 'FontSize', 14);

                % 1D cuts through center
        nexttile;
        % Find index of center
        [~, ix0] = min(abs(x - fitResult.x0));
        [~, iy0] = min(abs(y - fitResult.y0));

        % X-cut: density along y = y0
        orig_cut_x = state(iy0, :);
        fit_cut_x  = fitResult.fittedDensity(iy0, :);

        % Y-cut: density along x = x0
        orig_cut_y = state(:, ix0);
        fit_cut_y  = fitResult.fittedDensity(:, ix0);

        % Plot both cuts
        plot(x, orig_cut_x, 'w-', 'LineWidth', 1.5); hold on
        plot(x, fit_cut_x,  'r--', 'LineWidth', 1.5)
        plot(y, orig_cut_y, 'w-', 'LineWidth', 1.5)
        plot(y, fit_cut_y,  'r--', 'LineWidth', 1.5)
        hold off
        axis square;
        legend({'$n(x)$', '$n_{\mathrm{TF}}(x)$', '$n(y)$', '$n_{\mathrm{TF}}(y)$'}, ...
               'Interpreter', 'latex', 'FontSize', 12, 'Location', 'northeast')

        xlabel('$x$ or $y$ ($\mu$m)', 'Interpreter', 'latex', 'FontSize', 14)
        ylabel('Density', 'Interpreter', 'latex', 'FontSize', 14)
        title('1-D Cut', 'Interpreter', 'latex', 'FontSize', 16)
        grid on
        set(gca, 'FontSize', 14)
    end
end 

function fitResult = fitThomasFermiProfile2D(state, xgrid, ygrid)
    % Fits a 2D Thomas-Fermi profile to the input BEC state
    %
    % Input:
    %   state   - 2D array (real density or complex wavefunction)
    %   xgrid   - vector of x-coordinates (size must match size(state,2))
    %   ygrid   - vector of y-coordinates (size must match size(state,1))
    %
    % Output (struct):
    %   fitResult.fittedDensity - 2D array of the fitted TF profile
    %   fitResult.n0            - peak density
    %   fitResult.x0, y0        - center coordinates
    %   fitResult.Rx, Ry        - Thomas-Fermi radii
    %   fitResult.params        - parameter vector [n0, x0, y0, Rx, Ry]

    % Compute density if input is a wavefunction
    if ~isreal(state)
        density = abs(state).^2;
    else
        density = state;
    end

    % Set up grid
    [X, Y]   = meshgrid(xgrid, ygrid);
    xy       = [X(:), Y(:)];
    n        = density(:);

    % Mask: ignore near-zero values
    mask     = n > max(n)*1e-3;
    xy       = xy(mask, :);
    n        = n(mask);

    % Initial guess: [n0, x0, y0, Rx, Ry]
    n0_guess = max(n);
    x0_guess = mean(xgrid(:));
    y0_guess = mean(ygrid(:));
    Rx_guess = (max(xgrid)-min(xgrid)) / 2;
    Ry_guess = (max(ygrid)-min(ygrid)) / 2;
    p0       = [n0_guess, x0_guess, y0_guess, Rx_guess, Ry_guess];

    % Bounds
    lb       = [0, min(xgrid), min(ygrid), 0, 0];
    ub       = [Inf, max(xgrid), max(ygrid), Inf, Inf];

    % TF 2D function
    tf2d     = @(p, xy) max(0, p(1) * (1 - ((xy(:,1)-p(2)).^2 / p(4)^2 + (xy(:,2)-p(3)).^2 / p(5)^2)));

    % Fit
    options  = optimoptions('lsqcurvefit','Display','off');
    pfit     = lsqcurvefit(tf2d, p0, xy, n, lb, ub, options);

    % Evaluate fitted profile over full grid
    fittedDensity = reshape(tf2d(pfit, [X(:), Y(:)]), size(X));

    % Return result
    fitResult.fittedDensity = fittedDensity;
    fitResult.n0 = pfit(1);
    fitResult.x0 = pfit(2);
    fitResult.y0 = pfit(3);
    fitResult.Rx = pfit(4);
    fitResult.Ry = pfit(5);
    fitResult.params = pfit;
end