Merge pull request #7 from StFroese/ulfactory

add ULFactory

README.md

@@ -1,5 +1,5 @@
 # *asymp*T*otic* *l*I*kelihood-based* T*ests* *for* *da*R*k* *m*A*t*T*er* *s*E*arch*
-<img src="https://anaconda.org/conda-forge/titrate/badges/version.svg"> <img src="https://anaconda.org/conda-forge/titrate/badges/platforms.svg"> <img src="https://badge.fury.io/py/titrate.svg">
+[<img src="https://anaconda.org/conda-forge/titrate/badges/version.svg">](https://anaconda.org/conda-forge/titrate) <img src="https://anaconda.org/conda-forge/titrate/badges/platforms.svg"> [<img src="https://badge.fury.io/py/titrate.svg">](https://pypi.org/project/titrate/)
 
 This package is based on the paper [Asymptotic formulae for likelihood-based tests of new physics](https://arxiv.org/abs/1007.1727).
 
@@ -26,8 +26,8 @@ What have I learned?
 1. Adding Asimov datasets to gammapy :heavy_check_mark:
 2. Adding test statistics to gammapy :heavy_check_mark:
 3. Validation of test statistics :heavy_check_mark:
-4. Calculation of ULs
-5. Calculation of median ULs and bands with asymptotic formulae and asimov datasets (faster than toy MCs)
+4. Calculation of ULs :heavy_check_mark:
+5. Calculation of median ULs and bands with asymptotic formulae and asimov datasets (faster than toy MCs) :heavy_check_mark:
 
 ## Disclaimer
 This is not a finished version yet and can contain drastic changes

setup.cfg

@@ -1,6 +1,6 @@
 [metadata]
 name = titrate
-version = 0.3.0
+version = 0.4.0
 author = Stefan Fröse
 author_email = [email protected]
 description = asympTotic lIkelihood Tests for daRk mAtTer sEarch
@@ -27,6 +27,8 @@ install_requires =
     matplotlib
     gammapy
     joblib
+    astropy
+    h5py
 
 [options.extras_require]
 tests =

titrate/conftest.py

@@ -99,14 +99,20 @@ def ursa_major_ii_profile():
 
 
 @pytest.fixture(scope="module")
-def dm_models(observation, geometry2d, ursa_major_ii_profile):
+def jfact_map(geometry2d, ursa_major_ii_profile):
     jfactory = JFactory(
         geom=geometry2d,
         profile=ursa_major_ii_profile,
         distance=ursa_major_ii_profile.DISTANCE_GC,
     )
     jfactor = jfactory.compute_differential_jfactor()
     jfact_map = WcsNDMap(geom=geometry2d, data=jfactor.value, unit=jfactor.unit)
+
+    return jfact_map
+
+
+@pytest.fixture(scope="module")
+def dm_models(jfact_map):
     spatial_model = TemplateSpatialModel(jfact_map, normalize=False)
 
     spectral_model = DarkMatterAnnihilationSpectralModel(mass=50 * u.TeV, channel="b")

titrate/plotting.py

@@ -0,0 +1,91 @@
+import h5py
+import matplotlib.pyplot as plt
+from astropy import visualization as viz
+from astropy.table import QTable, unique
+
+
+class UpperLimitPlotter:
+    def __init__(self, path, channel, axes=None):
+        self.path = path
+        self.axes = axes if axes is not None else plt.gca()
+
+        try:
+            table = QTable.read(self.path, path=channel)
+        except OSError:
+            channels = list(h5py.File("/Users/stefan/Downloads/test.hdf5").keys())
+            channels = [ch for ch in channels if "meta" not in ch]
+            raise KeyError(
+                f"Channel {channel} not in dataframe. " f"Choose from {channels}"
+            )
+
+        self.channel = channel
+
+        masses = table["mass"]
+        uls = table["ul"]
+        median = table["median_ul"]
+        one_sigma_minus = table["1sigma_minus_ul"]
+        one_sigma_plus = table["1sigma_plus_ul"]
+        two_sigma_minus = table["2sigma_minus_ul"]
+        two_sigma_plus = table["2sigma_plus_ul"]
+
+        with viz.quantity_support():
+            self.plot_channel(
+                masses,
+                uls,
+                median,
+                one_sigma_minus,
+                one_sigma_plus,
+                two_sigma_minus,
+                two_sigma_plus,
+            )
+
+        self.axes.set_xscale("log")
+        self.axes.set_yscale("log")
+
+        cl_type = unique(table[table["channel"] == self.channel], keys="cl_type")[
+            "cl_type"
+        ][0]
+        cl = unique(table[table["channel"] == self.channel], keys="cl")["cl"][0]
+        self.axes.set_xlabel(f"m / {masses.unit:latex}")
+        self.axes.set_ylabel(
+            rf"$CL_{cl_type}^{{{cl}}}$ upper limit on $< \sigma v>$ / {uls.unit:latex}"
+        )
+
+        self.axes.set_title(f"Annihilation Upper Limits for channel {self.channel}")
+
+        self.axes.legend()
+
+    def plot_channel(
+        self,
+        masses,
+        uls,
+        median,
+        one_sigma_minus,
+        one_sigma_plus,
+        two_sigma_minus,
+        two_sigma_plus,
+    ):
+        self.axes.plot(masses, uls, color="tab:orange", label="Upper Limits")
+        self.axes.plot(masses, median, color="tab:blue", label="Expected Upper Limits")
+        self.axes.fill_between(
+            masses,
+            median,
+            one_sigma_plus,
+            color="tab:blue",
+            alpha=0.75,
+            label=r"$1\sigma$-region",
+        )
+        self.axes.fill_between(
+            masses, median, one_sigma_minus, color="tab:blue", alpha=0.75
+        )
+        self.axes.fill_between(
+            masses,
+            one_sigma_plus,
+            two_sigma_plus,
+            color="tab:blue",
+            alpha=0.5,
+            label=r"$2\sigma$-region",
+        )
+        self.axes.fill_between(
+            masses, one_sigma_minus, two_sigma_minus, color="tab:blue", alpha=0.5
+        )

titrate/tests/test_upperlimits.py

@@ -1,5 +1,8 @@
+import astropy.units as u
+import matplotlib.pyplot as plt
 import numpy as np
 import pytest
+from astropy.table import QTable
 
 
 @pytest.mark.parametrize("cl_type", ["s", "s+b"])
@@ -21,3 +24,44 @@ def test_ULCalculator(measurement_dataset, statistic, cl_type):
     assert bands["2sig"][0] < bands["1sig"][0]
     assert bands["med"] < bands["1sig"][1]
     assert bands["1sig"][1] < bands["2sig"][1]
+
+
+@pytest.fixture(scope="module")
+def upperlimits_file(jfact_map, measurement_dataset, tmp_path_factory):
+    from titrate.upperlimits import ULFactory
+
+    ulfactory = ULFactory(
+        measurement_dataset, ["b", "W"], 0.1 * u.TeV, 100 * u.TeV, 5, jfact_map
+    )
+    ulfactory.compute()
+
+    data = tmp_path_factory.mktemp("data")
+    ulfactory.save_results(f"{data}/ul.hdf5")
+
+    return f"{data}/ul.hdf5"
+
+
+@pytest.mark.parametrize("channel", ["b", "W"])
+def test_ULFactory(upperlimits_file, channel):
+    table = QTable.read(upperlimits_file, path=channel)
+    assert np.all(table["mass"] == np.geomspace(0.1, 100, 5) * u.TeV)
+    assert len(table["ul"]) == 5
+    assert len(table["median_ul"]) == 5
+    assert len(table["1sigma_minus_ul"]) == 5
+    assert len(table["1sigma_plus_ul"]) == 5
+    assert len(table["2sigma_minus_ul"]) == 5
+    assert len(table["2sigma_plus_ul"]) == 5
+    assert len(table["cl_type"]) == 5
+    assert np.all(table["cl_type"] == "s")
+    assert len(table["cl"]) == 5
+    assert np.all(table["cl"] == 0.95)
+    assert np.all(table["channel"] == sorted([channel] * 5)[::-1])
+
+
+def test_UpperLimitPlotter(upperlimits_file):
+    from titrate.plotting import UpperLimitPlotter
+
+    fig, axs = plt.subplots(nrows=1, ncols=2)
+
+    for channel, ax in zip(["b", "W"], np.array(axs).reshape(-1)):
+        UpperLimitPlotter(upperlimits_file, channel=channel, axes=ax)

titrate/upperlimits.py

@@ -1,18 +1,35 @@
+import astropy.units as u
 import numpy as np
+from astropy.table import QTable
+from gammapy.astro.darkmatter import DarkMatterAnnihilationSpectralModel
 from gammapy.modeling import Fit
+from gammapy.modeling.models import (
+    FoVBackgroundModel,
+    Models,
+    SkyModel,
+    TemplateSpatialModel,
+)
+from joblib import Parallel, delayed
 from scipy.interpolate import interp1d
 from scipy.optimize import brentq
 from scipy.stats import norm
 
 from titrate.datasets import AsimovMapDataset
 from titrate.statistics import QMuTestStatistic, QTildeMuTestStatistic
+from titrate.utils import copy_models_to_dataset
 
 STATISTICS = {"qmu": QMuTestStatistic, "qtildemu": QTildeMuTestStatistic}
+CS = DarkMatterAnnihilationSpectralModel.THERMAL_RELIC_CROSS_SECTION
 
 
 class ULCalculator:
     def __init__(
-        self, measurement_dataset, statistic="qmu", poi_name="", cl=0.95, cl_type="s"
+        self,
+        measurement_dataset,
+        statistic="qmu",
+        poi_name="scale",
+        cl=0.95,
+        cl_type="s",
     ):
         self.measurement_dataset = measurement_dataset
         self.poi_name = poi_name
@@ -36,18 +53,20 @@ def compute(self):
             poi_ul = 1e-2
         else:
             poi_ul = poi_best + 0.01 * poi_best
-        while self.pvalue(poi_ul, cl_type=self.cl_type) > 1 - self.cl:
+        while self.pvalue(poi_ul, cl_type=self.cl_type) > 1 - self.cl or np.isnan(
+            self.pvalue(poi_ul, cl_type=self.cl_type)
+        ):
             poi_ul *= 2
 
-        interp_ul_points = np.linspace(poi_ul / 10, poi_ul, 10)
+        interp_ul_points = np.linspace(poi_ul / 2, poi_ul, 10)
         interp_pvalues = np.array(
             [
                 self.pvalue(interp_ul, cl_type=self.cl_type)
                 for interp_ul in interp_ul_points
             ]
         ).ravel()
         interpolation = interp1d(interp_ul_points, interp_pvalues - 1 + self.cl)
-        poi_ul = brentq(interpolation, poi_ul / 10, poi_ul)
+        poi_ul = brentq(interpolation, poi_ul / 2, poi_ul)
 
         return poi_ul
 
@@ -114,3 +133,119 @@ def compute_band(self, sigma, n_sigma, cl_type):
             return sigma * (norm.ppf(self.cl) + n_sigma)
         elif cl_type == "s":
             return sigma * (norm.ppf(1 - (1 - self.cl) * norm.pdf(n_sigma)) + n_sigma)
+
+
+class ULFactory:
+    def __init__(
+        self,
+        measurement_dataset,
+        channels,
+        mass_min,
+        mass_max,
+        n_steps,
+        jfactor_map,
+        cl_type="s",
+        cl=0.95,
+        **kwargs,
+    ):
+        self.measurement_dataset = measurement_dataset
+        self.channels = channels
+        self.masses = np.geomspace(
+            mass_min.to_value("TeV"), mass_max.to_value("TeV"), n_steps
+        )
+        self.jfactor_map = jfactor_map
+        self.cl_type = cl_type
+        self.cl = cl
+        self.kwargs = kwargs
+        self.uls = None
+        self.expected_uls = None
+
+    def setup_models(self):
+        models = []
+        spatial_model = TemplateSpatialModel(self.jfactor_map, normalize=False)
+        bkg_model = FoVBackgroundModel(dataset_name="foo")
+        for channel in self.channels:
+            for mass in self.masses:
+                spectral_model = DarkMatterAnnihilationSpectralModel(
+                    mass=mass * u.TeV, channel=channel
+                )
+                sky_model = SkyModel(
+                    spatial_model=spatial_model,
+                    spectral_model=spectral_model,
+                    name=f"mass_{mass:.2f}TeV_channel_{channel}",
+                )
+                models.append(Models([sky_model, bkg_model]))
+
+        return models
+
+    def setup_calculator(self, models):
+        measurement_copy = self.measurement_dataset.copy()
+        copy_models_to_dataset(models, measurement_copy)
+        return ULCalculator(measurement_copy, **self.kwargs)
+
+    def compute_uls(self):
+        uls = Parallel(n_jobs=-1, verbose=0)(
+            delayed(self.setup_calculator(models).compute)()
+            for models in self.setup_models()
+        )
+        return uls
+
+    def compute_expected(self):
+        expected_uls = Parallel(n_jobs=-1, verbose=0)(
+            delayed(self.setup_calculator(models).expected_uls)()
+            for models in self.setup_models()
+        )
+        return expected_uls
+
+    def compute(self):
+        self.uls = self.compute_uls()
+        self.expected_uls = self.compute_expected()
+
+    def save_results(self, path, overwrite=False, **kwargs):
+        if self.uls is None or self.expected_uls is None:
+            raise ValueError("No results computed yet. Run compute() first.")
+
+        # prepare uls
+        n_channels = len(self.channels)
+        uls = np.array(self.uls).reshape(n_channels, -1)
+        median_uls = np.array([ul["med"] for ul in self.expected_uls]).reshape(
+            n_channels, -1
+        )
+        one_sigma_minus_uls = np.array(
+            [ul["1sig"][0] for ul in self.expected_uls]
+        ).reshape(n_channels, -1)
+        one_sigma_plus_uls = np.array(
+            [ul["1sig"][1] for ul in self.expected_uls]
+        ).reshape(n_channels, -1)
+        two_sigma_minus_uls = np.array(
+            [ul["2sig"][0] for ul in self.expected_uls]
+        ).reshape(n_channels, -1)
+        two_sigma_plus_uls = np.array(
+            [ul["2sig"][1] for ul in self.expected_uls]
+        ).reshape(n_channels, -1)
+
+        for ch_idx, channel in enumerate(self.channels):
+            # to dict
+            ul_dict = {
+                "mass": self.masses * u.TeV,
+                "channel": np.repeat(channel, len(self.masses)),
+                "cl_type": np.repeat(self.cl_type, len(self.masses)),
+                "cl": np.repeat(self.cl, len(self.masses)),
+                "ul": uls[ch_idx] * CS,
+                "median_ul": median_uls[ch_idx] * CS,
+                "1sigma_minus_ul": one_sigma_minus_uls[ch_idx] * CS,
+                "1sigma_plus_ul": one_sigma_plus_uls[ch_idx] * CS,
+                "2sigma_minus_ul": two_sigma_minus_uls[ch_idx] * CS,
+                "2sigma_plus_ul": two_sigma_plus_uls[ch_idx] * CS,
+            }
+            print(ul_dict)
+            qtable = QTable(ul_dict)
+            qtable.write(
+                path,
+                format="hdf5",
+                path=f"{channel}",
+                overwrite=overwrite,
+                append=True,
+                serialize_meta=True,
+                **kwargs,
+            )