From ac10ba0fa1209c60b5ee7af097d3bf560a6e2ae7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Henning=20Schulze=20Ei=C3=9Fing?= <h_schu55@uni-muenster.de>
Date: Fri, 11 Oct 2024 12:21:04 +0200
Subject: [PATCH] Draft for processing override

---
 fuse/context.py                            | 14 +++++++++++++
 fuse/plugins/__init__.py                   |  3 +++
 fuse/plugins/processing/__init__.py        |  2 ++
 fuse/plugins/processing/corrected_areas.py | 23 ++++++++++++++++++++++
 4 files changed, 42 insertions(+)
 create mode 100644 fuse/plugins/processing/__init__.py
 create mode 100644 fuse/plugins/processing/corrected_areas.py

diff --git a/fuse/context.py b/fuse/context.py
index 765a8580..6d687197 100644
--- a/fuse/context.py
+++ b/fuse/context.py
@@ -87,6 +87,11 @@
     fuse.truth_information.ClusterTagging,
 ]
 
+# Plugins to override the default processing plugins in straxen
+processing_plugins = [
+    fuse.processing.CorrectedAreasMC,
+]
+
 
 def microphysics_context(
     output_folder="./fuse_data", simulation_config_file="fuse_config_nt_sr1_dev.json"
@@ -194,6 +199,12 @@ def full_chain_context(
     for plugin in truth_information_plugins:
         st.register(plugin)
 
+    # Register processing plugins
+    log.info("Overriding processing plugins:")
+    for plugin in processing_plugins:
+        log.info(f"Registering {plugin}")
+        st.register(plugin)
+
     if corrections_version is not None:
         st.apply_xedocs_configs(version=corrections_version)
 
@@ -221,6 +232,9 @@ def full_chain_context(
     # Deregister plugins with missing dependencies
     st.deregister_plugins_with_missing_dependencies()
 
+    # Purge unused configs
+    st.purge_unused_configs()
+
     return st
 
 
diff --git a/fuse/plugins/__init__.py b/fuse/plugins/__init__.py
index cc298898..bebf6792 100644
--- a/fuse/plugins/__init__.py
+++ b/fuse/plugins/__init__.py
@@ -9,3 +9,6 @@
 
 from . import truth_information
 from .truth_information import *
+
+from . import processing
+from .processing import *
diff --git a/fuse/plugins/processing/__init__.py b/fuse/plugins/processing/__init__.py
new file mode 100644
index 00000000..40e00983
--- /dev/null
+++ b/fuse/plugins/processing/__init__.py
@@ -0,0 +1,2 @@
+from . import corrected_areas
+from .corrected_areas import *
diff --git a/fuse/plugins/processing/corrected_areas.py b/fuse/plugins/processing/corrected_areas.py
new file mode 100644
index 00000000..4586f057
--- /dev/null
+++ b/fuse/plugins/processing/corrected_areas.py
@@ -0,0 +1,23 @@
+import strax
+import straxen
+
+export, __all__ = strax.exporter()
+
+
+@export
+class CorrectedAreasMC(straxen.CorrectedAreas):
+    """Corrected areas plugin for MC data.
+
+    This plugin overwrites the cs1 and cs2 fields with the not-time-
+    corrected values as the effects are not simulated in the MC.
+    """
+
+    __version__ = "0.0.1"
+    child_plugin = True
+
+    def compute(self, events):
+        result = super().compute(events)
+        result["cs1"] = result["cs1_wo_timecorr"]
+        result["cs2"] = result["cs2_wo_timecorr"]
+
+        return result