New scripting method to efficiently run the solver over a wide parameter range on the cluster.

This commit is contained in:
Karthik 2025-05-02 01:29:57 +02:00
parent 4d10f6e515
commit 8bc8ab71d6
4 changed files with 177 additions and 11 deletions

View File

@ -0,0 +1,71 @@
function run_hybrid_worker(batchParams, batchIdx)
% Set up local cluster for parallel pool
cluster = parcluster('local');
nprocs = str2double(getenv('SLURM_CPUS_PER_TASK'));
if isnan(nprocs), nprocs = feature('numcores'); end
tmpdir = fullfile(getenv('TMPDIR'), sprintf('matlab_job_%d', batchIdx));
if ~exist(tmpdir, 'dir'); mkdir(tmpdir); end
cluster.JobStorageLocation = tmpdir;
pool = parpool(cluster, nprocs);
nJobs = size(batchParams, 1);
parfor k = 1:nJobs
% Unpack parameter tuple
a_s = batchParams(k, 1);
theta_deg = batchParams(k, 2);
phi_deg = batchParams(k, 3);
N_atoms = batchParams(k, 4);
theta_rad = deg2rad(theta_deg);
phi_rad = deg2rad(phi_deg);
% Create unique save directory
jobName = sprintf('aS_%03d_theta_%03d_phi_%03d_N_%d', a_s, theta_deg, phi_deg, N_atoms);
saveDir = fullfile('./Results/Data_3D/GradientDescent', jobName);
if ~exist(saveDir, 'dir')
mkdir(saveDir);
end
% Options for this run
OptionsStruct = struct;
OptionsStruct.NumberOfAtoms = N_atoms;
OptionsStruct.DipolarPolarAngle = theta_rad;
OptionsStruct.DipolarAzimuthAngle = phi_rad;
OptionsStruct.ScatteringLength = a_s;
OptionsStruct.TrapFrequencies = [50, 20, 150];
OptionsStruct.TrapPotentialType = 'Harmonic';
OptionsStruct.NumberOfGridPoints = [128, 256, 128];
OptionsStruct.Dimensions = [30, 50, 30];
OptionsStruct.UseApproximationForLHY = true;
OptionsStruct.IncludeDDICutOff = true;
OptionsStruct.CutoffType = 'Cylindrical';
OptionsStruct.SimulationMode = 'EnergyMinimization';
OptionsStruct.GradientDescentMethod = 'NonLinearCGD';
OptionsStruct.MaxIterationsForGD = 15000;
OptionsStruct.NoiseScaleFactor = 0.010;
OptionsStruct.PlotLive = false;
OptionsStruct.JobNumber = 0;
OptionsStruct.RunOnGPU = true;
OptionsStruct.SaveData = true;
OptionsStruct.SaveDirectory = saveDir;
options = Helper.convertstruct2cell(OptionsStruct);
sim = Simulator.DipolarGas(options{:});
pot = Simulator.Potentials(options{:});
sim.Potential = pot.trap();
NumberOfOutputs = 5;
[~, ~, ~, ~, ~, stats] = Helper.runWithProfiling(@() sim.run(), NumberOfOutputs, saveDir);
fprintf('Batch %d | Job %d: a_s = %d, theta = %d°, phi = %d°, N = %d | Time = %.2f s\n', ...
batchIdx, k, a_s, theta_deg, phi_deg, N_atoms, stats.runtime);
end
delete(pool);
end

View File

@ -0,0 +1,35 @@
function run_hybrid_worker_wrapper(batchIdx)
a_s_list = parse_environmental_variable('SCATTERING_LENGTH_RANGE', 85); % Scattering length
theta_list = parse_environmental_variable('POLAR_ANGLE_RANGE', 0); % Polar angle
phi_list = parse_environmental_variable('AZIMUTHAL_ANGLE_RANGE', 0); % Azimuthal angle
N_atoms_list = parse_environmental_variable('NUM_ATOMS_LIST', 90000); % Atom number
chunkSize = str2double(getenv('CHUNK_SIZE'));
% Create full parameter grid (extendable)
[A, T, P, N] = ndgrid(a_s_list, theta_list, phi_list, N_atoms_list);
paramGrid = [A(:), T(:), P(:), N(:)];
totalJobs = size(paramGrid, 1);
totalBatches = ceil(totalJobs / chunkSize);
if batchIdx > totalBatches
error('Batch index %d exceeds total batches (%d)', batchIdx, totalBatches);
end
firstIdx = (batchIdx - 1) * chunkSize + 1;
lastIdx = min(batchIdx * chunkSize, totalJobs);
% Call worker with this batch of parameters
batchParams = paramGrid(firstIdx:lastIdx, :);
Scripts.run_hybrid_worker(batchParams, batchIdx);
end
function vals = parse_environmental_variable(varName, default)
str = getenv(varName);
if isempty(str)
vals = default;
elseif startsWith(str, '[')
vals = str2num(str); %#ok<ST2NM>
else
vals = eval(str); % Trust only controlled environments
end
end

View File

@ -3,29 +3,24 @@ OptionsStruct = struct;
OptionsStruct.NumberOfAtoms = 90000;
OptionsStruct.DipolarPolarAngle = deg2rad(0);
OptionsStruct.DipolarAzimuthAngle = 0;
OptionsStruct.ScatteringLength = 95;
OptionsStruct.ScatteringLength = 85;
OptionsStruct.TrapFrequencies = [50, 20, 150];
OptionsStruct.TrapPotentialType = 'Harmonic';
OptionsStruct.NumberOfGridPoints = [64, 128, 64];
OptionsStruct.Dimensions = [30, 30, 30];
OptionsStruct.NumberOfGridPoints = [128, 256, 128];
OptionsStruct.Dimensions = [30, 50, 30];
OptionsStruct.UseApproximationForLHY = true;
OptionsStruct.IncludeDDICutOff = true;
OptionsStruct.CutoffType = 'Cylindrical';
OptionsStruct.SimulationMode = 'EnergyMinimization'; % 'ImaginaryTimeEvolution' | 'RealTimeEvolution' | 'EnergyMinimization'
OptionsStruct.GradientDescentMethod = 'NonLinearCGD'; % 'HeavyBall' | 'NonLinearCGD'
OptionsStruct.MaxIterationsForGD = 1E5;
OptionsStruct.TimeStepSize = 1E-3; % in s
OptionsStruct.MinimumTimeStepSize = 1E-6; % in s
OptionsStruct.TimeCutOff = 2E6; % in s
OptionsStruct.EnergyTolerance = 5E-10;
OptionsStruct.ResidualTolerance = 1E-08;
OptionsStruct.MaxIterationsForGD = 15000;
OptionsStruct.NoiseScaleFactor = 0.010;
OptionsStruct.PlotLive = true;
OptionsStruct.PlotLive = false;
OptionsStruct.JobNumber = 0;
OptionsStruct.RunOnGPU = false;
OptionsStruct.RunOnGPU = true;
OptionsStruct.SaveData = true;
OptionsStruct.SaveDirectory = './Results/Data_3D/GradientDescent'; % './Results/Data_3D/AnisotropicTrap/Tilted0'
options = Helper.convertstruct2cell(OptionsStruct);

View File

@ -0,0 +1,65 @@
# Define scan ranges (encoded as strings)
SCATTERING_LENGTH_RANGE="85"
POLAR_ANGLE_RANGE="0"
AZIMUTHAL_ANGLE_RANGE="0"
NUM_ATOMS_LIST="[90000]"
CHUNK_SIZE=1
# Convert parameter ranges into arrays
scatteringLengths=($(eval echo {$SCATTERING_LENGTH_RANGE}))
polarAngles=($(eval echo {$POLAR_ANGLE_RANGE}))
azimuthalAngles=($(eval echo {$AZIMUTHAL_ANGLE_RANGE}))
numAtoms=($(eval echo {$NUM_ATOMS_LIST}))
# Calculate the total number of jobs (combinations of parameters)
totalJobs=$((${#scatteringLengths[@]} * ${#polarAngles[@]} * ${#azimuthalAngles[@]} * ${#numAtoms[@]}))
# Calculate the number of array jobs based on total jobs and chunk size
numArrayJobs=$(( (totalJobs + CHUNK_SIZE - 1) / CHUNK_SIZE ))
# Print the total number of jobs and array jobs for debugging
echo "Total number of jobs: $totalJobs"
echo "Number of SLURM array jobs: $numArrayJobs"
# Create the SLURM job submission command
sbatch --export=SCATTERING_LENGTH_RANGE="$SCATTERING_LENGTH_RANGE",POLAR_ANGLE_RANGE="$POLAR_ANGLE_RANGE",AZIMUTHAL_ANGLE_RANGE="$AZIMUTHAL_ANGLE_RANGE",NUM_ATOMS_LIST="$NUM_ATOMS_LIST",CHUNK_SIZE=$CHUNK_SIZE << EOF
#!/bin/bash
########### Begin SLURM header ###########
#SBATCH --partition=gpu-single
#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --cpus-per-task=8
#SBATCH --gres=gpu:1
#SBATCH --mem=16G
#SBATCH --time=00:40:00
#SBATCH --job-name=simulation
#SBATCH --error=simulation_%A_%a.err
#SBATCH --output=simulation_%A_%a.out
#SBATCH --array=1-${numArrayJobs}
########### End SLURM header ##########
echo "Working Directory: $PWD"
echo "Running on host $HOSTNAME"
echo "Job id: $SLURM_JOB_ID"
echo "Job name: $SLURM_JOB_NAME"
echo "Number of nodes allocated to job: $SLURM_JOB_NUM_NODES"
echo "Number of GPUs allocated to job: $SLURM_GPUS"
# Load module
module load math/matlab/R2023a
echo Directory is `pwd`
echo "Initiating Job..."
# Inside SLURM job array
echo "Running SLURM job array from 1 to ${numArrayJobs}"
# Start MATLAB job with the batch index
matlab -nodisplay -nosplash -r "Scripts.run_hybrid_worker_wrapper(\$SLURM_ARRAY_TASK_ID)"
# notice for tests
echo "Job terminated successfully"
exit
EOF