Snakefile

import os

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 diffractive_vm_compile:
    input:
        ROOT_BUILD_DIR_PREFIX + "benchmarks/diffractive_vm/analysis/diffractive_vm_cxx.so",
        ROOT_BUILD_DIR_PREFIX + "benchmarks/diffractive_vm/analysis/plot_cxx.so",


rule diffractive_vm_get:
    input:
        lambda wildcards: S3.remote(f"eictest/EPIC/EVGEN/EXCLUSIVE/DIFFRACTIVE_{wildcards.PARTICLE.upper()}_ABCONV/Sartre/Coherent/sartre_bnonsat_Au_{wildcards.PARTICLE}_ab_eAu_1_{wildcards.INDEX}.hepmc.gz"),
    output:
        "input/diffractive_vm/sartre_{PARTICLE}_{INDEX}.hepmc",
    shell:
        """
gunzip -c {input} > {output}
"""


rule diffractive_vm_campaign_reco_get:
    input:
        lambda wildcards: S3.remote(f"eictest/EPIC/RECO/22.11.3/epic_arches/EXCLUSIVE/DIFFRACTIVE_{wildcards.PARTICLE.upper()}_ABCONV/Sartre/Coherent/sartre_bnonsat_Au_{wildcards.PARTICLE}_ab_eAu_1.0{wildcards.INDEX}.eicrecon.tree.edm4eic.root")
    output:
        "reco/epic_arches/campaign_22.11.3_sartre_{PARTICLE}_{INDEX}_eicrecon.edm4eic.root",
    shell:
        """
ln {input} {output}
"""


rule diffractive_vm_sim:
    input:
        "input/diffractive_vm/sartre_{PARTICLE}_{INDEX}.hepmc",
    output:
        "sim/{DETECTOR_CONFIG}/sartre_{PARTICLE}_{INDEX}.edm4hep.root",
    params:
        N_EVENTS=100
    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} \
  --outputFile {output}
"""


rule diffractive_vm_reco_eicrecon:
    input:
        "sim/{DETECTOR_CONFIG}/{file}.edm4hep.root",
    output:
        "reco/{DETECTOR_CONFIG}/{file}_eicrecon.edm4eic.root",
    shell:
        """
DETECTOR_CONFIG={wildcards.DETECTOR_CONFIG} eicrecon {input} -Ppodio:output_file={output}
"""


rule diffractive_vm_analysis:
    input:
        script="benchmarks/diffractive_vm/analysis/diffractive_vm.cxx",
        script_compiled=ROOT_BUILD_DIR_PREFIX + "benchmarks/diffractive_vm/analysis/diffractive_vm_cxx.so",
        data="reco/{DETECTOR_CONFIG}/{SIM}sartre_{PARTICLE}{SUFFIX}.edm4eic.root",
    params:
        EBEAM=18,
        PBEAM=110,
        BEAM_TAG="18on110",
    wildcard_constraints:
        SIM=".*", # can be empty
        PARTICLE="[^_]*",
        SUFFIX=".*", # can be empty
    output:
        config="results/{DETECTOR_CONFIG}/diffractive_vm/{SIM}sartre_{PARTICLE}{SUFFIX}/config.json",
        plots="results/{DETECTOR_CONFIG}/diffractive_vm/{SIM}sartre_{PARTICLE}{SUFFIX}/plots.root",
    shell:
        """
cat > {output.config} <<EOF
{{
  "rec_file": "{input.data}",
  "vm_name": "{wildcards.PARTICLE}",
  "detector": "{wildcards.DETECTOR_CONFIG}",
  "ebeam": {params.EBEAM},
  "pbeam": {params.PBEAM},
  "output_prefix": "$(dirname "{output.plots}")/plots",
  "test_tag": "{wildcards.PARTICLE}_{params.BEAM_TAG}"
}}
EOF
mkdir -p $(dirname "{output.plots}")
root -l -b -q '{input.script}+("{output.config}")'
"""


rule diffractive_vm_combine_sartre:
    input:
        lambda wildcards: [f"results/{wildcards.DETECTOR_CONFIG}/diffractive_vm/{wildcards.PREFIX}sartre_{wildcards.PARTICLE}_{ix:03d}{wildcards.SUFFIX}/plots.root" for ix in range(int(wildcards.NUM_FILES))],
    params:
        EBEAM=18,
        PBEAM=110,
        BEAM_TAG="18on110",
    wildcard_constraints:
        PREFIX=".*", # can be empty
        PARTICLE="[^_]*",
        NUM_FILES="\d+",
        SUFFIX=".*", # can be empty
    output:
        config="results/{DETECTOR_CONFIG}/diffractive_vm/{PREFIX}sartre_{PARTICLE}_combined_{NUM_FILES}{SUFFIX}/config.json",
        plots="results/{DETECTOR_CONFIG}/diffractive_vm/{PREFIX}sartre_{PARTICLE}_combined_{NUM_FILES}{SUFFIX}/plots.root",
    shell:
        """
cat > {output.config} <<EOF
{{
  "rec_file": "",
  "vm_name": "{wildcards.PARTICLE}",
  "detector": "{wildcards.DETECTOR_CONFIG}",
  "ebeam": {params.EBEAM},
  "pbeam": {params.PBEAM},
  "output_prefix": "$(dirname "{output.plots}")/plots",
  "test_tag": "{wildcards.PARTICLE}_{params.BEAM_TAG}"
}}
EOF
hadd {output.plots} {input}
"""

ruleorder: diffractive_vm_combine_sartre > diffractive_vm_analysis


rule diffractive_vm_plots:
    input:
        script="benchmarks/diffractive_vm/analysis/plot.cxx",
        script_compiled=ROOT_BUILD_DIR_PREFIX + "benchmarks/diffractive_vm/analysis/plot_cxx.so",
        config="{RESULT_PREFIX}/config.json",
    output:
        "{RESULT_PREFIX}/plots_benchmark-phi-dsigmadt.pdf",
        "{RESULT_PREFIX}/plots_benchmark-phi-t-resolution.pdf",
        "{RESULT_PREFIX}/plots_benchmark-phi-DIS-kinematics.pdf",
    shell:
        """
root -l -b -q '{input.script}+("{input.config}")'
"""


# Couple examples of invocation:

rule diffractive_vm_run_over_a_campaign:
    input:
        "results/epic_arches/diffractive_vm/campaign_22.11.3_sartre_phi_combined_9/plots_benchmark-phi-dsigmadt.pdf",
    message:
        "See output in {input[0]}"


rule diffractive_vm_run_locally:
    input:
        "results/" + os.environ["DETECTOR_CONFIG"] + "/diffractive_vm/sartre_phi_combined_9_eicrecon/plots_benchmark-phi-dsigmadt.pdf",
    message:
        "See output in {input[0]}"