ncx2

StFroese · Dec 16, 2024 · 1993823 · 1993823
1 parent 60e9765
commit 1993823
Show file tree

Hide file tree

Showing 3 changed files with 61 additions and 23 deletions.
diff --git a/titrate/plotting.py b/titrate/plotting.py
@@ -4,9 +4,12 @@
 from astropy import visualization as viz
 from astropy.table import QTable, unique
 from astropy.units import Quantity
+from gammapy.modeling import Fit
 
 from titrate.datasets import AsimovMapDataset, AsimovSpectralDataset
 from titrate.statistics import QMuTestStatistic, QTildeMuTestStatistic
+from titrate.utils import copy_dataset_with_models
+
 
 STATISTICS = {"qmu": QMuTestStatistic, "qtildemu": QTildeMuTestStatistic}
 
@@ -138,12 +141,32 @@ def __init__(
         self.path = path
         self.ax = ax if ax is not None else plt.gca()
         self.analysis = analysis
+        self.measurement_dataset = measurement_dataset
+        self.poi_name = poi_name
+
+        self.d_sig = copy_dataset_with_models(self.measurement_dataset)
+        self.d_sig.models.parameters[self.poi_name].value = 1e5
+        self.d_sig.models.parameters[self.poi_name].frozen = True
+        fit = Fit()
+        _ = fit.run(self.d_sig)
+        self.d_sig.models.parameters[self.poi_name].frozen = False
+
+        self.d_bkg = copy_dataset_with_models(self.measurement_dataset)
+        self.d_bkg.models.parameters[self.poi_name].value = 0
+        self.d_bkg.models.parameters[self.poi_name].frozen = True
+        fit = Fit()
+        _ = fit.run(self.d_bkg)
+        self.d_bkg.models.parameters[self.poi_name].frozen = False
 
         if self.analysis == "3d":
-            asimov_dataset = AsimovMapDataset.from_MapDataset(measurement_dataset)
+            self.asimov_sig_dataset = AsimovMapDataset.from_MapDataset(self.d_sig)
+            self.asimov_bkg_dataset = AsimovMapDataset.from_MapDataset(self.d_bkg)
         elif self.analysis == "1d":
-            asimov_dataset = AsimovSpectralDataset.from_SpectralDataset(
-                measurement_dataset
+            self.asimov_sig_dataset = AsimovSpectralDataset.from_SpectralDataset(
+                self.d_sig
+            )
+            self.asimov_bkg_dataset = AsimovSpectralDataset.from_SpectralDataset(
+                self.d_bkg
             )
 
         try:
@@ -181,13 +204,20 @@ def __init__(
         max_ts = max(toys_ts_diff.max(), toys_ts_same.max())
         bins = np.linspace(0, max_ts, 31)
         linspace = np.linspace(0, max_ts, 1000)
-        statistic = STATISTICS[statistic](asimov_dataset, poi_name)
+        statistic_sig = STATISTICS[statistic](self.asimov_sig_dataset, poi_name)
+        statistic_bkg = STATISTICS[statistic](self.asimov_bkg_dataset, poi_name)
         statistic_math_name = (
             r"q_\mu" if isinstance(statistic, QMuTestStatistic) else r"\tilde{q}_\mu"
         )
 
         self.plot(
-            linspace, bins, toys_ts_same, toys_ts_diff, statistic, statistic_math_name
+            linspace,
+            bins,
+            toys_ts_same,
+            toys_ts_diff,
+            statistic_sig,
+            statistic_bkg,
+            statistic_math_name,
         )
 
         self.ax.set_yscale("log")
@@ -199,37 +229,44 @@ def __init__(
         self.ax.legend()
 
     def plot(
-        self, linspace, bins, toys_ts_same, toys_ts_diff, statistic, statistic_math_name
+        self,
+        linspace,
+        bins,
+        toys_ts_same,
+        toys_ts_diff,
+        statistic_sig,
+        statistic_bkg,
+        statistic_math_name,
     ):
         plt.hist(
             toys_ts_diff,
             bins=bins,
             density=True,
             histtype="step",
             color="C0",
-            label=(r"$\mu=1$, $\mu^\prime=0$"),
+            label=(r"$\mu=10^5$, $\mu^\prime=0$"),
         )
         plt.hist(
             toys_ts_same,
             bins=bins,
             density=True,
             histtype="step",
             color="C1",
-            label=(r"$\mu=1$, $\mu^\prime=1$"),
+            label=(r"$\mu=10^5$, $\mu^\prime=10^5$"),
         )
 
         plt.plot(
             linspace,
-            statistic.asympotic_approximation_pdf(
-                poi_val=1, same=False, poi_true_val=0, ts_val=linspace
+            statistic_bkg.asympotic_approximation_pdf(
+                poi_val=1e5, same=False, poi_true_val=0, ts_val=linspace
             ),
             color="C0",
-            label=r"$\mu=1$, $\mu^\prime=0$",
+            label=r"$\mu=10^5$, $\mu^\prime=0$",
         )
         plt.plot(
             linspace,
-            statistic.asympotic_approximation_pdf(poi_val=1, ts_val=linspace),
+            statistic_sig.asympotic_approximation_pdf(poi_val=1e5, ts_val=linspace),
             color="C1",
             ls="--",
-            label=r"$\mu=1$, $\mu^\prime=1$",
+            label=r"$\mu=10^5$, $\mu^\prime=10^5$",
         )
diff --git a/titrate/statistics.py b/titrate/statistics.py
@@ -1,9 +1,8 @@
 import abc
-from scipy.stats import ncx2
 
 import numpy as np
 from gammapy.modeling import Fit
-from scipy.stats import kstwo, norm
+from scipy.stats import kstwo, ncx2, norm
 
 from titrate.datasets import AsimovMapDataset, AsimovSpectralDataset
 
@@ -253,7 +252,7 @@ def asympotic_approximation_pdf(
             nc = 0
         else:
             sigma = self.sigma(poi_val, poi_true_val)
-            nc = (poi_val - poi_true_val) / sigma
+            nc = (poi_val - poi_true_val) ** 2 / sigma**2
         return ncx2.pdf(ts_val, nc=nc, df=1)
 
         if same:

diff --git a/titrate/validation.py b/titrate/validation.py
@@ -38,7 +38,7 @@ def __init__(
 
         self.measurement_dataset = measurement_dataset
         self.d_sig = copy_dataset_with_models(self.measurement_dataset)
-        self.d_sig.models.parameters[self.poi_name].value = 1
+        self.d_sig.models.parameters[self.poi_name].value = 1e5
         self.d_sig.models.parameters[self.poi_name].frozen = True
         fit = Fit()
         _ = fit.run(self.d_sig)
@@ -96,12 +96,12 @@ def validate(self, n_toys=1000):
         ks_diff = kstest(
             self.toys_ts_diff[self.toys_ts_diff_valid],
             lambda x: stat_bkg.asympotic_approximation_cdf(
-                poi_val=1, same=False, poi_true_val=0, ts_val=x
+                poi_val=1e5, same=False, poi_true_val=0, ts_val=x
             ),
         )
         ks_same = kstest(
             self.toys_ts_same[self.toys_ts_same_valid],
-            lambda x: stat_sig.asympotic_approximation_cdf(poi_val=1, ts_val=x),
+            lambda x: stat_sig.asympotic_approximation_cdf(poi_val=1e5, ts_val=x),
         )
 
         valid = ks_diff > 0.05 and ks_same > 0.05
@@ -110,8 +110,10 @@ def validate(self, n_toys=1000):
 
     def generate_datasets(self, n_toys):
         if self.path is None:
-            toys_ts_diff, toys_ts_diff_valid = self.toys_ts(n_toys, 1, 0)
-            toys_ts_same, toys_ts_same_valid = self.toys_ts(n_toys, 1, 1)
+            toys_ts_diff, toys_ts_diff_valid = self.toys_ts(n_toys, 1e5, 0, self.d_bkg)
+            toys_ts_same, toys_ts_same_valid = self.toys_ts(
+                n_toys, 1e5, 1e5, self.d_sig
+            )
         else:
             (
                 toys_ts_diff,
@@ -126,12 +128,12 @@ def generate_datasets(self, n_toys):
         self.toys_ts_same_valid = toys_ts_same_valid
 
     @lru_cache
-    def toys_ts(self, n_toys, poi_val, poi_true_val):
+    def toys_ts(self, n_toys, poi_val, poi_true_val, dataset):
         with ProcessPoolExecutor(self.max_workers) as pool:
             futures = [
                 pool.submit(
                     calc_ts_toyMC,
-                    self.measurement_dataset,
+                    dataset,
                     self.statistic,
                     poi_val,
                     poi_true_val,