Add wimprates rate generator to nestWIMPSource

flamedisx/nest/__init__.py

@@ -1,4 +1,5 @@
 from .parameter_calc import *
+from .get_energy_hist import *
 
 from .lxe_blocks.energy_spectrum import *
 from .lxe_blocks.secondary_quanta_generation import *

flamedisx/nest/config/xlzd.ini

@@ -1,26 +1,43 @@
 [NEST]
 
-; Detector conditions
-; DUMMY VALUES (defined in constructor instead)
-drift_field_config = 300.
-gas_field_config = 0.
-elife_config = 0.
-g1_config = 0.
-temperature_config = 174.1
-pressure_config = 1.79
-num_pmts_config = 902
-double_pe_fraction_config = 0.
-g1_gas_config = 0.
-s2Fano_config = 0.
+
+; Detector parameters
+temperature_config = 174.1 ; K
+pressure_config = 1.79 ; bar
+drift_field_config = 80. ; DUMMY VALUE, DEFINED IN CONSTRUCTOR INSTEAD
+gas_field_config = 6. ; DUMMY VALUE, DEFINED IN CONSTRUCTOR INSTEAD
+
+g1_config = 0.31 ; ; DUMMY VALUE, DEFINED IN CONSTRUCTOR INSTEAD
+elife_config = 13000000. ; DUMMY VALUE, DEFINED IN CONSTRUCTOR INSTEAD
+
 spe_res_config = 0.38
-spe_thr_config = 0.
+spe_thr_config = 0.375
 spe_eff_config = 1.
-; REAL VALUES
+
+num_pmts_config = 1155
+double_pe_fraction_config = 0.2
+
+coin_level_config = 4
+
+g1_gas_config = 0.10
+s2Fano_config = 2.
+
+s2_thr_config = 198.
+
 S1_noise_config = 0.
 S2_noise_config = 0.
 
-; Geometry parameters
-; DUMMY VALUES (either redundant or overriden)
+
+; Selection parameters
+
+cS1_min_config = 3.
+cS1_max_config = 30.
+log10_cS2_min_config = 3.
+log10_cS2_max_config = 5.0
+
+
+; Dummy values - either redundant or overriden
+
 radius_config = 0.
 z_topDrift_config = 0.
 dt_min_config = 0.
@@ -34,8 +51,6 @@ S1_min_config = 0.
 S1_max_config = 0.
 S2_min_config = 0.
 S2_max_config = 0.
-s2_thr_config = 0.
-coin_level_config = 3
 
 
 [80t]

flamedisx/nest/get_energy_hist.py

@@ -0,0 +1,88 @@
+from concurrent.futures import ProcessPoolExecutor
+
+import numericalunits as nu
+import numpy as np
+import pandas as pd
+
+import flamedisx as fd
+import multihist as mh
+import wimprates as wr
+
+
+export, __all__ = fd.exporter()
+
+
+@export
+def get_energy_hist(
+    wimp_mass=40.0,
+    sigma=1e-45,
+    min_E=1e-2,
+    max_E=80.0,
+    n_energy_bins=800,
+    min_time="2019-09-01T08:28:00",
+    max_time="2020-09-01T08:28:00",
+    n_time_bins=25,
+    modulation=True,
+    migdal_model=None,
+):
+
+    wimprates_kwargs = dict(
+        mw=wimp_mass,
+        sigma_nucleon=sigma,
+    )
+
+    if migdal_model is not None:
+        if migdal_model not in ["Ibe", "Cox", "Cox_dipole"]:
+            raise ValueError(
+                "Invalid Migdal model. Choose from 'Ibe', 'Cox', 'Cox_dipole'"
+            )
+        dipole = False
+        if migdal_model == "Cox_dipole":
+            migdal_model = "Cox"
+            dipole = True
+
+        wimprates_kwargs.update(
+            dict(
+                detection_mechanism="migdal",
+                migdal_model=migdal_model,
+                dipole=dipole,
+            )
+        )
+
+    energy_bin_edges = np.linspace(min_E, max_E, n_energy_bins + 1)
+    energy_bin_width = (energy_bin_edges[1] - energy_bin_edges[0]) * nu.keV
+    energy_values = (energy_bin_edges[1:] + energy_bin_edges[:-1]) / 2
+
+    time_bin_edges = (
+        pd.date_range(min_time, max_time, periods=n_time_bins + 1).to_julian_date()
+        - 2451545.0  # Convert to J2000
+    )
+    times = (time_bin_edges[:-1] + time_bin_edges[1:]) / 2
+
+    scale = (
+        energy_bin_width / nu.keV
+    )  # Convert from [per keV] to [per energy_bin_width]
+
+    rates_list = []
+    if modulation:
+        with ProcessPoolExecutor() as executor:
+            for time in times:
+                rates_list.append(
+                    executor.submit(
+                        wr.rate_wimp_std, energy_values, t=time, **wimprates_kwargs
+                    )
+                )
+
+            rates_list = [future.result() for future in rates_list]
+    else:
+        rates = wr.rate_wimp_std(energy_values, **wimprates_kwargs)
+        for _ in times:
+            rates_list.append(rates)
+
+    RATES = np.array(rates_list) * scale
+
+    hist = mh.Histdd.from_histogram(
+        RATES, [time_bin_edges, energy_bin_edges], axis_names=("time", "energy")
+    )
+
+    return hist

flamedisx/nest/lxe_sources.py

@@ -1,19 +1,18 @@
-import tensorflow as tf
-
-import numpy as np
-import wimprates as wr
-from multihist import Histdd
-
+from concurrent.futures import ProcessPoolExecutor
 import configparser
+import math as m
 import os
 
-import pickle as pkl
+import numericalunits as nu
+import numpy as np
 import pandas as pd
+import tensorflow as tf
 
 import flamedisx as fd
+import multihist as mh
+import wimprates as wr
 from .. import nest as fd_nest
 
-import math as m
 pi = tf.constant(m.pi)
 
 export, __all__ = fd.exporter()
@@ -566,38 +565,109 @@ class nestWIMPSource(nestNRSource):
     """SI WIMP signal source.
     Reads in energy spectrum from .pkl file, generated with LZ's DMCalc.
     Normalised with an exposure of 1 tonne year and a cross section of 1e-45 cm^2.
-    Temporally varying between 2019-09-01T08:28:00 and 2020-09-01T08:28:00.
+    By default, temporally varying between 2019-09-01T08:28:00 and 2020-09-01T08:28:00.
     """
 
     model_blocks = (fd_nest.WIMPEnergySpectrum,) + nestNRSource.model_blocks[1:]
 
-    def __init__(self, *args, wimp_mass=40, fid_mass=1., livetime=1., **kwargs):
-        if ('detector' not in kwargs):
-            kwargs['detector'] = 'default'
-
-        energy_edges = np.geomspace(0.7, 50, 100)
-        e_centers = 0.5 * (energy_edges[1:] + energy_edges[:-1])
+    def __init__(
+        self,
+        *args,
+        wimp_mass=40,
+        sigma=1e-45,
+        fid_mass=1.0,
+        min_E=1e-2,
+        max_E=80.0,
+        n_energy_bins=800,
+        min_time="2019-09-01T08:28:00",
+        max_time="2020-09-01T08:28:00",
+        livetime=1.0,
+        n_time_bins=25,
+        modulation=False,
+        **kwargs
+    ):
+
+        if (livetime < 0 or livetime > 1) and not isinstance(livetime, float):
+            raise ValueError("Livetime should be a percentage between 0 and 1")
+
+        if "detector" not in kwargs:
+            kwargs["detector"] = "default"
+
+        self.energy_hist = fd_nest.get_energy_hist(
+            wimp_mass=wimp_mass,
+            sigma=sigma,
+            min_E=min_E,
+            max_E=max_E,
+            n_energy_bins=n_energy_bins,
+            min_time=min_time,
+            max_time=max_time,
+            n_time_bins=n_time_bins,
+            modulation=modulation,
+        )
+
+        scale = fid_mass * livetime
+        self.energy_hist *= scale
+
+        self.n_time_bins = len(self.energy_hist.bin_centers()[0])
+        e_centers = self.energy_hist.bin_centers()[1]
         self.energies = fd.np_to_tf(e_centers)
 
-        self.n_time_bins = 24
-        times = np.linspace(wr.j2000(self.model_blocks[0].t_start.value),
-                            wr.j2000(self.model_blocks[0].t_stop.value),
-                            self.n_time_bins + 1)
-        time_centers = 0.5 * (times[1:] + times[:-1])
+        self.array_columns = (("energy_spectrum", len(e_centers)),)
+
+        super().__init__(*args, **kwargs)
 
-        wimp_kwargs = dict(mw=wimp_mass,
-                           sigma_nucleon=1e-45)
-        spectra = np.array([wr.rate_wimp_std(t=t,
-                                             es=e_centers,
-                                             **wimp_kwargs)
-                            * np.diff(energy_edges)
-                            for t in time_centers])
-        assert spectra.shape == (len(time_centers), len(e_centers))
 
-        self.energy_hist = Histdd.from_histogram(
-            spectra,
-            bin_edges=(times, energy_edges)) * (fid_mass * livetime)
+@export
+class nestMigdalSource(nestERSource):
+
+    model_blocks = (fd_nest.WIMPEnergySpectrum,) + nestERSource.model_blocks[1:]
+
+    def __init__(
+        self,
+        *args,
+        wimp_mass=40,
+        sigma=1e-45,
+        fid_mass=1.0,
+        min_E=1e-2,
+        max_E=80.0,
+        n_energy_bins=800,
+        min_time="2019-09-01T08:28:00",
+        max_time="2020-09-01T08:28:00",
+        livetime=1.0,
+        n_time_bins=25,
+        modulation=True,
+        migdal_model="Cox",
+        **kwargs
+    ):
+        if migdal_model not in ["Ibe", "Cox", "Cox_dipole"]:
+            raise ValueError("Invalid Migdal model. Choose from 'Ibe', 'Cox', 'Cox_dipole'")
+
+        if (livetime < 0 or livetime > 1) and not isinstance(livetime, float):
+            raise ValueError("Livetime should be a percentage between 0 and 1")
+
+        if "detector" not in kwargs:
+            kwargs["detector"] = "default"
+
+        self.energy_hist = fd_nest.get_energy_hist(
+            wimp_mass=wimp_mass,
+            sigma=sigma,
+            min_E=min_E,
+            max_E=max_E,
+            n_energy_bins=n_energy_bins,
+            min_time=min_time,
+            max_time=max_time,
+            n_time_bins=n_time_bins,
+            modulation=modulation,
+            migdal_model=migdal_model,
+        )
+
+        scale = fid_mass * livetime
+        self.energy_hist *= scale
+
+        self.n_time_bins = len(self.energy_hist.bin_centers()[0])
+        e_centers = self.energy_hist.bin_centers()[1]
+        self.energies = fd.np_to_tf(e_centers)
 
-        self.array_columns = (('energy_spectrum', len(e_centers)),)
+        self.array_columns = (("energy_spectrum", len(e_centers)),)
 
-        super().__init__(*args, **kwargs)
+        super().__init__(*args, **kwargs)

flamedisx/nest/parameter_calc.py

@@ -181,7 +181,9 @@ def calculate_s1_mean_mult(spe_res):
 
 @export
 def get_coin_table(coin_level, num_pmts, spe_res, spe_thr, spe_eff, double_pe_fraction):
+
     assert coin_level <= 4, 'This logic will not work well for coincidence levels higher than 4'
+
     coin_dict = dict()
     coin_table = []