From a4f96097a6f855ebd4eb93b966d3af0e27543f13 Mon Sep 17 00:00:00 2001
From: josh0-jrg <joshua.robert.green@gmail.com>
Date: Tue, 24 Sep 2024 04:10:12 -0700
Subject: [PATCH] Added in temp drift map

---
 flamedisx/nest/lxe_blocks/energy_spectrum.py | 29 ++++++++++++++++----
 flamedisx/nest/lxe_sources.py                | 10 +++++--
 2 files changed, 32 insertions(+), 7 deletions(-)

diff --git a/flamedisx/nest/lxe_blocks/energy_spectrum.py b/flamedisx/nest/lxe_blocks/energy_spectrum.py
index ad935aba2..0519498e9 100644
--- a/flamedisx/nest/lxe_blocks/energy_spectrum.py
+++ b/flamedisx/nest/lxe_blocks/energy_spectrum.py
@@ -10,7 +10,23 @@
 export, __all__ = fd.exporter()
 o = tf.newaxis
 
-
+def apply_drift_map(data,drift_map,s2x_map,z_topDrift,drift_velocity):
+    #Replace with drift-time map
+    if drift_map is not None:
+        print("using drift-map")
+        data['drift_time']=drift_map(np.array([data['r'],data['z']]).T)
+    elif s2x_map is not None:
+        #leave this in for now
+        print("using S2x map")
+        costheta=data['x']/data['r']
+        sinetheta=data['y']/data['r']
+        data['s2x']=drift_map(np.array([data['r'],data['z']]).T)
+        data['x']=data['s2x']*costheta
+        data['y']=data['s2x']*sinetheta
+    else: 
+        print("using linear relation")
+        data['drift_time'] = (z_topDrift-data['z']) / drift_velocity
+    return data
 @export
 class EnergySpectrum(fd.FirstBlock):
     dimensions = ('energy',)
@@ -18,7 +34,7 @@ class EnergySpectrum(fd.FirstBlock):
         'energies',
         'radius', 'z_top', 'z_bottom', 'z_topDrift',
         'drift_velocity',
-        't_start', 't_stop')
+        't_start', 't_stop','drift_map')
 
     # The default boundaries are at points where the WIMP wind is at its
     # average speed.
@@ -32,6 +48,8 @@ class EnergySpectrum(fd.FirstBlock):
     # Just a dummy 0-10 keV spectrum
     energies = tf.cast(tf.linspace(0., 10., 1000),
                        dtype=fd.float_type())
+    #drift map variabls
+    drift_map=None
 
     def domain(self, data_tensor):
         assert isinstance(self.energies, tf.Tensor)  # see WIMPsource for why
@@ -114,7 +132,8 @@ def draw_positions(self, n_events, **params):
                                       size=n_events)
         data['x'], data['y'] = fd.pol_to_cart(data['r'], data['theta'])
 
-        data['drift_time'] = (self.z_topDrift-data['z']) / self.drift_velocity
+        #Replace with drift-time map
+        data=apply_drift_map(data,self.drift_map,None,self.z_topDrift, self.drift_velocity)
         return data
 
     def draw_time(self, n_events, **params):
@@ -190,7 +209,7 @@ def validate_fix_truth(self, d):
             else:
                 raise ValueError("When fixing position, give (x, y, z), "
                                  "or (r, theta, z).")
-            d['drift_time'] = (self.z_topDrift-d['z']) / self.drift_velocity
+            data=apply_drift_map(data,self.drift_map,None,self.z_topDrift, self.drift_velocity)
         elif 'event_time' not in d and 'energy' not in d:
             # Neither position, time, nor energy given
             raise ValueError(f"Dict should contain at least ['x', 'y', 'z'] "
@@ -348,7 +367,7 @@ def draw_positions(self, n_events, **params):
         else:
             data['r'], data['theta'] = fd.cart_to_pol(data['x'], data['y'])
 
-        data['drift_time'] = (self.z_topDrift-data['z']) / self.drift_velocity
+        data=apply_drift_map(data,self.drift_map,None,self.z_topDrift, self.drift_velocity)
         return data
 
 
diff --git a/flamedisx/nest/lxe_sources.py b/flamedisx/nest/lxe_sources.py
index 30c494c58..4e664902c 100644
--- a/flamedisx/nest/lxe_sources.py
+++ b/flamedisx/nest/lxe_sources.py
@@ -7,7 +7,8 @@
 
 import flamedisx as fd
 from .. import nest as fd_nest
-
+import numpy as np
+from scipy import interpolate
 import math as m
 pi = tf.constant(m.pi)
 
@@ -21,8 +22,13 @@
 
 
 class nestSource(fd.BlockModelSource):
-    def __init__(self, *args, detector='default', **kwargs):
+    def __init__(self, *args, detector='default',drift_map_path=None, **kwargs):
         assert detector in ('default', 'lz','lz_SR3')
+        
+        if drift_map_path is not None:
+            print("Loading in Drift map:\n",drift_map_path)
+            drift_map=np.loadtxt(drift_map_path,delimiter=',').T
+            self.drift_map=interpolate.LinearNDInterpolator(drift_map[:2].T/10,drift_map[3]*1e3)
 
         self.detector = detector