imdd_silas/Functions/Job_Processing/submitJobs.m

function [results, wh] = submitJobs(run_ids, dsp_options, submit_mode, submit_options)
% SUBMITJOBS Submits dsp_runid jobs for processing (parallel or serial)
%
%   [results, wh] = submitJobs(run_ids, dsp_options, submit_mode, submit_options)
%
%   run_ids        : scalar or vector of run IDs
%   dsp_options    : struct of dsp_runid name/value options
%   submit_mode    : processingMode.parallel or .serial
%   submit_options : struct with fields
%       .waitbar (logical)
%       .wh      (DataStorage object)
%       .storePackages (string array, optional package whitelist)
%
%   results : cell(nJobsPerRunId, nRunIds)
%   wh      : updated DataStorage. For multiple run_ids, run_id is added as
%             the first storage axis.

arguments
    run_ids        int32               = 0
    dsp_options    struct              = struct()
    submit_mode    processingMode      = processingMode.serial
    submit_options.waitbar (1,1) logical = true
    submit_options.wh = DataStorage(struct())
    submit_options.storePackages string = string.empty()
end

% Normalize
run_ids = run_ids(:)';
nRunIds       = numel(run_ids);
sweepWh       = submit_options.wh;
validateNoRunIdSweepParameter(sweepWh);

nJobsPerRunId = sweepWh.getLastLinIndice();
totalJobs     = nRunIds * nJobsPerRunId;

wh = buildStorageWarehouse(sweepWh, run_ids);
results = cell(nJobsPerRunId, nRunIds);
jobs = repmat(struct('args', {{}}, 'label', "", 'meta', struct()), 1, totalJobs);

jobCounter = 0;
for r = 1:nRunIds
    for k = 1:nJobsPerRunId
        jobCounter = jobCounter + 1;
        userParameters = buildUserParameters(k, sweepWh);
        jobOptions = dsp_options;
        jobOptions.userParameters = userParameters;

        jobs(jobCounter).args = [{run_ids(r)}, structToNameValue(jobOptions)];
        jobs(jobCounter).label = sprintf('RunID %d, Job %d', run_ids(r), k);
        jobs(jobCounter).meta.runIndex = r;
        jobs(jobCounter).meta.jobIndex = k;
        jobs(jobCounter).meta.sweepJobIndex = k;
        jobs(jobCounter).meta.storageJobIndex = buildStorageJobIndex(run_ids(r), userParameters, wh);
        jobs(jobCounter).meta.run_id = run_ids(r);
    end
end

linearResults = runBatch(@dsp_runid, jobs, ...
    "mode", submit_mode, ...
    "waitbar", submit_options.waitbar, ...
    "waitbarMessage", "Processing Jobs...", ...
    "numWorkers", 11, ...
    "idleTimeout", 300, ...
    "cancelExistingQueue", true, ...
    "resultHandler", @storeResult, ...
    "errorHandler", @handleError);

for idx = 1:totalJobs
    r = jobs(idx).meta.runIndex;
    k = jobs(idx).meta.jobIndex;
    results{k, r} = linearResults{idx};
end

%% Local helpers %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

    function handleError(ME, job, ~)
        fprintf('[RunID %d, Job %d] ERROR [%s]: %s\n', job.meta.run_id, job.meta.jobIndex, ...
            ME.identifier, ME.message);
        for st = ME.stack'
            fprintf('    %s:%d   (%s)\n', st.file, st.line, st.name);
        end
        fprintf('Full report:\n%s\n', getReport(ME,'extended'));
    end

    function userParameters = buildUserParameters(jobIndex, wh)
        userParameters = struct();
        if ~isempty(wh.getDimension())
            [vals, names] = wh.getPhysIndicesByLinIndex(jobIndex);
            for pi = 1:numel(names)
                userParameters.(names{pi}) = vals{pi};
            end
        end
    end

    function validateNoRunIdSweepParameter(wh)
        if isfield(wh.inputParams, "run_id")
            error("submitJobs:RunIdParameterConflict", ...
                "Do not define userParameters.run_id. submitJobs manages run_id as a storage axis.");
        end
    end

    function storageWh = buildStorageWarehouse(sweepWh, runIds)
        if isscalar(runIds)
            storageWh = sweepWh;
            return
        end

        storageParameters = struct();
        storageParameters.run_id = runIds;

        sweepParameterNames = fieldnames(sweepWh.inputParams);
        for parameterIdx = 1:numel(sweepParameterNames)
            parameterName = sweepParameterNames{parameterIdx};
            storageParameters.(parameterName) = sweepWh.inputParams.(parameterName);
        end

        storageWh = DataStorage(storageParameters);
    end

    function storageJobIndex = buildStorageJobIndex(runId, userParameters, wh)
        storageParameterNames = wh.fn;

        if isempty(storageParameterNames)
            storageJobIndex = 1;
            return
        end

        storageSubscripts = cell(1, numel(storageParameterNames));
        for parameterIdx = 1:numel(storageParameterNames)
            parameterName = char(storageParameterNames(parameterIdx));

            if strcmp(parameterName, "run_id")
                parameterValue = runId;
            else
                parameterValue = userParameters.(parameterName);
            end

            storageSubscripts{parameterIdx} = wh.getIndexByPhys(parameterName, parameterValue);
        end

        if isscalar(storageSubscripts)
            storageJobIndex = storageSubscripts{1};
        else
            storageJobIndex = sub2ind(wh.getStorageSize(), storageSubscripts{:});
        end
    end

    function nameValue = structToNameValue(options)
        names = fieldnames(options);
        nameValue = cell(1, 2*numel(names));

        for i = 1:numel(names)
            nameValue{2*i - 1} = string(names{i});
            nameValue{2*i} = options.(names{i});
        end
    end

    function storeResult(val, job, ~)
        if isempty(wh) || ~isstruct(val)
            return
        end

        storageJobIndex = job.meta.storageJobIndex;
        resultFields = fieldnames(val);

        for resultIdx = 1:numel(resultFields)
            storageName = resultFields{resultIdx};

            if ~shouldStore(storageName)
                continue
            end

            if isempty(val.(storageName))
                continue
            end

            ensureStorage(storageName);
            wh.addValueToStorageByLinIdx(val.(storageName), storageName, storageJobIndex);
        end
    end

    function tf = shouldStore(storageName)
        if isempty(submit_options.storePackages)
            tf = true;
            return
        end

        tf = any(string(storageName) == submit_options.storePackages);
    end

    function ensureStorage(storageName)
        if ~isfield(wh.sto, storageName)
            wh.addStorage(storageName);
        end
    end

end