ecal_gaps: enable caching for simulation/reconstruction

eic · Nov 24, 2024 · a6a770b · a6a770b
1 parent 417e954
commit a6a770b
Show file tree

Hide file tree

Showing 3 changed files with 46 additions and 32 deletions.
diff --git a/Snakefile b/Snakefile
@@ -1,5 +1,35 @@
 configfile: "snakemake.yml"
 
+import functools
+import os
+from snakemake.logging import logger
+
+
+@functools.cache
+def get_spack_package_hash(package_name):
+    import json
+    try:
+        ver_info = json.loads(subprocess.check_output(["spack", "find", "--json", package_name]))
+        return ver_info[0]["package_hash"]
+    except FileNotFoundError as e:
+        logger.warning("Spack is not installed")
+        return ""
+    except subprocess.CalledProcessError as e:
+        print(e)
+        return ""
+
+
+@functools.cache
+def find_epic_libraries():
+    import ctypes.util
+    # if library is not found (not avaliable) we return an empty list to let DAG still evaluate
+    libs = []
+    lib = ctypes.util.find_library("epic")
+    if lib is not None:
+        libs.append(os.environ["DETECTOR_PATH"] + "/../../lib/" + lib)
+    return libs
+
+
 include: "benchmarks/backgrounds/Snakefile"
 include: "benchmarks/backwards_ecal/Snakefile"
 include: "benchmarks/barrel_ecal/Snakefile"

diff --git a/benchmarks/backwards_ecal/Snakefile b/benchmarks/backwards_ecal/Snakefile
@@ -4,33 +4,6 @@ def get_n_events(wildcards):
     n_events = int(n_events // (energy ** 0.5))
     return n_events
 
-import functools
-import json
-import ctypes.util
-import warnings
-from snakemake.logging import logger
-
-@functools.cache
-def get_spack_package_hash(package_name):
-    try:
-        ver_info = json.loads(subprocess.check_output(["spack", "find", "--json", package_name]))
-        return ver_info[0]["package_hash"]
-    except FileNotFoundError as e:
-        logger.warning("Spack is not installed")
-        return ""
-    except subprocess.CalledProcessError as e:
-        print(e)
-        return ""
-
-@functools.cache
-def find_epic_libraries():
-    # if library is not found (not avaliable) we return an empty list to let DAG still evaluate
-    libs = []
-    lib = ctypes.util.find_library("epic")
-    if lib is not None:
-        libs.append(os.environ["DETECTOR_PATH"] + "/../../lib/" + lib)
-    return libs
-
 
 rule backwards_ecal_sim:
     input:

diff --git a/benchmarks/ecal_gaps/Snakefile b/benchmarks/ecal_gaps/Snakefile
@@ -3,8 +3,9 @@ import os
 
 rule ecal_gaps_sim:
     input:
-        steering_file=ancient("EPIC/EVGEN/SINGLE/{PARTICLE}/{ENERGY}/{PHASE_SPACE}/{PARTICLE}_{ENERGY}_{PHASE_SPACE}.steer"),
+        steering_file="EPIC/EVGEN/SINGLE/{PARTICLE}/{ENERGY}/{PHASE_SPACE}/{PARTICLE}_{ENERGY}_{PHASE_SPACE}.steer",
         warmup="warmup/{DETECTOR_CONFIG}.edm4hep.root",
+        geometry_lib=find_epic_libraries(),
     output:
         "sim_output/ecal_gaps/{DETECTOR_CONFIG}/{PARTICLE}/{ENERGY}/{PHASE_SPACE}/{PARTICLE}_{ENERGY}_{PHASE_SPACE}.{INDEX}.edm4hep.root",
     log:
@@ -15,19 +16,25 @@ rule ecal_gaps_sim:
         PHASE_SPACE="(3to50|45to135|130to177)deg",
         INDEX="\d{4}",
     params:
-        N_EVENTS=1000
+        N_EVENTS=1000,
+        SEED=lambda wildcards: "1" + wildcards.INDEX,
+        DETECTOR_PATH=os.environ["DETECTOR_PATH"],
+        DETECTOR_CONFIG=lambda wildcards: wildcards.DETECTOR_CONFIG,
+        DD4HEP_HASH=get_spack_package_hash("dd4hep"),
+        NPSIM_HASH=get_spack_package_hash("npsim"),
+    cache: True
     shell:
         """
 set -m # monitor mode to prevent lingering processes
 exec ddsim \
   --runType batch \
   --enableGun \
   --steeringFile "{input.steering_file}" \
-  --random.seed 1{wildcards.INDEX} \
+  --random.seed {params.SEED} \
   --filter.tracker edep0 \
   -v WARNING \
   --numberOfEvents {params.N_EVENTS} \
-  --compactFile $DETECTOR_PATH/{wildcards.DETECTOR_CONFIG}.xml \
+  --compactFile {params.DETECTOR_PATH}/{params.DETECTOR_CONFIG}.xml \
   --outputFile {output}
 """
 
@@ -41,9 +48,13 @@ rule ecal_gaps_recon:
         "sim_output/ecal_gaps/{DETECTOR_CONFIG}/{PARTICLE}/{ENERGY}/{PHASE_SPACE}/{PARTICLE}_{ENERGY}_{PHASE_SPACE}.{INDEX}.eicrecon.tree.edm4eic.root.log",
     wildcard_constraints:
         INDEX="\d{4}",
+    params:
+        DETECTOR_CONFIG=lambda wildcards: wildcards.DETECTOR_CONFIG,
+        EICRECON_HASH=get_spack_package_hash("eicrecon"),
+    cache: True
     shell: """
 set -m # monitor mode to prevent lingering processes
-exec env DETECTOR_CONFIG={wildcards.DETECTOR_CONFIG} \
+exec env DETECTOR_CONFIG={params.DETECTOR_CONFIG} \
   eicrecon {input} -Ppodio:output_file={output} \
   -Ppodio:output_collections=EcalEndcapNRecHits,EcalBarrelScFiRecHits,EcalBarrelImagingRecHits,EcalEndcapPRecHits,MCParticles
 """