Save more fields into isolated peaks

XENONnT · Apr 27, 2024 · 302a26d · 302a26d
1 parent d0fa9b7
commit 302a26d
Show file tree

Hide file tree

Showing 4 changed files with 8 additions and 18 deletions.
diff --git a/axidence/context.py b/axidence/context.py
@@ -123,6 +123,8 @@ def replication_tree(st, suffixes=["Paired", "Salted"], assign_attributes=None,
     provides and data_kind as the plugins in original tree.
     """
     if assign_attributes is None:
+        # this is due to some features are assigned in `infer_dtype` of the original plugins:
+        # https://github.com/XENONnT/straxen/blob/e555c7dcada2743d2ea627ea49df783e9dba40e3/straxen/plugins/events/event_basics.py#L69
         assign_attributes = {"EventBasics": ["peak_properties", "posrec_save"]}
 
     snakes = ["_" + strax.camel_to_snake(suffix) for suffix in suffixes]

diff --git a/axidence/plugins/isolated/isolated_s1.py b/axidence/plugins/isolated/isolated_s1.py
@@ -4,7 +4,7 @@
 import straxen
 
 from ...utils import needed_dtype, copy_dtype
-from ...dtypes import positioned_peak_dtype
+from ...dtypes import positioned_peak_dtype, correlation_fields
 
 
 class IsolatedS1(Plugin):
@@ -24,7 +24,7 @@ class IsolatedS1(Plugin):
     data_kind = "isolated_s1"
 
     isolated_peaks_fields = straxen.URLConfig(
-        default=np.dtype(positioned_peak_dtype()).names,
+        default=list(np.dtype(positioned_peak_dtype()).names) + correlation_fields,
         type=(list, tuple),
         help="Needed fields in isolated peaks",
     )

diff --git a/axidence/plugins/isolated/isolated_s2.py b/axidence/plugins/isolated/isolated_s2.py
@@ -4,7 +4,7 @@
 import straxen
 
 from ...utils import needed_dtype, copy_dtype
-from ...dtypes import positioned_peak_dtype
+from ...dtypes import positioned_peak_dtype, correlation_fields, event_level_fields
 
 
 class IsolatedS2(Plugin):
@@ -29,13 +29,13 @@ class IsolatedS2(Plugin):
     data_kind = "isolated_s2"
 
     isolated_peaks_fields = straxen.URLConfig(
-        default=np.dtype(positioned_peak_dtype()).names,
+        default=list(np.dtype(positioned_peak_dtype()).names) + correlation_fields,
         type=(list, tuple),
         help="Needed fields in isolated peaks",
     )
 
     isolated_events_fields = straxen.URLConfig(
-        default=[],
+        default=event_level_fields,
         type=(list, tuple),
         help="Needed fields in isolated events",
     )

diff --git a/axidence/plugins/pairing/peaks_paired.py b/axidence/plugins/pairing/peaks_paired.py
@@ -8,7 +8,7 @@
 from straxen import PeakProximity
 
 from ...utils import copy_dtype
-from ...dtypes import peak_positions_dtype, positioned_peak_dtype
+from ...dtypes import peak_positions_dtype
 from ...plugin import ExhaustPlugin, RunMetaPlugin
 
 
@@ -25,18 +25,6 @@ class PeaksPaired(ExhaustPlugin, RunMetaPlugin):
         help="Seed for pairing",
     )
 
-    isolated_peaks_fields = straxen.URLConfig(
-        default=np.dtype(positioned_peak_dtype()).names,
-        type=(list, tuple),
-        help="Needed fields in isolated peaks",
-    )
-
-    isolated_events_fields = straxen.URLConfig(
-        default=[],
-        type=(list, tuple),
-        help="Needed fields in isolated events",
-    )
-
     real_run_start = straxen.URLConfig(
         default=None,
         type=(int, None),