From 0ec7ee5ddc034bbdb665bb6de5c535e22f4e0c47 Mon Sep 17 00:00:00 2001
From: Sebastian Achim Mueller <sebastian-achim.mueller@mpi-hd.mpg.de>
Date: Wed, 21 Feb 2024 16:42:08 +0100
Subject: [PATCH] more central control over caches

---
 plenoirf/production/__init__.py               | 22 +++++-
 .../draw_primaries_and_pointings.py           | 67 +++++++------------
 ...wer_and_collect_cherenkov_light_in_grid.py | 19 +-----
 3 files changed, 46 insertions(+), 62 deletions(-)

diff --git a/plenoirf/production/__init__.py b/plenoirf/production/__init__.py
index 6a4548e..7bed619 100644
--- a/plenoirf/production/__init__.py
+++ b/plenoirf/production/__init__.py
@@ -81,15 +81,31 @@ def run_job_in_dir(job, work_dir):
         job = draw_pointing_range.run_job(job=job, logger=logger)
 
     with jll.TimeDelta(logger, "draw_primaries_and_pointings"):
-        job = draw_primaries_and_pointings.run_job(job=job, logger=logger)
+        job = checkpoint.checkpoint(
+            job=job,
+            logger=logger,
+            func=draw_primaries_and_pointings.run_job,
+            cache_path=opj(
+                job["paths"]["work_dir"],
+                "draw_primaries_and_pointings",
+                "__job_cache__",
+            ),
+        )
 
     job["event_table"] = event_table.structure.init_table_dynamicsizerecarray()
 
     with jll.TimeDelta(
         logger, "simulate_shower_and_collect_cherenkov_light_in_grid"
     ):
-        job = simulate_shower_and_collect_cherenkov_light_in_grid.run_job(
-            job=job, logger=logger
+        job = checkpoint.checkpoint(
+            job=job,
+            logger=logger,
+            func=simulate_shower_and_collect_cherenkov_light_in_grid.run_job,
+            cache_path=opj(
+                job["paths"]["work_dir"],
+                "simulate_shower_and_collect_cherenkov_light_in_grid",
+                "__job_cache__",
+            ),
         )
 
     with jll.TimeDelta(logger, "inspect_cherenkov_pool"):
diff --git a/plenoirf/production/draw_primaries_and_pointings.py b/plenoirf/production/draw_primaries_and_pointings.py
index dee676f..5f7ee43 100644
--- a/plenoirf/production/draw_primaries_and_pointings.py
+++ b/plenoirf/production/draw_primaries_and_pointings.py
@@ -14,7 +14,6 @@
 from .. import bookkeeping
 from .. import tar_append
 from .. import utils
-from . import job_io
 
 
 def draw_primaries_and_pointings(
@@ -187,51 +186,37 @@ def draw_primaries_and_pointings(
 
 
 def run_job(job, logger):
-    cache_path = os.path.join(
-        job["paths"]["work_dir"],
-        "draw_primaries_and_pointings",
-        "__job_cache__",
+    logger.info("draw_primaries_and_pointings, open AllSky")
+    allsky = magnetic_deflection.allsky.AllSky(
+        job["paths"]["magnetic_deflection_allsky"]
     )
 
-    if os.path.exists(cache_path) and job["cache"]:
-        logger.info("draw_primaries_and_pointings, read cache")
-        return job_io.read(path=cache_path)
-    else:
-        logger.info("draw_primaries_and_pointings, open AllSky")
-        allsky = magnetic_deflection.allsky.AllSky(
-            job["paths"]["magnetic_deflection_allsky"]
-        )
-
-        logger.info("draw_primaries_and_pointings, draw primaries")
-
-        drw, debug = draw_primaries_and_pointings(
-            prng=job["prng"],
-            run_id=job["run_id"],
-            site_particle_magnetic_deflection=allsky,
-            pointing_range=job["run"]["pointing_range"],
-            field_of_view_half_angle_rad=job["instrument"][
-                "field_of_view_half_angle_rad"
-            ],
-            num_events=job["num_events"],
-            allsky_query_mode=job["config"]["magnetic_deflection"][
-                "query_mode"
-            ],
-            event_uids_for_debugging=job["run"]["event_uids_for_debugging"],
-            logger=logger,
-        )
-
-        job["run"].update(drw)
+    logger.info("draw_primaries_and_pointings, draw primaries")
+
+    drw, debug = draw_primaries_and_pointings(
+        prng=job["prng"],
+        run_id=job["run_id"],
+        site_particle_magnetic_deflection=allsky,
+        pointing_range=job["run"]["pointing_range"],
+        field_of_view_half_angle_rad=job["instrument"][
+            "field_of_view_half_angle_rad"
+        ],
+        num_events=job["num_events"],
+        allsky_query_mode=job["config"]["magnetic_deflection"][
+            "query_mode"
+        ],
+        event_uids_for_debugging=job["run"]["event_uids_for_debugging"],
+        logger=logger,
+    )
 
-        logger.info("draw_primaries_and_pointings, export debug info")
+    job["run"].update(drw)
 
-        write_draw_primaries_and_pointings_debug(
-            path=job["paths"]["debug"]["draw_primary_and_pointing"],
-            debug=debug,
-        )
+    logger.info("draw_primaries_and_pointings, export debug info")
 
-        if job["cache"]:
-            logger.info("draw_primaries_and_pointings, write cache")
-            job_io.write(path=cache_path, job=job)
+    write_draw_primaries_and_pointings_debug(
+        path=job["paths"]["debug"]["draw_primary_and_pointing"],
+        debug=debug,
+    )
 
     return job
 
diff --git a/plenoirf/production/simulate_shower_and_collect_cherenkov_light_in_grid.py b/plenoirf/production/simulate_shower_and_collect_cherenkov_light_in_grid.py
index c490df5..3436842 100644
--- a/plenoirf/production/simulate_shower_and_collect_cherenkov_light_in_grid.py
+++ b/plenoirf/production/simulate_shower_and_collect_cherenkov_light_in_grid.py
@@ -17,28 +17,11 @@
 from .. import utils
 
 from . import transform_cherenkov_bunches
-from . import job_io
 from . import cherenkov_bunch_storage
 
 
 def run_job(job, logger):
-    cache_path = os.path.join(
-        job["paths"]["work_dir"],
-        "simulate_shower_and_collect_cherenkov_light_in_grid",
-        "__job_cache__",
-    )
-
-    if os.path.exists(cache_path) and job["cache"]:
-        logger.info("corsika and grid, read cache")
-        return job_io.read(path=cache_path)
-    else:
-        logger.info("corsika and grid, run corsika")
-        job = corsika_and_grid(job=job, logger=logger)
-
-        if job["cache"]:
-            logger.info("corsika and grid, write cache")
-            job_io.write(path=cache_path, job=job)
-
+    job = corsika_and_grid(job=job, logger=logger)
     return job