steps/utils/markers.js

import * as scran from "scran.js";

export const summaries2int = { "min": 0, "mean": 1, "min_rank": 4 };

export function unserializeGroupStats(handle, permuter, { no_summaries = false, compute_auc = true } = {}) {
    let output = {};
    for (const x of [ "means", "detected" ]) {
        output[x] = permuter(handle.open(x, { load: true }).values);
    }

    for (const i of [ "lfc", "delta_detected", "auc", "cohen" ]) {
        if (i == "auc" && !compute_auc) {
            continue;
        }

        if (no_summaries) {
            output[i] = handle.open(i, { load: true }).values;
        } else {
            let rhandle = handle.open(i);
            let current = {};
            for (const j of Object.keys(rhandle.children)) {
                current[j] = permuter(rhandle.open(j, { load: true }).values);
            }
            output[i] = current;
        }
    }

    return output;
}

export function fillGroupStats(object, i, vals) {
    object.means(i, { copy: false }).set(vals.means);
    object.detected(i, { copy: false }).set(vals.detected);

    for (const [s, v] of Object.entries(vals.cohen)) {
        object.cohen(i, { summary: summaries2int[s], copy: false }).set(v);
    }

    for (const [s, v] of Object.entries(vals.lfc)) {
        object.lfc(i, { summary: summaries2int[s], copy: false }).set(v);
    }

    for (const [s, v] of Object.entries(vals.delta_detected)) {
        object.deltaDetected(i, { summary: summaries2int[s], copy: false }).set(v);
    }

    if ("auc" in vals) {
        for (const [s, v] of Object.entries(vals.auc)) {
            object.auc(i, { summary: summaries2int[s], copy: false }).set(v);
        }
    }
}

/**
 * Report marker results for a given group or cluster, ordered so that the strongest candidate markers appear first.
 *
 * @param {ScoreMarkersResults} results - The marker results object generated by the `scoreMarkers` function in **scran.js**.
 * @param {number} group - Integer specifying the group or cluster of interest.
 * Any number can be used if it was part of the `groups` passed to `scoreMarkers`.
 * @param {string} rankEffect - Summarized effect size to use for ranking markers.
 * This should follow the format of `<effect>-<summary>` where `<effect>` may be `lfc`, `cohen`, `auc` or `delta_detected`,
 * and `<summary>` may be `min`, `mean` or `min-rank`.
 *
 * @return An object containing the marker statistics for the selection, sorted by the specified effect and summary size from `rankEffect`.
 * This contains:
 *   - `means`: a Float64Array of length equal to the number of genes, containing the mean expression within the selection.
 *   - `detected`: a Float64Array of length equal to the number of genes, containing the proportion of cells with detected expression inside the selection.
 *   - `lfc`: a Float64Array of length equal to the number of genes, containing the log-fold changes for the comparison between cells inside and outside the selection.
 *   - `delta_detected`: a Float64Array of length equal to the number of genes, containing the difference in the detected proportions between cells inside and outside the selection.
 */
export function formatMarkerResults(results, group, rankEffect) {
    if (!rankEffect || rankEffect === undefined) {
        rankEffect = "cohen-min-rank";
    }

    var ordering;
    {
        // Choosing the ranking statistic. Do NOT do any Wasm allocations
        // until 'ranking' is fully consumed!
        let ranking;
        let increasing = false;
      
        let index = 1;
        if (rankEffect.match(/-min$/)) {
            index = 0;
        } else if (rankEffect.match(/-min-rank$/)) {
            increasing = true;
            index = 4;
        }

        if (rankEffect.match(/^cohen-/)) {
            ranking = results.cohen(group, { summary: index, copy: false });
        } else if (rankEffect.match(/^auc-/)) {
            ranking = results.auc(group, { summary: index, copy: false });
        } else if (rankEffect.match(/^lfc-/)) {
            ranking = results.lfc(group, { summary: index, copy: false });
        } else if (rankEffect.match(/^delta-d-/)) {
            ranking = results.deltaDetected(group, { summary: index, copy: false });
        } else {
            throw "unknown rank type '" + rankEffect + "'";
        }
  
        // Computing the ordering based on the ranking statistic.
        ordering = new Int32Array(ranking.length);
        for (var i = 0; i < ordering.length; i++) {
            ordering[i] = i;
        }
        if (increasing) {
            ordering.sort((f, s) => (ranking[f] - ranking[s]));
        } else {
            ordering.sort((f, s) => (ranking[s] - ranking[f]));
        }
    }
  
    // Apply that ordering to each statistic of interest.
    var reorder = function(stats) {
        var thing = new Float64Array(stats.length);
        for (var i = 0; i < ordering.length; i++) {
            thing[i] = stats[ordering[i]];
        }
        return thing;
    };
  
    var stat_detected = reorder(results.detected(group, { copy: false }));
    var stat_mean = reorder(results.means(group, { copy: false }));
    var stat_lfc = reorder(results.lfc(group, { summary: 1, copy: false }));
    var stat_delta_d = reorder(results.deltaDetected(group, { summary: 1, copy: false }));

    return {
        "ordering": ordering,
        "means": stat_mean,
        "detected": stat_detected,
        "lfc": stat_lfc,
        "delta_detected": stat_delta_d
    };
}

export function locateVersusCache(left, right, cache) {
    let left_small = left < right;

    let bigg = (left_small ? right : left);
    if (!(bigg in cache)) {
        cache[bigg] = {};
    }
    let biggversus = cache[bigg];

    let smal = (left_small ? left : right); 
    let rerun = !(smal in biggversus);
    if (rerun) {
        biggversus[smal] = {};
    }

    return { 
        cached: biggversus[smal],
        run: rerun,
        left_small: left_small
    };
}

export function freeVersusResults(cache) {
    if (cache) {
        for (const v of Object.values(cache)) {
            for (const v2 of Object.values(v)) {
                for (const m of Object.values(v2)) {
                    scran.free(m);
                }
            }
        }
        for (const k of Object.keys(cache)) {
            delete cache[k];
        }
    }
}

export function computeVersusResults(matrices, clusters, block, keep, cache, lfc_threshold, compute_auc) {
    let new_block = null;
    if (block !== null) {
        new_block = scran.subsetBlock(block, keep);
        scran.dropUnusedBlock(new_block);
    }

    for (const modality of matrices.available()) {
        let modmat = matrices.get(modality);
        let sub;
        try {
            sub = scran.subsetColumns(modmat, keep);
            cache[modality] = scran.scoreMarkers(sub, clusters, { block: new_block, lfcThreshold: lfc_threshold, computeAuc: compute_auc });
        } finally {
            scran.free(sub);
        }
    }
}