7049225217

pending completion

Build # 7049225217

Build Type

Pull #333

github

Committed by

web-flow

Commit Message

Merge d8ca38eb7 into 956b8ca59

Pull Request Pull Request #333: Implementation of the KDE/photospline PSF (9yr/11yr NT samples) as spatial PDF for flarestack

Run Details

24 of 171 new or added lines in 4 files covered. (14.04%)

4442 of 5948 relevant lines covered (74.68%)

1.49 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

69.61

/flarestack/core/spatial_pdf.py

import numpy as np
import healpy as hp
import os
import logging
from scipy.stats import norm
from scipy.interpolate import interp1d
from numpy.lib.recfunctions import append_fields
from flarestack.core.astro import angular_distance
from flarestack.shared import bkg_spline_path
from flarestack.utils.make_SoB_splines import load_bkg_spatial_spline
from photospline import SplineTable

logger = logging.getLogger(__name__)


class SpatialPDF:
    """General SpatialPDF holder class. Has separate signal and background
    spatial PDF objects.
    """

    def __init__(self, spatial_pdf_dict, season):
        self.signal = SignalSpatialPDF.create(spatial_pdf_dict)
        self.background = BackgroundSpatialPDF.create(spatial_pdf_dict, season)

        self.simulate_distribution = self.signal.simulate_distribution
        self.signal_spatial = self.signal.signal_spatial
        self.rotate_to_position = self.signal.rotate_to_position

        self.background_spatial = self.background.background_spatial


# ==============================================================================
# Signal Spatial PDFs
# ==============================================================================


class SignalSpatialPDF:
    """Base Signal Spatial PDF class."""

    subclasses: dict[str, object] = {}

    def __init__(self, spatial_pdf_dict):
        pass

    @staticmethod
    def simulate_distribution(source, data):
        return data

    @staticmethod
    def signal_spatial(source, events):
        return

    @classmethod
    def register_subclass(cls, spatial_pdf_name):
        """Adds a new subclass of SpatialPDF, with class name equal to
        "spatial_pdf_name".
        """

        def decorator(subclass):
            cls.subclasses[spatial_pdf_name] = subclass
            return subclass

        return decorator

    @classmethod
    def create(cls, s_pdf_dict):
        try:
            s_pdf_name = s_pdf_dict["spatial_pdf_name"]
        except KeyError:
            s_pdf_name = "circular_gaussian"

        if s_pdf_name not in cls.subclasses:
            raise ValueError(
                "Bad Signal Spatial PDF name {0} \n"
                "Available names are {1}".format(s_pdf_name, cls.subclasses)
            )

        return cls.subclasses[s_pdf_name](s_pdf_dict)

    @staticmethod
    def rotate(ra1, dec1, ra2, dec2, ra3, dec3):
        """Rotate ra1 and dec1 in a way that ra2 and dec2 will exactly map
        onto ra3 and dec3, respectively. All angles are treated as radians.
        Essentially rotates the events, so that they behave as if they were
        originally incident on the source.

        :param ra1: Event Right Ascension
        :param dec1: Event Declination
        :param ra2: True Event Right Ascension
        :param dec2: True Event Declination
        :param ra3: Source Right Ascension
        :param dec3: Source Declination
        :return: Returns new Right Ascensions and Declinations
        """
        # Turns Right Ascension/Declination into Azimuth/Zenith for healpy
        phi1 = ra1 - np.pi
        zen1 = np.pi / 2.0 - dec1
        phi2 = ra2 - np.pi
        zen2 = np.pi / 2.0 - dec2
        phi3 = ra3 - np.pi
        zen3 = np.pi / 2.0 - dec3

        # Rotate each ra1 and dec1 towards the pole?
        x = np.array(
            [
                hp.rotator.rotateDirection(
                    hp.rotator.get_rotation_matrix((dp, -dz, 0.0))[0], z, p
                )
                for z, p, dz, dp in zip(zen1, phi1, zen2, phi2)
            ]
        )

        # Rotate **all** these vectors towards ra3, dec3 (source_path)
        zen, phi = hp.rotator.rotateDirection(
            np.dot(
                hp.rotator.get_rotation_matrix((-phi3, 0, 0))[0],
                hp.rotator.get_rotation_matrix((0, zen3, 0.0))[0],
            ),
            x[:, 0],
            x[:, 1],
        )

        dec = np.pi / 2.0 - zen

        ra = phi + np.pi
        return np.atleast_1d(ra), np.atleast_1d(dec)

    def rotate_to_position(self, ev, ra, dec):
        """Modifies the events by reassigning the Right Ascension and
        Declination of the events. Rotates the events, so that they are
        distributed as if they originated from the source. Removes the
        additional Monte Carlo information from sampled events, so that they
        appear like regular data.

        The fields removed are:
            True Right Ascension,
            True Declination,
            True Energy,
            OneWeight

        :param ev: Events
        :param ra: Source Right Ascension (radians)
        :param dec: Source Declination (radians)
        :return: Events (modified)
        """
        names = ev.dtype.names

        # Rotates the events to lie on the source
        ev["ra"], rot_dec = self.rotate(
            ev["ra"], np.arcsin(ev["sinDec"]), ev["trueRa"], ev["trueDec"], ra, dec
        )

        if "dec" in names:
            ev["dec"] = rot_dec
        ev["sinDec"] = np.sin(rot_dec)

        # Deletes the Monte Carlo information from sampled events
        non_mc = [
            name for name in names if name not in ["trueRa", "trueDec", "trueE", "ow"]
        ]
        ev = ev[non_mc].copy()

        return ev


@SignalSpatialPDF.register_subclass("circular_gaussian")
class CircularGaussian(SignalSpatialPDF):
    def simulate_distribution(self, source, data):
        data["ra"] = np.pi + norm.rvs(size=len(data)) * data["sigma"]
        data["dec"] = norm.rvs(size=len(data)) * data["sigma"]
        data["sinDec"] = np.sin(data["dec"])
        data = append_fields(
            data,
            ["trueRa", "trueDec"],
            [np.ones_like(data["dec"]) * np.pi, np.zeros_like(data["dec"])],
        ).copy()

        data = self.rotate_to_position(data, source["ra_rad"], source["dec_rad"]).copy()

        return data.copy()

    @staticmethod
    def signal_spatial(source, cut_data):
        """Calculates the angular distance between the source and the
        coincident dataset. Uses a Gaussian PDF function, centered on the
        source. Returns the value of the Gaussian at the given distances.

        :param source: Single Source
        :param cut_data: Subset of Dataset with coincident events
        :return: Array of Spatial PDF values
        """
        distance = angular_distance(
            cut_data["ra"], cut_data["dec"], source["ra_rad"], source["dec_rad"]
        )
        space_term = (
            1.0
            / (2.0 * np.pi * cut_data["sigma"] ** 2.0)
            * np.exp(-0.5 * (distance / cut_data["sigma"]) ** 2.0)
        )

        return space_term


@SignalSpatialPDF.register_subclass("northern_tracks_kde")
class NorthernTracksKDE(SignalSpatialPDF):
    """Spatial PDF class for use with the KDE-smoothed MC PDFs introduced for the 10-yr NT analysis.
    Current limitation: In the NT analysis, this PDF depends on the spectral index gamma . Here
    a fixed gamma is used (either by setting 'spatial_pdf_data' to the location of the
    4D-photospline-tables of the PDF and 'spatial_pdf_index' to the preferred gamma value, or
    (more efficiently) setting 'spatial_pdf_data' directly to the corresponding 3D-spline table.
    """

    KDEspline = None
    SplineIs4D = False

    def __init__(self, spatial_pdf_dict):
        super().__init__(spatial_pdf_dict)
        assert "spatial_pdf_data" in spatial_pdf_dict.keys() and os.path.exists(
            spatial_pdf_dict["spatial_pdf_data"]
        )
        KDEfile = spatial_pdf_dict["spatial_pdf_data"]

        NorthernTracksKDE.KDEspline = SplineTable(KDEfile)
        if NorthernTracksKDE.KDEspline.ndim == 3:
            NorthernTracksKDE.SplineIs4D = False
        elif NorthernTracksKDE.KDEspline.ndim == 4:
            NorthernTracksKDE.SplineIs4D = True
        else:
            raise RuntimeError(
                f"{KDEfile} does not seem to be a valid photospline table for the PSF"
            )

        logger.info("Using KDE spatial PDF from file {0}.".format(KDEfile))

    def _inverse_cdf(sigma, logE, gamma, npoints=100):
        psi_range = np.linspace(0.001, 0.5, npoints)
        d_psi = psi_range[1] - psi_range[0]
        if NorthernTracksKDE.SplineIs4D:
            psi_pdf = (
                d_psi
                / (log(10) * psi_range)
                * self.KDEspline.evaluate_simple(
                    [log10(sigma), logE, log10(psi_range), NorthernTracksKDE.gamma]
                )
            )
        else:
            psi_pdf = (
                d_psi
                / (log(10) * psi_range)
                * self.KDEspline.evaluate_simple([log10(sigma), logE, log10(psi_range)])
            )
        psi_cdf = np.insert(psi_pdf.cumsum(), 0, 0)
        psi_range = np.insert(psi_range, 0, 0)
        psi_cdf /= psi_cdf[-1]
        psi_cdf_unique, unique_indices = np.unique(psi_cdf, return_index=True)
        psi_range_unique = psi_range[unique_indices]
        return interp1d(psi_cdf_unique, psi_range_unique, "cubic")

    def simulate_distribution(self, source, data, gamma=2.0):
        nevents = len(data)
        phi = np.random.rand(nevents) * 2.0 * np.pi
        distance = np.random.rand(nevents)
        for _i in range(nevents):
            event_icdf = self._inverse_cdf(data["sigma"][_i], data["logE"][_i], gamma)
            distance = event_icdf(distance)

        data["ra"] = np.pi + distance * np.cos(phi)
        data["dec"] = distance * np.sin(phi)
        data["sinDec"] = np.sin(data["dec"])
        data = append_fields(
            data,
            ["trueRa", "trueDec"],
            [np.ones_like(data["dec"]) * np.pi, np.zeros_like(data["dec"])],
        ).copy()

        data = self.rotate_to_position(data, source["ra_rad"], source["dec_rad"]).copy()

        return data.copy()

    @staticmethod
    def signal_spatial(source, cut_data, gamma=2.0):
        """Calculates the angular distance between the source and the
        coincident dataset. This class provides an interface for the KDE-smoothed MC PDF introduced for the 10yr NT analysis.
        Returns the value of the PDF at the given distances between the source and the events.

        :param source: Single Source
        :param cut_data: Subset of Dataset with coincident events
        :return: Array of Spatial PDF values
        """

        # logger.debug(f"signal_spatial called with gamma={gamma}.")

        distance = angular_distance(
            cut_data["ra"], cut_data["dec"], source["ra_rad"], source["dec_rad"]
        )

        if NorthernTracksKDE.SplineIs4D:
            space_term = NorthernTracksKDE.KDEspline.evaluate_simple(
                [
                    np.log10(cut_data["sigma"]),
                    cut_data["logE"],
                    np.log10(distance),
                    gamma,
                ]
            )
        else:
            space_term = NorthernTracksKDE.KDEspline.evaluate_simple(
                [np.log10(cut_data["sigma"]), cut_data["logE"], np.log10(distance)]
            )

        space_term /= 2 * np.pi * np.log(10) * (distance**2)

        return space_term


# ==============================================================================
# Background Spatial PDFs
# ==============================================================================


class BackgroundSpatialPDF:
    subclasses: dict[str, object] = {}

    def __init__(self, spatial_pdf_dict, season):
        pass

    @classmethod
    def register_subclass(cls, bkg_spatial_name):
        """Adds a new subclass of BackgroundSpatialPDF, with class name equal to
        "spatial_pdf_name".
        """

        def decorator(subclass):
            cls.subclasses[bkg_spatial_name] = subclass
            return subclass

        return decorator

    @classmethod
    def create(cls, s_pdf_dict, season):
        try:
            s_pdf_name = s_pdf_dict["bkg_spatial_pdf"]
        except KeyError:
            s_pdf_name = "zenith_spline"

        if s_pdf_name not in cls.subclasses:
            raise ValueError(
                "Bad Background Spatial PDF name {0} \n"
                "Available names are {1}".format(s_pdf_name, cls.subclasses)
            )

        return cls.subclasses[s_pdf_name](s_pdf_dict, season)

    def background_spatial(self, events):
        return np.ones(len(events))


@BackgroundSpatialPDF.register_subclass("uniform")
class UniformPDF(BackgroundSpatialPDF):
    """A highly-simplified spatial PDF in which events are distributed
    uniformly over the celestial sphere.
    """

    def background_spatial(self, events):
        space_term = (1.0 / (4.0 * np.pi)) * np.ones(len(events))
        return space_term


@BackgroundSpatialPDF.register_subclass("uniform_solid_angle")
class UniformSolidAngle(BackgroundSpatialPDF):
    """Generic class for a background PDF that is uniform over some fixed
    area, and 0 otherwise. Requires an argument to be passed in the
    'spatial_pdf_dict', with key 'background_solid_angle'. In the limit of a
    solid angle of 4 pi, this becomes identical to the UniformPDF class.
    """

    def __init__(self, spatial_pdf_dict, season):
        BackgroundSpatialPDF.__init__(self, spatial_pdf_dict, season)

        try:
            self.solid_angle = spatial_pdf_dict["background_solid_angle"]
        except KeyError:
            raise KeyError(
                "No solid angle passed to UniformSolidAngle class.\n"
                "Please include an entry 'background_solid_angle' "
                "in the 'spatial_pdf_dict'."
            )

        if self.solid_angle > 4 * np.pi:
            raise ValueError(
                "Solid angle {0} was provided, but this is "
                "larger than 4pi. Please provide a valid solid "
                "angle."
            )

    def background_spatial(self, events):
        space_term = (1.0 / self.solid_angle) * np.ones(len(events))
        return space_term


@BackgroundSpatialPDF.register_subclass("zenith_spline")
class ZenithSpline(BackgroundSpatialPDF):
    """A 1D background spatial PDF, in which the background is assumed to be
    uniform in azimuth, but varying as a function of zenith. A
    spline is used to parameterise this distribution.
    """

    def __init__(self, spatial_pdf_dict, season):
        BackgroundSpatialPDF.__init__(self, spatial_pdf_dict, season)
        self.bkg_f = self.create_background_function(season)

    @staticmethod
    def create_background_function(season):
        # Checks if background spatial spline has been created

        if not os.path.isfile(bkg_spline_path(season)):
            season.make_background_spatial()

        return load_bkg_spatial_spline(season)

    def background_spatial(self, events):
        space_term = (1.0 / (2.0 * np.pi)) * np.exp(self.bkg_f(events["sinDec"]))

        return space_term

1	import numpy as np	3✔
2	import healpy as hp	3✔
3	import os	3✔
4	import logging	3✔
5	from scipy.stats import norm	3✔
6	from scipy.interpolate import interp1d	3✔
7	from numpy.lib.recfunctions import append_fields	3✔
8	from flarestack.core.astro import angular_distance	3✔
9	from flarestack.shared import bkg_spline_path	3✔
10	from flarestack.utils.make_SoB_splines import load_bkg_spatial_spline	3✔
11	from photospline import SplineTable	3✔
12
13	logger = logging.getLogger(__name__)	3✔
14
15
16	class SpatialPDF:	3✔
17	"""General SpatialPDF holder class. Has separate signal and background
18	spatial PDF objects.
19	"""
20
21	def __init__(self, spatial_pdf_dict, season):	3✔
22	self.signal = SignalSpatialPDF.create(spatial_pdf_dict)	3✔
23	self.background = BackgroundSpatialPDF.create(spatial_pdf_dict, season)	3✔
24
25	self.simulate_distribution = self.signal.simulate_distribution	3✔
26	self.signal_spatial = self.signal.signal_spatial	3✔
27	self.rotate_to_position = self.signal.rotate_to_position	3✔
28
29	self.background_spatial = self.background.background_spatial	3✔
30
31
32	# ==============================================================================
33	# Signal Spatial PDFs
34	# ==============================================================================
35
36
37	class SignalSpatialPDF:	3✔
38	"""Base Signal Spatial PDF class."""
39
40	subclasses: dict[str, object] = {}	3✔
41
42	def __init__(self, spatial_pdf_dict):	3✔
43	pass	3✔
44
45	@staticmethod	3✔
46	def simulate_distribution(source, data):	3✔
47	return data	×
48
49	@staticmethod	3✔
50	def signal_spatial(source, events):	3✔
51	return	×
52
53	@classmethod	3✔
54	def register_subclass(cls, spatial_pdf_name):	3✔
55	"""Adds a new subclass of SpatialPDF, with class name equal to
56	"spatial_pdf_name".
57	"""
58
59	def decorator(subclass):	3✔
60	cls.subclasses[spatial_pdf_name] = subclass	3✔
61	return subclass	3✔
62
63	return decorator	3✔
64
65	@classmethod	3✔
66	def create(cls, s_pdf_dict):	3✔
67	try:	3✔
68	s_pdf_name = s_pdf_dict["spatial_pdf_name"]	3✔
69	except KeyError:	3✔
70	s_pdf_name = "circular_gaussian"	3✔
71
72	if s_pdf_name not in cls.subclasses:	3✔
73	raise ValueError(	×
74	"Bad Signal Spatial PDF name {0} \n"
75	"Available names are {1}".format(s_pdf_name, cls.subclasses)
76	)
77
78	return cls.subclasses[s_pdf_name](s_pdf_dict)	3✔
79
80	@staticmethod	3✔
81	def rotate(ra1, dec1, ra2, dec2, ra3, dec3):	3✔
82	"""Rotate ra1 and dec1 in a way that ra2 and dec2 will exactly map
83	onto ra3 and dec3, respectively. All angles are treated as radians.
84	Essentially rotates the events, so that they behave as if they were
85	originally incident on the source.
86
87	:param ra1: Event Right Ascension
88	:param dec1: Event Declination
89	:param ra2: True Event Right Ascension
90	:param dec2: True Event Declination
91	:param ra3: Source Right Ascension
92	:param dec3: Source Declination
93	:return: Returns new Right Ascensions and Declinations
94	"""
95	# Turns Right Ascension/Declination into Azimuth/Zenith for healpy
96	phi1 = ra1 - np.pi	3✔
97	zen1 = np.pi / 2.0 - dec1	3✔
98	phi2 = ra2 - np.pi	3✔
99	zen2 = np.pi / 2.0 - dec2	3✔
100	phi3 = ra3 - np.pi	3✔
101	zen3 = np.pi / 2.0 - dec3	3✔
102
103	# Rotate each ra1 and dec1 towards the pole?
104	x = np.array(	3✔
105	[
106	hp.rotator.rotateDirection(
107	hp.rotator.get_rotation_matrix((dp, -dz, 0.0))[0], z, p
108	)
109	for z, p, dz, dp in zip(zen1, phi1, zen2, phi2)
110	]
111	)
112
113	# Rotate all these vectors towards ra3, dec3 (source_path)
114	zen, phi = hp.rotator.rotateDirection(	3✔
115	np.dot(
116	hp.rotator.get_rotation_matrix((-phi3, 0, 0))[0],
117	hp.rotator.get_rotation_matrix((0, zen3, 0.0))[0],
118	),
119	x[:, 0],
120	x[:, 1],
121	)
122
123	dec = np.pi / 2.0 - zen	3✔
124
125	ra = phi + np.pi	3✔
126	return np.atleast_1d(ra), np.atleast_1d(dec)	3✔
127
128	def rotate_to_position(self, ev, ra, dec):	3✔
129	"""Modifies the events by reassigning the Right Ascension and
130	Declination of the events. Rotates the events, so that they are
131	distributed as if they originated from the source. Removes the
132	additional Monte Carlo information from sampled events, so that they
133	appear like regular data.
134
135	The fields removed are:
136	True Right Ascension,
137	True Declination,
138	True Energy,
139	OneWeight
140
141	:param ev: Events
142	:param ra: Source Right Ascension (radians)
143	:param dec: Source Declination (radians)
144	:return: Events (modified)
145	"""
146	names = ev.dtype.names	3✔
147
148	# Rotates the events to lie on the source
149	ev["ra"], rot_dec = self.rotate(	3✔
150	ev["ra"], np.arcsin(ev["sinDec"]), ev["trueRa"], ev["trueDec"], ra, dec
151	)
152
153	if "dec" in names:	3✔
154	ev["dec"] = rot_dec	3✔
155	ev["sinDec"] = np.sin(rot_dec)	3✔
156
157	# Deletes the Monte Carlo information from sampled events
158	non_mc = [	3✔
159	name for name in names if name not in ["trueRa", "trueDec", "trueE", "ow"]
160	]
161	ev = ev[non_mc].copy()	3✔
162
163	return ev	3✔
164
165
166	@SignalSpatialPDF.register_subclass("circular_gaussian")	3✔
167	class CircularGaussian(SignalSpatialPDF):	3✔
168	def simulate_distribution(self, source, data):	3✔
169	data["ra"] = np.pi + norm.rvs(size=len(data)) * data["sigma"]	3✔
170	data["dec"] = norm.rvs(size=len(data)) * data["sigma"]	3✔
171	data["sinDec"] = np.sin(data["dec"])	3✔
172	data = append_fields(	3✔
173	data,
174	["trueRa", "trueDec"],
175	[np.ones_like(data["dec"]) * np.pi, np.zeros_like(data["dec"])],
176	).copy()
177
178	data = self.rotate_to_position(data, source["ra_rad"], source["dec_rad"]).copy()	3✔
179
180	return data.copy()	3✔
181
182	@staticmethod	3✔
183	def signal_spatial(source, cut_data):	3✔
184	"""Calculates the angular distance between the source and the
185	coincident dataset. Uses a Gaussian PDF function, centered on the
186	source. Returns the value of the Gaussian at the given distances.
187
188	:param source: Single Source
189	:param cut_data: Subset of Dataset with coincident events
190	:return: Array of Spatial PDF values
191	"""
192	distance = angular_distance(	3✔
193	cut_data["ra"], cut_data["dec"], source["ra_rad"], source["dec_rad"]
194	)
195	space_term = (	3✔
196	1.0
197	/ (2.0 * np.pi * cut_data["sigma"] ** 2.0)
198	* np.exp(-0.5 * (distance / cut_data["sigma"]) ** 2.0)
199	)
200
201	return space_term	3✔
202
203
204	@SignalSpatialPDF.register_subclass("northern_tracks_kde")	3✔
205	class NorthernTracksKDE(SignalSpatialPDF):	3✔
206	"""Spatial PDF class for use with the KDE-smoothed MC PDFs introduced for the 10-yr NT analysis.
207	Current limitation: In the NT analysis, this PDF depends on the spectral index gamma . Here
208	a fixed gamma is used (either by setting 'spatial_pdf_data' to the location of the
209	4D-photospline-tables of the PDF and 'spatial_pdf_index' to the preferred gamma value, or
210	(more efficiently) setting 'spatial_pdf_data' directly to the corresponding 3D-spline table.
211	"""
212
213	KDEspline = None	3✔
214	SplineIs4D = False	3✔
215
216	def __init__(self, spatial_pdf_dict):	3✔
NEW 217	super().__init__(spatial_pdf_dict)	×
NEW 218	assert "spatial_pdf_data" in spatial_pdf_dict.keys() and os.path.exists(	×
219	spatial_pdf_dict["spatial_pdf_data"]
220	)
NEW 221	KDEfile = spatial_pdf_dict["spatial_pdf_data"]	×
222
NEW 223	NorthernTracksKDE.KDEspline = SplineTable(KDEfile)	×
NEW 224	if NorthernTracksKDE.KDEspline.ndim == 3:	×
NEW 225	NorthernTracksKDE.SplineIs4D = False	×
NEW 226	elif NorthernTracksKDE.KDEspline.ndim == 4:	×
NEW 227	NorthernTracksKDE.SplineIs4D = True	×
228	else:
NEW 229	raise RuntimeError(	×
230	f"{KDEfile} does not seem to be a valid photospline table for the PSF"
231	)
232
NEW 233	logger.info("Using KDE spatial PDF from file {0}.".format(KDEfile))	×
234
235	def _inverse_cdf(sigma, logE, gamma, npoints=100):	3✔
NEW 236	psi_range = np.linspace(0.001, 0.5, npoints)	×
NEW 237	d_psi = psi_range[1] - psi_range[0]	×
NEW 238	if NorthernTracksKDE.SplineIs4D:	×
NEW 239	psi_pdf = (	×
240	d_psi
241	/ (log(10) * psi_range)
242	* self.KDEspline.evaluate_simple(
243	[log10(sigma), logE, log10(psi_range), NorthernTracksKDE.gamma]
244	)
245	)
246	else:
NEW 247	psi_pdf = (	×
248	d_psi
249	/ (log(10) * psi_range)
250	* self.KDEspline.evaluate_simple([log10(sigma), logE, log10(psi_range)])
251	)
NEW 252	psi_cdf = np.insert(psi_pdf.cumsum(), 0, 0)	×
NEW 253	psi_range = np.insert(psi_range, 0, 0)	×
NEW 254	psi_cdf /= psi_cdf[-1]	×
NEW 255	psi_cdf_unique, unique_indices = np.unique(psi_cdf, return_index=True)	×
NEW 256	psi_range_unique = psi_range[unique_indices]	×
NEW 257	return interp1d(psi_cdf_unique, psi_range_unique, "cubic")	×
258
259	def simulate_distribution(self, source, data, gamma=2.0):	3✔
NEW 260	nevents = len(data)	×
NEW 261	phi = np.random.rand(nevents) * 2.0 * np.pi	×
NEW 262	distance = np.random.rand(nevents)	×
NEW 263	for _i in range(nevents):	×
NEW 264	event_icdf = self._inverse_cdf(data["sigma"][_i], data["logE"][_i], gamma)	×
NEW 265	distance = event_icdf(distance)	×
266
NEW 267	data["ra"] = np.pi + distance * np.cos(phi)	×
NEW 268	data["dec"] = distance * np.sin(phi)	×
NEW 269	data["sinDec"] = np.sin(data["dec"])	×
NEW 270	data = append_fields(	×
271	data,
272	["trueRa", "trueDec"],
273	[np.ones_like(data["dec"]) * np.pi, np.zeros_like(data["dec"])],
274	).copy()
275
NEW 276	data = self.rotate_to_position(data, source["ra_rad"], source["dec_rad"]).copy()	×
277
NEW 278	return data.copy()	×
279
280	@staticmethod	3✔
281	def signal_spatial(source, cut_data, gamma=2.0):	3✔
282	"""Calculates the angular distance between the source and the
283	coincident dataset. This class provides an interface for the KDE-smoothed MC PDF introduced for the 10yr NT analysis.
284	Returns the value of the PDF at the given distances between the source and the events.
285
286	:param source: Single Source
287	:param cut_data: Subset of Dataset with coincident events
288	:return: Array of Spatial PDF values
289	"""
290
291	# logger.debug(f"signal_spatial called with gamma={gamma}.")
292
NEW 293	distance = angular_distance(	×
294	cut_data["ra"], cut_data["dec"], source["ra_rad"], source["dec_rad"]
295	)
296
NEW 297	if NorthernTracksKDE.SplineIs4D:	×
NEW 298	space_term = NorthernTracksKDE.KDEspline.evaluate_simple(	×
299	[
300	np.log10(cut_data["sigma"]),
301	cut_data["logE"],
302	np.log10(distance),
303	gamma,
304	]
305	)
306	else:
NEW 307	space_term = NorthernTracksKDE.KDEspline.evaluate_simple(	×
308	[np.log10(cut_data["sigma"]), cut_data["logE"], np.log10(distance)]
309	)
310
NEW 311	space_term /= 2 * np.pi * np.log(10) * (distance**2)	×
312
NEW 313	return space_term	×
314
315
316	# ==============================================================================
317	# Background Spatial PDFs
318	# ==============================================================================
319
320
321	class BackgroundSpatialPDF:	3✔
322	subclasses: dict[str, object] = {}	3✔
323
324	def __init__(self, spatial_pdf_dict, season):	3✔
325	pass	3✔
326
327	@classmethod	3✔
328	def register_subclass(cls, bkg_spatial_name):	3✔
329	"""Adds a new subclass of BackgroundSpatialPDF, with class name equal to
330	"spatial_pdf_name".
331	"""
332
333	def decorator(subclass):	3✔
334	cls.subclasses[bkg_spatial_name] = subclass	3✔
335	return subclass	3✔
336
337	return decorator	3✔
338
339	@classmethod	3✔
340	def create(cls, s_pdf_dict, season):	3✔
341	try:	3✔
342	s_pdf_name = s_pdf_dict["bkg_spatial_pdf"]	3✔
343	except KeyError:	3✔
344	s_pdf_name = "zenith_spline"	3✔
345
346	if s_pdf_name not in cls.subclasses:	3✔
347	raise ValueError(	×
348	"Bad Background Spatial PDF name {0} \n"
349	"Available names are {1}".format(s_pdf_name, cls.subclasses)
350	)
351
352	return cls.subclasses[s_pdf_name](s_pdf_dict, season)	3✔
353
354	def background_spatial(self, events):	3✔
355	return np.ones(len(events))	×
356
357
358	@BackgroundSpatialPDF.register_subclass("uniform")	3✔
359	class UniformPDF(BackgroundSpatialPDF):	3✔
360	"""A highly-simplified spatial PDF in which events are distributed
361	uniformly over the celestial sphere.
362	"""
363
364	def background_spatial(self, events):	3✔
365	space_term = (1.0 / (4.0 * np.pi)) * np.ones(len(events))	×
366	return space_term	×
367
368
369	@BackgroundSpatialPDF.register_subclass("uniform_solid_angle")	3✔
370	class UniformSolidAngle(BackgroundSpatialPDF):	3✔
371	"""Generic class for a background PDF that is uniform over some fixed
372	area, and 0 otherwise. Requires an argument to be passed in the
373	'spatial_pdf_dict', with key 'background_solid_angle'. In the limit of a
374	solid angle of 4 pi, this becomes identical to the UniformPDF class.
375	"""
376
377	def __init__(self, spatial_pdf_dict, season):	3✔
378	BackgroundSpatialPDF.__init__(self, spatial_pdf_dict, season)	×
379
380	try:	×
381	self.solid_angle = spatial_pdf_dict["background_solid_angle"]	×
382	except KeyError:	×
383	raise KeyError(	×
384	"No solid angle passed to UniformSolidAngle class.\n"
385	"Please include an entry 'background_solid_angle' "
386	"in the 'spatial_pdf_dict'."
387	)
388
389	if self.solid_angle > 4 * np.pi:	×
390	raise ValueError(	×
391	"Solid angle {0} was provided, but this is "
392	"larger than 4pi. Please provide a valid solid "
393	"angle."
394	)
395
396	def background_spatial(self, events):	3✔
397	space_term = (1.0 / self.solid_angle) * np.ones(len(events))	×
398	return space_term	×
399
400
401	@BackgroundSpatialPDF.register_subclass("zenith_spline")	3✔
402	class ZenithSpline(BackgroundSpatialPDF):	3✔
403	"""A 1D background spatial PDF, in which the background is assumed to be
404	uniform in azimuth, but varying as a function of zenith. A
405	spline is used to parameterise this distribution.
406	"""
407
408	def __init__(self, spatial_pdf_dict, season):	3✔
409	BackgroundSpatialPDF.__init__(self, spatial_pdf_dict, season)	3✔
410	self.bkg_f = self.create_background_function(season)	3✔
411
412	@staticmethod	3✔
413	def create_background_function(season):	3✔
414	# Checks if background spatial spline has been created
415
416	if not os.path.isfile(bkg_spline_path(season)):	3✔
417	season.make_background_spatial()	3✔
418
419	return load_bkg_spatial_spline(season)	3✔
420
421	def background_spatial(self, events):	3✔
422	space_term = (1.0 / (2.0 * np.pi)) * np.exp(self.bkg_f(events["sinDec"]))	3✔
423
424	return space_term	3✔

icecube / flarestack / 7049225217

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous