Calculations/ODT-Calculator/+Calculator/computeTrapPotential.m

function [Positions, IdealTrappingPotential, TrappingPotential, TrapDepthsInKelvin, ExtractedTrapFrequencies] = computeTrapPotential(w_x, w_z, Power, options)
    
    alpha       = options.Polarizability;

    % Unpack the options struct
    axis        = options.Axis;
    extent      = options.Extent;
    gravity     = options.Gravity;
    astigmatism = options.Astigmatism;
    modulation  = options.Modulation;
    crossed     = options.Crossed;
    
    % Apply modulation if necessary
    if modulation
        aspect_ratio = options.AspectRatio;
        current_ar   = w_x / w_z;
        w_x          = w_x * (aspect_ratio / current_ar);
    end
    
    % Compute ideal trap depth
    TrapDepth              = Calculator.computeTrapDepth(w_x, w_z, Power, alpha);
    IdealTrapDepthInKelvin = (TrapDepth / physconst('boltzmann')) * 1E6; 
    
    % Default projection axis
    projection_axis = [0; 1; 0]; % Default Y-axis

    % Set projection axis based on the given axis option
    if axis == 1
        projection_axis = [1; 0; 0]; % X-axis
    elseif axis == 2
        projection_axis = [0; 1; 0]; % Y-axis
    elseif axis == 3
        projection_axis = [0; 0; 1]; % Z-axis
    else
        projection_axis = [1; 1; 1]; % Diagonal
    end

    % Define positions along the specified extent
    x_Positions = (-extent:0.05:extent) * 1E-6;
    y_Positions = (-extent:0.05:extent) * 1E-6;
    z_Positions = (-extent:0.05:extent) * 1E-6;
    
    % Stack positions along the projection axis
    Positions = [x_Positions; y_Positions; z_Positions] .* projection_axis;
    
    % Define the waists for the beams
    if crossed
        Waists = [w_x(1), w_z(1); w_x(2), w_z(2)];
        IdealTrappingPotential = Potentials.generateCrossedBeamPotential(Positions, Waists, Power, options); % Calculate the ideal trapping potential
    elseif modulation
        Waists = [w_x, w_z];
        IdealTrappingPotential = Potentials.generateModulatedSingleGaussianBeamPotential(Positions, Waists, Power, options);
    else 
        Waists = [w_x, w_z];
        IdealTrappingPotential = Potentials.generateSingleGaussianBeamPotential(Positions, Waists, Power, options);
    end
    
    IdealTrappingPotential = IdealTrappingPotential .* (ones(3, length(IdealTrappingPotential)) .* projection_axis);
    IdealTrappingPotential = (IdealTrappingPotential / physconst('boltzmann')) * 1E6; % Convert to uK
    
    % Initialize the TrappingPotential variable
    TrappingPotential = [];
    
    % Handle different combinations of gravity and astigmatism
    if gravity && ~astigmatism
        % Gravity effect only
        m = options.Mass;
        gravity_axis = [0; 0; 1]; % Z-axis for gravity
        if options.TiltGravity
            R = Helpers.generateRotationMatrix(options.tilt_axis, deg2rad(options.theta));
            gravity_axis = R * gravity_axis;
        end
        gravity_axis_positions = [x_Positions; y_Positions; z_Positions] .* gravity_axis;
        if crossed
            TrappingPotential = Potentials.generateCrossedBeamPotential(Positions, Waists, Power, options);
        elseif modulation
            TrappingPotential = Potentials.generateModulatedSingleGaussianBeamPotential(Positions, Waists, Power, options);
        else
            TrappingPotential = Potentials.generateSingleGaussianBeamPotential(Positions, Waists, Power, options);
        end
        TrappingPotential = TrappingPotential .* (ones(3, length(TrappingPotential)) .* projection_axis) + Potentials.generateGravitationalPotential(gravity_axis_positions, m);
        TrappingPotential = (TrappingPotential / physconst('boltzmann')) * 1E6; % Convert to uK

    elseif ~gravity && astigmatism
        % Astigmatism effect only
        if crossed
            TrappingPotential = Potentials.generateAstigmaticCrossedBeamPotential(Positions, Waists, Power, options);
        else
            TrappingPotential = Potentials.generateAstigmaticSingleGaussianBeamPotential(Positions, Waists, Power, options);
        end
        TrappingPotential = TrappingPotential .* (ones(3, length(TrappingPotential)) .* projection_axis);
        TrappingPotential = (TrappingPotential / physconst('boltzmann')) * 1E6; % Convert to uK

    elseif gravity && astigmatism
        % Both gravity and astigmatism
        m = options.Mass;
        gravity_axis = [0; 0; 1]; % Z-axis for gravity
        if options.TiltGravity
            R = Helpers.generateRotationMatrix(options.tilt_axis, deg2rad(options.theta));
            gravity_axis = R * gravity_axis;
        end
        gravity_axis_positions = [x_Positions; y_Positions; z_Positions] .* gravity_axis;
        if crossed
            TrappingPotential = Potentials.generateAstigmaticCrossedBeamPotential(Positions, Waists, Power, options);
        else
            TrappingPotential = Potentials.generateAstigmaticSingleGaussianBeamPotential(Positions, Waists, Power, options);
        end
        TrappingPotential = TrappingPotential .* (ones(3, length(TrappingPotential)) .* projection_axis) + Potentials.generateGravitationalPotential(gravity_axis_positions, m);
        TrappingPotential = (TrappingPotential / physconst('boltzmann')) * 1E-6; % Convert to uK

    else
        % Default to ideal trapping potential
        TrappingPotential = IdealTrappingPotential;
    end

    % Calculate inflection points and effective trap depth
    infls = find(diff(sign(gradient(gradient(TrappingPotential(axis,:))))));
    
    try
        if TrappingPotential(axis, 1) > TrappingPotential(axis, end)
            EffectiveTrapDepthInKelvin = max(TrappingPotential(axis, infls(2):end)) - min(TrappingPotential(axis, infls(1):infls(2)));
        elseif TrappingPotential(axis, 1) < TrappingPotential(axis, end)
            EffectiveTrapDepthInKelvin = max(TrappingPotential(axis, 1:infls(1))) - min(TrappingPotential(axis, infls(1):infls(2)));
        else
            EffectiveTrapDepthInKelvin = IdealTrapDepthInKelvin;
        end
    catch
        EffectiveTrapDepthInKelvin = NaN;
    end

    % Trap depths in Kelvin
    TrapDepthsInKelvin = [IdealTrapDepthInKelvin, EffectiveTrapDepthInKelvin];

    % Extract trap frequencies from the ideal trapping potential
    [v, dv] = Calculator.extractTrapFrequency(Positions, IdealTrappingPotential, options);
    if isinf(v), v = NaN; end
    if isinf(dv), dv = NaN; end
    IdealTrapFrequency = [v, dv];

    % Extract trap frequencies from the actual trapping potential
    if options.ExtractTrapFrequencies
        [v, dv] = Calculator.extractTrapFrequency(Positions, TrappingPotential, options);
        if isinf(v), v = NaN; end
        if isinf(dv), dv = NaN; end
        TrapFrequency = [v, dv];
        ExtractedTrapFrequencies = {IdealTrapFrequency, TrapFrequency};
    else
        ExtractedTrapFrequencies = {IdealTrapFrequency};
    end

end