Snakefile

import os
import shutil

from snakemake.remote.S3 import RemoteProvider as S3RemoteProvider


S3 = S3RemoteProvider(
    endpoint_url="https://eics3.sdcc.bnl.gov:9000",
    access_key_id=os.environ["S3_ACCESS_KEY"],
    secret_access_key=os.environ["S3_SECRET_KEY"],
)


rule demp_compile:
    input:
        ROOT_BUILD_DIR_PREFIX + "benchmarks/demp/analysis/demp_analysis_cxx.so",
        ROOT_BUILD_DIR_PREFIX + "benchmarks/demp/analysis/demp_plots_cxx.so",

rule demp_get:
    input:
        S3.remote("eictest/EPIC/EVGEN/EXCLUSIVE/DEMP/{EBEAM}on{PBEAM}/eic_DEMPGen_{EBEAM}on{PBEAM}_ip6_pi+_1B_{INDEX}.hepmc"),
    output:
        "input/demp/eic_DEMPGen_{EBEAM}on{PBEAM}_ip6_pi+_1B_{INDEX}.hepmc",
    run:
        shutil.move(input[0], output[0])


rule demp_sim:
    input:
        hepmc="input/demp/eic_DEMPGen_{EBEAM}on{PBEAM}_ip6_pi+_1B_{INDEX}.hepmc",
        warmup="warmup/{DETECTOR_CONFIG}.edm4hep.root",
    output:
        "sim/{DETECTOR_CONFIG}/demp_{EBEAM}on{PBEAM}_{INDEX}.edm4hep.root",
    params:
        N_EVENTS=100
    wildcard_constraints:
        EBEAM="\d+",
        PBEAM="\d+",
        INDEX="\d+",
    shell:
        """
ddsim \
  --runType batch \
  --part.minimalKineticEnergy 100*GeV  \
  --filter.tracker edep0 \
  -v WARNING \
  --numberOfEvents {params.N_EVENTS} \
  --compactFile $DETECTOR_PATH/{wildcards.DETECTOR_CONFIG}.xml \
  --inputFiles {input.hepmc} \
  --outputFile {output}
"""


rule demp_reco:
    input:
        "sim/{DETECTOR_CONFIG}/demp_{EBEAM}on{PBEAM}_{INDEX}.edm4hep.root",
    output:
        "reco/{DETECTOR_CONFIG}/demp_{EBEAM}on{PBEAM}_{INDEX}.edm4eic.root",
    wildcard_constraints:
        EBEAM="\d+",
        PBEAM="\d+",
        INDEX="\d+",
    shell:
        """
DETECTOR_CONFIG={wildcards.DETECTOR_CONFIG} eicrecon {input} -Ppodio:output_file={output}
"""


rule demp_analysis:
    input:
        script="benchmarks/demp/analysis/demp_analysis.cxx",
        script_compiled=ROOT_BUILD_DIR_PREFIX + "benchmarks/demp/analysis/demp_analysis_cxx.so",
        data="reco/{DETECTOR_CONFIG}/demp_{EBEAM}on{PBEAM}_{INDEX}.edm4eic.root",
    output:
        config="results/{DETECTOR_CONFIG}/demp/demp_{EBEAM}on{PBEAM}_{INDEX}/config.json",
        hists="results/{DETECTOR_CONFIG}/demp/demp_{EBEAM}on{PBEAM}_{INDEX}/hists.root",
    wildcard_constraints:
        EBEAM="\d+",
        PBEAM="\d+",
        INDEX="\d+",
    shell:
        """
cat > {output.config} <<EOF
{{
  "rec_file": "{input.data}",
  "detector": "{wildcards.DETECTOR_CONFIG}",
  "ebeam": {wildcards.EBEAM},
  "pbeam": {wildcards.PBEAM},
  "output_prefix": "$(dirname "{output.hists}")/hists"
}}
EOF
root -l -b -q '{input.script}+("{output.config}")'
"""


rule demp_combine:
    input:
        lambda wildcards: [f"results/{wildcards.DETECTOR_CONFIG}/demp/demp_{wildcards.EBEAM}on{wildcards.PBEAM}_{ix}/hists.root" for ix in range(1,int(wildcards.NUM_FILES)+1)],
    output:
        config="results/{DETECTOR_CONFIG}/demp/demp_{EBEAM}on{PBEAM}_combined_{NUM_FILES}/config.json",
        hists="results/{DETECTOR_CONFIG}/demp/demp_{EBEAM}on{PBEAM}_combined_{NUM_FILES}/hists.root",
    wildcard_constraints:
        EBEAM="\d+",
        PBEAM="\d+",
        NUM_FILES="\d+",
    shell:
        """
cat > {output.config} <<EOF
{{
  "hists_file": "{output.hists}",
  "detector": "{wildcards.DETECTOR_CONFIG}",
  "ebeam": {wildcards.EBEAM},
  "pbeam": {wildcards.PBEAM},
  "output_prefix": "$(dirname "{output.hists}")/plots"
}}
EOF
hadd {output.hists} {input}
"""


rule demp_plots:
    input:
        script="benchmarks/demp/analysis/demp_plots.cxx",
        script_compiled=ROOT_BUILD_DIR_PREFIX + "benchmarks/demp/analysis/demp_plots_cxx.so",
        config="results/{DETECTOR_CONFIG}/demp/demp_{EBEAM}on{PBEAM}_combined_{NUM_FILES}/config.json",
    output:
        "results/{DETECTOR_CONFIG}/demp/demp_{EBEAM}on{PBEAM}_combined_{NUM_FILES}/plots.pdf"
    wildcard_constraints:
        EBEAM="\d+",
        PBEAM="\d+",
        NUM_FILES="\d+",
    shell:
        """
root -l -b -q '{input.script}+("{input.config}")'
"""

# Example of invocation
rule demp_run_locally:
    input:
        "results/" + os.environ["DETECTOR_CONFIG"] + "/demp/demp_5on41_combined_5/plots.pdf",
    message:
        "See output in {input[0]}"