plotPR.jl

# =============================================================================
#  plotPR.jl — Evaluate GRN predictions against gold standards and plot PR curves
#
#  What:
#    Evaluates one or more inferred GRNs against gold-standard interaction sets
#    by computing precision-recall (PR) and ROC metrics, then generating
#    publication-quality PR curve plots. Supports multiple networks and
#    multiple gold standards in a single run. Optionally generates per-TF
#    PR curves and AUPR bar plots.
#
#  Required inputs:
#    outNetFiles    — inferred GRN file(s) as legend label → file path dict
#    gsParam        — gold-standard file(s) as name → file path dict
#    prTargGeneFile — target gene list used to restrict evaluation universe
#                     (set to "" to use all genes in the network)
#    gsRegsFile     — regulator list to restrict evaluation to shared TFs
#                     (set to "" to use all regulators)
#
#  Expected outputs (written relative to each network file's directory):
#    PR_noPotRegs/<gsName>/          — PR data files (if gsRegsFile = "")
#    PR_withPotRegs/<gsName>/        — PR data files (if gsRegsFile is set)
#    dirOutPlot/<figBaseName>_*.png  — PR curve plots and optional AUPR bar plots
#
#  In-memory vs. path storage (Section 1):
#    By default, Section 1 stores the full computePR result dict in prFilesByGS.
#    This avoids all disk reloads in Section 1B and plotting — most efficient for
#    a single session.
#    To store .jld paths instead (lower memory; reloads from disk downstream),
#    swap the one marked line in Section 1 — see the comment there.
#    All downstream functions (plotPRCurves, plotAUPR, extractMetricsAtLimit,
#    loadPRData) accept both transparently.
#
#  Usage:
#    julia examples/plotPR.jl
#    or configure the USER CONFIG section and run step-by-step in the REPL
#
#  Installation:
#    pkg> dev /path/to/InferelatorJL      # local development
#    pkg> add "https://github.com/org/InferelatorJL.jl"   # published release
#    Tip: load Revise before this file to pick up source edits without restarting:
#         using Revise; using InferelatorJL
# =============================================================================

using InferelatorJL
import InferelatorJL: computePR, plotPRCurves, plotAUPR, loadPRData,
                      extractMetricsAtLimit, saveSummaryTables

using OrderedCollections

# =========================================================
#  USER CONFIG — edit this section
# =========================================================

# Output directory for plots
dirOutPlot = "/path/to/plots"

# Base name for saved figures (set to "" to use gold-standard name only)
figBaseName = ""

# Network files to compare: legend label => file path
# Add one entry per network you want to compare
outNetFiles = OrderedDict(
    "TFA"      => "/path/to/output/networkLambda.../TFA/edges_subset.tsv",
    "TFmRNA"   => "/path/to/output/networkLambda.../TFmRNA/edges_subset.tsv",
    "Combined" => "/path/to/output/networkLambda.../Combined/combined_max.tsv",
)

# Gold-standard files: name => file path
# Add one entry per gold standard you want to evaluate against
gsParam = OrderedDict(
    "gsName" => "/path/to/goldStandard.tsv",
)

# Evaluation inputs
prTargGeneFile = "/path/to/target_genes.txt"   # set to "" to use all genes
gsRegsFile     = "/path/to/potential_regs.txt" # set to "" to use all TFs
breakTies      = true   # true (recommended): average indicators over tied scores → smoother, order-independent curve
                        # false: use raw binary indicators — curve depends on file sort order; use when comparing
                        #        with tools that do not break ties, or when all scores are unique
summaryLimits = []               # recall limits for post-hoc summary tables (Section 1B); [] skips tables
                                 # e.g. [0.01, 0.05, 0.1] — re-run Section 1B freely without rerunning computePR

# Plot parameters
lineTypes    = []   # e.g. ["-", "--", "-."] — one per dataset; [] uses defaults
lineWidths   = []   # e.g. [1.5, 1.0, 2.0] — one per dataset; [] uses defaults
lineColors   = []   # e.g. ["#e41a1c", "steelblue"] — matplotlib color strings; [] uses default palette
xLimitRecall = 0.1  # x-axis recall limit for plots; also used as the partial AUPR limit stored in .jld and shown in legends
yStepSize    = [0.1, 0.1]  # y-axis tick step size, one per gold standard; e.g. [0.1, 0.2]
yScaleType   = "linear"    # y-axis scale: "linear", "sqrt", or "cubert"
yZoomPR      = [[0.3, 0.9], [], []]  # one entry per gold standard:
                                    #   []         → full y-axis (0–1)
                                    #   [0.8]      → clip y-axis to 0–0.8
                                    #   [0.3, 0.9] → broken y-axis with gap between 0.3 and 0.9
# heightRatios = [0.5, 3.0]          # optional: override broken-axis panel heights [top, bottom]; auto-computed from yZoomPR by default
isInside     = false   # legend inside the plot
plotAUPRflag = false   # set to true to also generate AUPR bar plots
combinePlot  = true    # set to false to skip Section 2 (metrics only, no plots)
doPerTF      = true    # compute and save per-TF PR metrics to .jld
plotPerTF    = false   # plot individual per-TF curves inside computePR — expensive with many TFs; use tfList in Section 3 instead
tfList       = []      # list of TF names for per-TF curves; [] skips Section 3

# =========================================================
#  EXECUTION — no edits needed below this line
# =========================================================

mkpath(dirOutPlot)

# --- Helper: resolve per-GS plot parameters ---
function getPlotParams(i, gsName; figBaseName, yZoomPR, yStepSize)
    saveNamePR       = isempty(figBaseName) ? "$(gsName)" : "$(figBaseName)_$(gsName)"
    currentYzoomPR   = (length(yZoomPR) >= i && !isempty(yZoomPR[i])) ? yZoomPR[i] : Float64[]
    currentYstepSize = (length(yStepSize) >= i && !isempty(yStepSize[i])) ? yStepSize[i] : nothing
    return saveNamePR, currentYzoomPR, currentYstepSize
end

# ----- 1. Compute PR/ROC metrics and save results ------------------------------
# Run once — results are also written to .jld files automatically by computePR.
# To skip recomputation on re-runs, comment out this section and populate
# prFilesByGS manually with .jld paths:
#   prFilesByGS = OrderedDict("gsName" => OrderedDict("TFA" => "/path/to/.jld", ...))
@info "---- 1. Calculating Performance Metrics for the Networks -----"

parts = String[]
isempty(gsRegsFile)     || push!(parts, "regs")
isempty(prTargGeneFile) || push!(parts, "targs")
prSuffix = isempty(parts) ? "" : "_" * join(parts, "_")

# prFilesByGS holds one entry per (GS, network) pair.
# Each value is the full in-memory result dict from computePR — fastest for downstream use.
# To save memory and load from disk instead, swap the assignment line below (see comment).
prFilesByGS = OrderedDict{String, OrderedDict{String, Any}}()

for (legendLabel, outNetFile) in outNetFiles
    @info "Processing network" network=legendLabel file=outNetFile
    filepath = dirname(outNetFile)

    for (gsName, gsFile) in gsParam
        dirPR = joinpath(filepath, "PR" * prSuffix, gsName)
        mkpath(dirPR)
        @info "Using GS" gs=gsName saveDir=dirPR

        res = computePR(gsFile, outNetFile;
                        gsRegsFile       = gsRegsFile,
                        targGeneFile     = prTargGeneFile,
                        breakTies        = breakTies,
                        partialAUPRlimit = xLimitRecall,
                        doPerTF          = doPerTF,
                        plotPerTF        = plotPerTF,
                        saveDir          = dirPR)

        if res[:savedFile] === nothing
            @warn "computePR returned no saved file — skipping" network=legendLabel gs=gsName
            continue
        end

        !haskey(prFilesByGS, gsName) && (prFilesByGS[gsName] = OrderedDict{String, Any}())

        prFilesByGS[gsName][legendLabel] = res               # in-memory: no reloads downstream
        # prFilesByGS[gsName][legendLabel] = res[:savedFile]  # path: lower memory, loads from disk downstream
    end
end

# ----- 1B. Extract summary metrics at multiple recall limits ------------------
# Fast: uses in-memory results (or loads .jld once per entry if paths were stored).
# Re-run freely — adjust summaryLimits in USER CONFIG without rerunning Section 1.
if !isempty(summaryLimits)
    @info "---- 1B. Extracting Summary Metrics -----"

    netNames = collect(keys(outNetFiles))
    gsNames  = collect(keys(gsParam))

    auprFull        = OrderedDict{String, OrderedDict{String, Float64}}()
    auprAtLimit     = Dict(l => OrderedDict{String, OrderedDict{String, Float64}}() for l in summaryLimits)
    auprNormAtLimit = Dict(l => OrderedDict{String, OrderedDict{String, Float64}}() for l in summaryLimits)
    precAtLimit     = Dict(l => OrderedDict{String, OrderedDict{String, Float64}}() for l in summaryLimits)
    precNormAtLimit = Dict(l => OrderedDict{String, OrderedDict{String, Float64}}() for l in summaryLimits)

    for (gsName, netDict) in prFilesByGS
        auprFull[gsName] = OrderedDict{String, Float64}()
        for l in summaryLimits
            auprAtLimit[l][gsName]     = OrderedDict{String, Float64}()
            auprNormAtLimit[l][gsName] = OrderedDict{String, Float64}()
            precAtLimit[l][gsName]     = OrderedDict{String, Float64}()
            precNormAtLimit[l][gsName] = OrderedDict{String, Float64}()
        end

        for (legendLabel, source) in netDict
            # include 1.0 to get full AUPR in the same pass — single load regardless of source type
            allM = extractMetricsAtLimit(source, vcat(summaryLimits, [1.0]))
            auprFull[gsName][legendLabel] = allM[1.0].aupr

            for l in summaryLimits
                m = allM[l]
                auprAtLimit[l][gsName][legendLabel]     = m.aupr
                auprNormAtLimit[l][gsName][legendLabel] = m.auprNormalized
                precAtLimit[l][gsName][legendLabel]     = m.precAtLimit
                precNormAtLimit[l][gsName][legendLabel] = m.precNormalized
            end
        end
    end

    tables = [("aupr_full", auprFull)]
    for l in summaryLimits
        push!(tables, ("aupr_partial_$(l)",    auprAtLimit[l]))
        push!(tables, ("aupr_normalized_$(l)", auprNormAtLimit[l]))
        push!(tables, ("prec_at_limit_$(l)",   precAtLimit[l]))
        push!(tables, ("prec_normalized_$(l)", precNormAtLimit[l]))
    end

    saveSummaryTables(tables, gsNames, netNames, joinpath(dirOutPlot, "summary"))
end

# Example filter — uncomment and adjust to restrict which networks/GSs are plotted:
# keepKeys = ["TFA", "TFmRNA"]
# prFilesByGS = OrderedDict(
#     gs => OrderedDict(k => v for (k, v) in nets if k in keepKeys)
#     for (gs, nets) in prFilesByGS
# )

# ----- 2. Global PR curves --------------------------------------------
@info "---- 2. Generating Global PR Curves ----"
if combinePlot
    for (i, (gsName, listFilePR)) in enumerate(prFilesByGS)
        @info "Plotting PR curves" gs=gsName

        saveNamePR, currentYzoomPR, currentYstepSize = getPlotParams(i, gsName;
                                                                      figBaseName = figBaseName,
                                                                      yZoomPR     = yZoomPR,
                                                                      yStepSize   = yStepSize)

        plotPRCurves(listFilePR, dirOutPlot, saveNamePR;
                     xLimitRecall = xLimitRecall,
                     yZoomPR      = currentYzoomPR,
                     yStepSize    = currentYstepSize,
                     yScale       = yScaleType,
                     isInside     = isInside,
                     lineColors   = lineColors,
                     lineTypes    = lineTypes,
                     lineWidths   = lineWidths,
                     # heightRatios = heightRatios,   # uncomment + set in USER CONFIG to override auto ratio
                     mode         = :global)

        if plotAUPRflag
            singleGS     = OrderedDict(gsName => listFilePR)
            saveNameAUPR = isempty(figBaseName) ? "$(gsName)" : "$(figBaseName)_$(gsName)"

            for (figSize, saveLegend) in [((5, 4), true), ((1.5, 1.5), false)]
                plotAUPR(singleGS, dirOutPlot;
                         saveName       = saveNameAUPR,
                         metricType     = "partial",
                         figSize        = figSize,
                         axisTitleSize  = 9,
                         tickLabelSize  = 7,
                         legendFontSize = 9,
                         tickRotation   = 45,
                         plotType       = "bar",
                         saveLegend     = saveLegend)
            end
        end
        @info "Plots completed" gs=gsName
    end
end

# ----- 3. Per-TF PR curves --------------------------------------------
if !isempty(tfList) && doPerTF
    @info "----- 3. Generating Per-TF PR Curves -----"
    for (i, (gsName, resultsDict)) in enumerate(prFilesByGS)
        @info "Plotting per-TF PR curves" gs=gsName

        saveNamePR, currentYzoomPR, currentYstepSize = getPlotParams(i, gsName;
                                                                      figBaseName = figBaseName,
                                                                      yZoomPR     = yZoomPR,
                                                                      yStepSize   = yStepSize)
        saveNamePR = "perTF_$(saveNamePR)"
        tfListPR   = OrderedDict()

        resCache = Dict{String, Any}()
        for (runName, source) in resultsDict
            resCache[runName] = loadPRData(source; mode = :perTF)
        end

        for (runName, res) in resCache
            res === nothing && continue
            tfIndex = Dict(tf => j for (j, tf) in enumerate(res[:gsRegs]))
            for tf in tfList
                idx = get(tfIndex, tf, nothing)
                idx === nothing && continue

                label =
                    length(tfList) == 1 && length(resultsDict) > 1 ? runName :
                    length(resultsDict) == 1 && length(tfList) > 1 ? tf :
                    "$runName - $tf"

                tfListPR[label] = Dict(
                    :precisions => res[:precisions][idx],
                    :recalls    => res[:recalls][idx],
                    :randPR     => res[:randPR][idx]
                )
            end
        end

        if isempty(tfListPR)
            @warn "No matching TFs found in per-TF results — skipping plot" gs=gsName
            continue
        end

        plotPRCurves(tfListPR, dirOutPlot, saveNamePR;
                     xLimitRecall = xLimitRecall,
                     yZoomPR      = currentYzoomPR,
                     yStepSize    = currentYstepSize,
                     yScale       = yScaleType,
                     isInside     = isInside,
                     lineColors   = lineColors,
                     lineTypes    = lineTypes,
                     lineWidths   = lineWidths)
    end
end

@info "Completed — plots generated for all gold standards"