diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index f35aa6a8f0..845d299e2c 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -21,6 +21,19 @@ jobs: ruff check . ruff format --check . + typecheck: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v7 + - name: Set up Python + uses: actions/setup-python@v7 + with: + python-version: "3.11" + - name: Install pyright + run: pip install ".[dev,extras]" + - name: Run pyright + run: pyright + pytest: runs-on: ubuntu-latest strategy: diff --git a/pyproject.toml b/pyproject.toml index d713bb3bbb..a29ace71c8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -43,6 +43,7 @@ Issues = "https://github.com/icecube/skyllh/issues" dev = [ "pre-commit", "pytest>=9.0.0", + "pyright", "ruff>=0.16", ] docs = [ @@ -95,3 +96,13 @@ quote-style = "single" minversion = "9.0" addopts = ["-ra"] testpaths = ["tests"] + +[tool.pyright] +pythonVersion = "3.11" +typeCheckingMode = "standard" +reportUnnecessaryTypeIgnoreComment = "warning" + +include = [ + "skyllh", + "tests", +] diff --git a/skyllh/__init__.py b/skyllh/__init__.py index ee5b905164..67bf192c85 100644 --- a/skyllh/__init__.py +++ b/skyllh/__init__.py @@ -1,5 +1,9 @@ import logging import multiprocessing as mp +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from skyllh.datasets import create_datasets __all__ = [ 'create_datasets', diff --git a/skyllh/analyses/i3/publicdata_ps/aeff.py b/skyllh/analyses/i3/publicdata_ps/aeff.py index 2d185cf659..98d1bc85d2 100644 --- a/skyllh/analyses/i3/publicdata_ps/aeff.py +++ b/skyllh/analyses/i3/publicdata_ps/aeff.py @@ -1,3 +1,6 @@ +from collections.abc import Sequence +from typing import cast + import numpy as np from scipy import interpolate @@ -13,27 +16,29 @@ ) -def load_effective_area_array(pathfilenames): +def load_effective_area_array( + pathfilenames: str | Sequence[str], +) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]: """Loads the (nbins_decnu, nbins_log10enu)-shaped 2D effective area array from the given data file. Parameters ---------- - pathfilename : str | list of str - The file name of the data file. + pathfilenames + The file name(s) of the data file(s). Returns ------- - aeff_decnu_log10enu : (nbins_decnu, nbins_log10enu)-shaped 2D ndarray + aeff_decnu_log10enu The ndarray holding the effective area for each (dec_nu,log10(E_nu/GeV)) bin. - decnu_binedges_lower : (nbins_decnu,)-shaped ndarray + decnu_binedges_lower The ndarray holding the lower bin edges of the dec_nu axis. - decnu_binedges_upper : (nbins_decnu,)-shaped ndarray + decnu_binedges_upper The ndarray holding the upper bin edges of the dec_nu axis. - log10_enu_binedges_lower : (nbins_log10enu,)-shaped ndarray + log10_enu_binedges_lower The ndarray holding the lower bin edges of the log10(E_nu/GeV) axis. - log10_enu_binedges_upper : (nbins_log10enu,)-shaped ndarray + log10_enu_binedges_upper The ndarray holding the upper bin edges of the log10(E_nu/GeV) axis. """ loader = create_FileLoader(pathfilenames=pathfilenames) @@ -95,25 +100,32 @@ class PDAeff: the public data. """ - def __init__(self, pathfilenames, src_dec=None, min_log10enu=None, max_log10enu=None, **kwargs): + def __init__( + self, + pathfilenames: str | list[str], + src_dec: float | None = None, + min_log10enu: float | None = None, + max_log10enu: float | None = None, + **kwargs, + ): """Creates an effective area instance by loading the effective area data from the given file. Parameters ---------- - pathfilenames : str | list of str - The path file names of the effective area data file(s) which should + pathfilenames + The path file name(s) of the effective area data file(s) which should be used for this public data effective area instance. - src_dec : float | None + src_dec The source declination in radians for which detection probabilities should get pre-calculated using the ``get_detection_prob_for_decnu`` method. - min_log10enu : float | None + min_log10enu The minimum log10(E_nu/GeV) value that should be used for calculating the detection probability. If None, the lowest available neutrino energy bin edge of the effective area is used. - max_log10enu : float | None + max_log10enu The maximum log10(E_nu/GeV) value that should be used for calculating the detection probability. If None, the highest available neutrino energy bin edge of the @@ -185,7 +197,7 @@ def decnu_bincenters(self): return get_bincenters_from_binedges(self._decnu_binedges) @property - def n_decnu_bins(self): + def n_decnu_bins(self) -> int: """(read-only) The number of bins of the neutrino declination axis.""" return len(self._decnu_binedges) - 1 @@ -218,7 +230,7 @@ def log10_enu_bincenters(self): return get_bincenters_from_binedges(self._log10_enu_binedges) @property - def n_log10_enu_bins(self): + def n_log10_enu_bins(self) -> int: """(read-only) The number of bins of the log10 neutrino energy axis.""" return len(self._log10_enu_binedges) - 1 @@ -229,31 +241,31 @@ def aeff_decnu_log10enu(self): """ return self._aeff_decnu_log10enu - def create_sin_decnu_log10_enu_spline(self): + def create_sin_decnu_log10_enu_spline(self) -> FctSpline2D: """DEPRECATED! Creates a FctSpline2D object representing a 2D spline of the effective area in sin(dec_nu)-log10(E_nu/GeV)-space. Returns ------- - spl : FctSpline2D instance + spl The FctSpline2D instance representing a spline in the sin(dec_nu)-log10(E_nu/GeV)-space. """ spl = FctSpline2D(self._aeff_decnu_log10enu, self.sin_decnu_binedges, self.log10_enu_binedges) return spl - def get_aeff_for_decnu(self, decnu): + def get_aeff_for_decnu(self, decnu: float) -> np.ndarray: """Retrieves the effective area as function of log10_enu. Parameters ---------- - decnu : float + decnu The true neutrino declination. Returns ------- - aeff : (n,)-shaped numpy ndarray + aeff The effective area in cm^2 for the given true neutrino declination as a function of log10 true neutrino energy. """ @@ -263,26 +275,33 @@ def get_aeff_for_decnu(self, decnu): return aeff - def get_detection_prob_for_decnu(self, decnu, enu_min, enu_max, enu_range_min, enu_range_max): + def get_detection_prob_for_decnu( + self, + decnu: float, + enu_min: float | np.ndarray, + enu_max: float | np.ndarray, + enu_range_min: float, + enu_range_max: float, + ) -> np.ndarray: """Calculates the detection probability for given true neutrino energy ranges for a given neutrino declination. Parameters ---------- - decnu : float + decnu The neutrino declination in radians. - enu_min : float | ndarray of float + enu_min The minimum energy in GeV. - enu_max : float | ndarray of float + enu_max The maximum energy in GeV. - enu_range_min : float + enu_range_min The minimum energy in GeV of the entire energy range. - enu_range_max : float + enu_range_max The maximum energy in GeV of the entire energy range. Returns ------- - det_prob : ndarray of float + det_prob The neutrino energy detection probabilities for the given true enegry ranges. """ @@ -323,13 +342,13 @@ def get_detection_prob_for_decnu(self, decnu, enu_min, enu_max, enu_range_min, e spl = interpolate.splrep(x, y, k=1, s=0) - norm = interpolate.splint(enu_range_min, enu_range_max, spl) + norm = cast(float, interpolate.splint(enu_range_min, enu_range_max, spl)) enu_min = np.atleast_1d(enu_min) enu_max = np.atleast_1d(enu_max) det_prob = np.empty((len(enu_min),), dtype=np.double) for i in range(len(enu_min)): - det_prob[i] = interpolate.splint(enu_min[i], enu_max[i], spl) / norm + det_prob[i] = cast(float, interpolate.splint(enu_min[i], enu_max[i], spl)) / norm return det_prob diff --git a/skyllh/analyses/i3/publicdata_ps/backgroundpdf.py b/skyllh/analyses/i3/publicdata_ps/backgroundpdf.py index 320f12bd09..0023d0293f 100644 --- a/skyllh/analyses/i3/publicdata_ps/backgroundpdf.py +++ b/skyllh/analyses/i3/publicdata_ps/backgroundpdf.py @@ -30,9 +30,8 @@ from skyllh.core.storage import ( DataFieldRecordArray, ) -from skyllh.core.timing import ( - TaskTimer, -) +from skyllh.core.timing import TaskTimer, TimeLord +from skyllh.core.trialdata import TrialDataManager class PDBackgroundI3EnergyPDF(EnergyPDF, IsBackgroundPDF, UsesBinning): @@ -51,46 +50,45 @@ class PDBackgroundI3EnergyPDF(EnergyPDF, IsBackgroundPDF, UsesBinning): def __init__( self, - data_logE, - data_sinDec, - data_mcweight, - data_physicsweight, - logE_binning, - sinDec_binning, - smoothing_filter, - kde_smoothing=False, + data_logE: np.ndarray, + data_sinDec: np.ndarray, + data_mcweight: np.ndarray, + data_physicsweight: np.ndarray, + logE_binning: BinningDefinition, + sinDec_binning: BinningDefinition, + smoothing_filter: SmoothingFilter | None = None, + kde_smoothing: bool = False, **kwargs, ): """Creates a new IceCube energy PDF object for the public data. Parameters ---------- - data_logE : instance of ndarray + data_logE The 1d ndarray holding the log10(E) values of the events. - data_sinDec : instance of ndarray + data_sinDec The 1d ndarray holding the sin(dec) values of the events. - data_mcweight : instance of ndarray + data_mcweight The 1d ndarray holding the monte-carlo weights of the events. The final data weight will be the product of data_mcweight and data_physicsweight. - data_physicsweight : instance of ndarray + data_physicsweight The 1d ndarray holding the physics weights of the events. The final data weight will be the product of data_mcweight and data_physicsweight. - logE_binning : instance of BinningDefinition + logE_binning The binning definition for the log10(E) axis. - sinDec_binning : instance of BinningDefinition + sinDec_binning The binning definition for the sin(declination) axis. - smoothing_filter : instance of SmoothingFilter | None + smoothing_filter The smoothing filter to use for smoothing the energy histogram. If None, no smoothing will be applied. - kde_smoothing : bool + kde_smoothing Deprecated: use of ``kde_smoothing=True`` is deprecated and will be removed in a future version. Apply a kde smoothing to the energy pdf for each bin in sin(dec). This is useful for signal injections, because it ensures that the background is not zero when injecting high energy events. - Default: False. """ super().__init__(pmm=None, **kwargs) @@ -264,21 +262,21 @@ def initialize_for_new_trial(self, tdm, tl=None, **kwargs): with TaskTimer(tl, 'Evaluating logE-sinDec histogram.'): self._pd = self._pdf_spline(tdm['log_energy'], tdm['sin_dec'], grid=False) - def assert_is_valid_for_trial_data(self, tdm, tl=None): + def assert_is_valid_for_trial_data(self, tdm: TrialDataManager, tl: TimeLord | None = None, **kwargs) -> None: """Checks if this energy PDF covers the entire value range of the trail data events. Parameters ---------- - tdm : instance of TrialDataManager + tdm The instance of TrialDataManager holding the trial data events. The following data fields need to exist: - log_energy : float + log_energy The base-10 logarithm of the reconstructed energy value. - sin_dec : float + sin_dec The sine of the declination value of the event. - tl : instance of TimeLord | None + tl The optional instance of TimeLord to measure timing information. Raises @@ -317,34 +315,36 @@ def assert_is_valid_for_trial_data(self, tdm, tl=None): f'{sindecmu_axis.vmax:g}!' ) - def get_pd(self, tdm, params_recarray=None, tl=None): + def get_pd( + self, tdm: TrialDataManager, params_recarray: np.ndarray | None = None, tl: TimeLord | None = None + ) -> tuple[np.ndarray, dict]: """Calculates the energy probability density (in 1/log10(E/GeV)) of each trial data event. Parameters ---------- - tdm : instance of TrialDataManager + tdm The TrialDataManager instance holding the data events for which the probability should be calculated for. The following data fields must exist: - log_energy : float + log_energy The base-10 logarithm of the energy value of the event. - sin_dec : float + sin_dec The sin(declination) value of the event. - params_recarray : None + params_recarray Unused interface parameter. - tl : instance of TimeLord | None + tl The optional TimeLord instance that should be used to measure timing information. Returns ------- - pd : instance of ndarray + pd The (N_selected_events,)-shaped numpy ndarray holding the energy probability density value for each trial data event. - grads : dict + grads The dictionary holding the gradients of the probability density w.r.t. each global fit parameter. By definition this PDF does not depend on any fit parameter, hence, this dictionary is empty. @@ -359,36 +359,43 @@ class PDDataBackgroundI3EnergyPDF(PDBackgroundI3EnergyPDF): the experimental data of the public data. """ - def __init__(self, data_exp, logE_binning, sinDec_binning, smoothing_filter=None, kde_smoothing=False, **kwargs): + def __init__( + self, + data_exp: DataFieldRecordArray, + logE_binning: BinningDefinition, + sinDec_binning: BinningDefinition, + smoothing_filter: SmoothingFilter | None = None, + kde_smoothing: bool = False, + **kwargs, + ): """Constructs a new IceCube energy background PDF from experimental data. Parameters ---------- - data_exp : instance of DataFieldRecordArray + data_exp The array holding the experimental data. The following data fields must exist: - log_energy : float + log_energy The base-10 logarithm of the reconstructed energy value of the data event. - sin_dec : float + sin_dec The sine of the reconstructed declination of the data event. - logE_binning : instance of BinningDefinition + logE_binning The binning definition for the binning in log10(E). - sinDec_binning : instance of BinningDefinition + sinDec_binning The binning definition for the sin(declination). - smoothing_filter : instance of SmoothingFilter | None + smoothing_filter The smoothing filter to use for smoothing the energy histogram. If None, no smoothing will be applied. - kde_smoothing : bool + kde_smoothing Deprecated: use of ``kde_smoothing=True`` is deprecated and will be removed in a future version. Apply a kde smoothing to the energy pdf for each bin in sin(dec). This is useful for signal injections, because it ensures that the background is not zero when injecting high energy events. - Default: False. """ if not isinstance(data_exp, DataFieldRecordArray): raise TypeError( @@ -422,20 +429,26 @@ class PDMCBackgroundI3EnergyPDF(EnergyPDF, IsBackgroundPDF, UsesBinning): data and a monte-carlo background flux model. """ - def __init__(self, pdf_log10emu_sindecmu, log10emu_binning, sindecmu_binning, **kwargs): + def __init__( + self, + pdf_log10emu_sindecmu: np.ndarray, + log10emu_binning: BinningDefinition, + sindecmu_binning: BinningDefinition, + **kwargs, + ): """Constructs a new background energy PDF with the given PDF data and binning. Parameters ---------- - pdf_log10emu_sindecmu : instance of numpy ndarray + pdf_log10emu_sindecmu The (n_log10emu, n_sindecmu)-shaped 2D numpy ndarray holding the PDF values in unit 1/log10(E_mu/GeV). A copy of this data will be created and held within this class instance. - log10emu_binning : instance of BinningDefinition + log10emu_binning The binning definition for the binning in log10(E_mu/GeV). - sindecmu_binning : instance of BinningDefinition + sindecmu_binning The binning definition for the binning in sin(dec_mu). """ if not isinstance(pdf_log10emu_sindecmu, np.ndarray): @@ -467,21 +480,21 @@ def __init__(self, pdf_log10emu_sindecmu, log10emu_binning, sindecmu_binning, ** self.add_binning(log10emu_binning, name='log_energy') self.add_binning(sindecmu_binning, name='sin_dec') - def assert_is_valid_for_trial_data(self, tdm, tl=None): + def assert_is_valid_for_trial_data(self, tdm: TrialDataManager, tl: TimeLord | None = None, **kwargs) -> None: """Checks if this energy PDF covers the entire value range of the trail data events. Parameters ---------- - tdm : instance of TrialDataManager + tdm The instance of TrialDataManager holding the trial data events. The following data fields need to exist: - log_energy : float + log_energy The base-10 logarithm of the reconstructed energy value. - sin_dec : float + sin_dec The sine of the declination value of the event. - tl : instance of TimeLord | None + tl The optional instance of TimeLord to measure timing information. Raises @@ -491,7 +504,7 @@ def assert_is_valid_for_trial_data(self, tdm, tl=None): PDF. """ log10emu = tdm['log_energy'] - log10emu_axis = self.get_axis(0) + log10emu_axis = self._axes[0] if np.min(log10emu) < log10emu_axis.vmin: raise ValueError( f'The minimum log10emu value {np.min(log10emu):g} of the trial ' @@ -502,11 +515,11 @@ def assert_is_valid_for_trial_data(self, tdm, tl=None): raise ValueError( f'The maximum log10emu value {np.max(log10emu):g} of the trial ' 'data is larger than the maximum value of the PDF ' - f'{log10emu_axis.vmax}:g!' + f'{log10emu_axis.vmax:g}!' ) sindecmu = tdm['sin_dec'] - sindecmu_axis = self.get_axis(1) + sindecmu_axis = self._axes[1] if np.min(sindecmu) < sindecmu_axis.vmin: raise ValueError( f'The minimum sindecmu value {np.min(sindecmu):g} of the trial ' @@ -520,32 +533,34 @@ def assert_is_valid_for_trial_data(self, tdm, tl=None): f'{sindecmu_axis.vmax:g}!' ) - def get_pd(self, tdm, params_recarray=None, tl=None): + def get_pd( + self, tdm: TrialDataManager, params_recarray: np.ndarray | None = None, tl: TimeLord | None = None + ) -> tuple[np.ndarray, dict]: """Gets the probability density for the given trial data events. Parameters ---------- - tdm : instance of TrialDataManager + tdm The instance of TrialDataManager holding the trial data events. The following data fields need to exist: - log_energy : float + log_energy The base-10 logarithm of the reconstructed energy value. - sin_dec : float + sin_dec The sine of the declination value of the event. - params_recarray : None + params_recarray Unused interface argument. - tl : instance of TimeLord | None + tl The optional instance of TimeLord that should be used to measure timing information. Returns ------- - pd : instance of ndarray + pd The (N_selected_events,)-shaped numpy ndarray holding the probability density value for each event. - grads : dict + grads The dictionary holding the gradients of the probability density w.r.t. each global fit parameter. By definition this PDF does not depend on any fit parameter, hence, this dictionary is empty. diff --git a/skyllh/analyses/i3/publicdata_ps/bkg_flux.py b/skyllh/analyses/i3/publicdata_ps/bkg_flux.py index 0598a63403..df15067787 100644 --- a/skyllh/analyses/i3/publicdata_ps/bkg_flux.py +++ b/skyllh/analyses/i3/publicdata_ps/bkg_flux.py @@ -11,61 +11,63 @@ ) -def get_dOmega(dec_min, dec_max): +def get_dOmega(dec_min: float | np.ndarray, dec_max: float | np.ndarray) -> float | np.ndarray: """Calculates the solid angle given two declination angles. Parameters ---------- - dec_min : float | array of float + dec_min The smaller declination angle. - dec_max : float | array of float + dec_max The larger declination angle. Returns ------- - solidangle : float | array of float + solidangle The solid angle corresponding to the two given declination angles. """ return 2 * np.pi * (np.sin(dec_max) - np.sin(dec_min)) -def southpole_zen2dec(zen): +def southpole_zen2dec(zen: np.ndarray) -> np.ndarray: """Converts zenith angles at the South Pole to declination angles. Parameters ---------- - zen : (n,)-shaped 1d numpy ndarray + zen The numpy ndarray holding the zenith angle values in radians. Returns ------- - dec : (n,)-shaped 1d numpy ndarray + dec The numpy ndarray holding the declination angle values in radians. """ dec = zen - np.pi / 2 return dec -def get_flux_atmo_decnu_log10enu(flux_pathfilename, log10_enu_max=9): +def get_flux_atmo_decnu_log10enu( + flux_pathfilename: str, log10_enu_max: float = 9 +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: """Constructs the atmospheric flux map function f_atmo(log10(E_nu/GeV),dec_nu) in unit 1/(GeV cm^2 sr s). Parameters ---------- - flux_pathfilename : str + flux_pathfilename The pathfilename of the file containing the MCEq fluxes. - log10_enu_max : float + log10_enu_max The log10(E/GeV) value of the maximum neutrino energy to be considered. Returns ------- - flux_atmo : (n_dec, n_e_grid)-shaped 2D numpy ndarray - The numpy ndarray holding the the atmospheric neutrino flux function in + flux_atmo + The (n_dec, n_e_grid)-shaped 2D numpy ndarray holding the the atmospheric neutrino flux function in unit 1/(GeV cm^2 sr s). - decnu_binedges : (n_decnu+1,)-shaped 1D numpy ndarray - The numpy ndarray holding the dec_nu bin edges. - log10_enu_binedges : (n_enu+1,)-shaped 1D numpy ndarray - The numpy ndarray holding the neutrino energy bin edges in log10. + decnu_binedges + The (n_decnu+1,)-shaped 1D numpy ndarray holding the dec_nu bin edges. + log10_enu_binedges + The (n_enu+1,)-shaped 1D numpy ndarray holding the neutrino energy bin edges in log10. """ with open(flux_pathfilename, 'rb') as f: ((e_grid, zenith_angle_binedges), flux_def) = pickle.load(f) @@ -104,7 +106,7 @@ def get_flux_atmo_decnu_log10enu(flux_pathfilename, log10_enu_max=9): return (f_atmo, decnu_binedges, log10_enu_binedges) -def get_flux_astro_decnu_log10enu(decnu_binedges, log10_enu_binedges): +def get_flux_astro_decnu_log10enu(decnu_binedges: np.ndarray, log10_enu_binedges: np.ndarray) -> np.ndarray: """Constructs the astrophysical neutrino flux function f_astro(log10(E_nu/GeV),dec_nu) in unit 1/(GeV cm^2 sr s). @@ -112,16 +114,16 @@ def get_flux_astro_decnu_log10enu(decnu_binedges, log10_enu_binedges): Parameters ---------- - decnu_binedges : (n_decnu+1,)-shaped 1D numpy ndarray - The numpy ndarray holding the dec_nu bin edges. - log10_enu_binedges : (n_enu+1,)-shaped 1D numpy ndarray - The numpy ndarray holding the log10 values of the neutrino energy bin + decnu_binedges + The (n_decnu+1,)-shaped 1D numpy ndarray holding the dec_nu bin edges. + log10_enu_binedges + The (n_enu+1,)-shaped 1D numpy ndarray holding the log10 values of the neutrino energy bin edges in GeV. Returns ------- - f_astro : (n_decnu, n_log10enu)-shaped 2D numpy ndarray - The numpy ndarray holding the astrophysical flux values in unit + f_astro + The (n_decnu, n_log10enu)-shaped 2D numpy ndarray holding the astrophysical flux values in unit 1/(GeV cm^2 sr s). References @@ -141,25 +143,27 @@ def get_flux_astro_decnu_log10enu(decnu_binedges, log10_enu_binedges): return f_astro -def convert_flux_bkg_to_pdf_bkg(f_bkg, decnu_binedges, log10_enu_binedges): +def convert_flux_bkg_to_pdf_bkg( + f_bkg: np.ndarray, decnu_binedges: np.ndarray, log10_enu_binedges: np.ndarray +) -> np.ndarray: """Converts the given background flux function f_bkg into a background flux PDF in unit 1/(log10(E/GeV) rad). Parameters ---------- - f_bkg : (n_decnu, n_enu)-shaped 2D numpy ndarray - The numpy ndarray holding the background flux values in unit + f_bkg + The (n_decnu, n_enu)-shaped 2D numpy ndarray holding the background flux values in unit 1/(GeV cm^2 s sr). - decnu_binedges : (n_decnu+1,)-shaped 1D numpy ndarray - The numpy ndarray holding the dec_nu bin edges in radians. - log10_enu_binedges : (n_enu+1,)-shaped 1D numpy ndarray - The numpy ndarray holding the log10 values of the neutrino energy bin + decnu_binedges + The (n_decnu+1,)-shaped 1D numpy ndarray holding the dec_nu bin edges in radians. + log10_enu_binedges + The (n_enu+1,)-shaped 1D numpy ndarray holding the log10 values of the neutrino energy bin edges in GeV. Returns ------- - p_bkg : (n_decnu, n_enu)-shaped 2D numpy ndarray - The numpy ndarray holding the background flux pdf values. + p_bkg + The (n_decnu, n_enu)-shaped 2D numpy ndarray holding the background flux pdf values. """ d_decnu = np.diff(decnu_binedges) d_log10_enu = np.diff(log10_enu_binedges) @@ -174,26 +178,28 @@ def convert_flux_bkg_to_pdf_bkg(f_bkg, decnu_binedges, log10_enu_binedges): return p_bkg -def get_pd_atmo_decnu_Enu(flux_pathfilename, log10_true_e_max=9): +def get_pd_atmo_decnu_Enu( + flux_pathfilename: str, log10_true_e_max: float = 9 +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: """Constructs the atmospheric neutrino PDF p_atmo(E_nu,dec_nu) in unit 1/(GeV rad). Parameters ---------- - flux_pathfilename : str + flux_pathfilename The pathfilename of the file containing the MCEq flux. - log10_true_e_max : float + log10_true_e_max The log10(E/GeV) value of the maximum true energy to be considered. Returns ------- - pd_atmo : (n_dec, n_e_grid)-shaped 2D numpy ndarray - The numpy ndarray holding the the atmospheric neutrino PDF in unit + pd_atmo + The (n_decnu, n_e_grid)-shaped 2D numpy ndarray holding the the atmospheric neutrino PDF in unit 1/(GeV rad). - decnu_binedges : (n_decnu+1,)-shaped 1D numpy ndarray - The numpy ndarray holding the dec_nu bin edges. - log10_e_grid_edges : (n_e_grid+1,)-shaped 1D numpy ndarray - The numpy ndarray holding the energy bin edges in log10. + decnu_binedges + The (n_decnu+1,)-shaped 1D numpy ndarray holding the dec_nu bin edges. + log10_e_grid_edges + The (n_e_grid+1,)-shaped 1D numpy ndarray holding the energy bin edges in log10. """ with open(flux_pathfilename, 'rb') as f: ((e_grid, zenith_angle_binedges), flux_def) = pickle.load(f) @@ -242,23 +248,23 @@ def get_pd_atmo_decnu_Enu(flux_pathfilename, log10_true_e_max=9): return (pd_atmo, decnu_binedges, log10_e_grid_edges) -def get_pd_atmo_E_nu_sin_dec_nu(flux_pathfilename): +def get_pd_atmo_E_nu_sin_dec_nu(flux_pathfilename: str) -> tuple[np.ndarray, np.ndarray, np.ndarray]: """Constructs the atmospheric energy PDF p_atmo(E_nu|sin(dec_nu)) in unit 1/GeV. Parameters ---------- - flux_pathfilename : str + flux_pathfilename The pathfilename of the file containing the MCEq flux. Returns ------- - pd_atmo : (n_sin_dec, n_e_grid)-shaped 2D numpy ndarray - The numpy ndarray holding the the atmospheric energy PDF in unit 1/GeV. - sin_dec_binedges : numpy ndarray + pd_atmo + The (n_sin_dec, n_e_grid)-shaped 2D numpy ndarray holding the the atmospheric energy PDF in unit 1/GeV. + sin_dec_binedges The (n_sin_dec+1,)-shaped 1D numpy ndarray holding the sin(dec) bin edges. - log10_e_grid_edges : numpy ndarray + log10_e_grid_edges The (n_e_grid+1,)-shaped 1D numpy ndarray holding the energy bin edges in log10. """ @@ -307,23 +313,23 @@ def get_pd_atmo_E_nu_sin_dec_nu(flux_pathfilename): return (pd_atmo, sin_dec_binedges, log10_e_grid_edges) -def get_pd_astro_E_nu_sin_dec_nu(sin_dec_binedges, log10_e_grid_edges): +def get_pd_astro_E_nu_sin_dec_nu(sin_dec_binedges: np.ndarray, log10_e_grid_edges: np.ndarray) -> np.ndarray: """Constructs the astrophysical energy PDF p_astro(E_nu|sin(dec_nu)) in unit 1/GeV. It uses the best fit from the IceCube publication [1]. Parameters ---------- - sin_dec_binedges : (n_sin_dec+1,)-shaped 1D numpy ndarray - The numpy ndarray holding the sin(dec) bin edges. - log10_e_grid_edges : (n_e_grid+1,)-shaped 1D numpy ndarray - The numpy ndarray holding the log10 values of the energy bin edges in + sin_dec_binedges + The (n_sin_dec+1,)-shaped 1D numpy ndarray holding the sin(dec) bin edges. + log10_e_grid_edges + The (n_e_grid+1,)-shaped 1D numpy ndarray holding the log10 values of the energy bin edges in GeV of the energy grid. Returns ------- - pd_astro : (n_sin_dec, n_e_grid)-shaped 2D numpy ndarray - The numpy ndarray holding the energy probability density values + pd_astro + The (n_sin_dec, n_e_grid)-shaped 2D numpy ndarray holding the energy probability density values p(E_nu|sin_dec_nu) in unit 1/GeV. References @@ -350,26 +356,26 @@ def get_pd_astro_E_nu_sin_dec_nu(sin_dec_binedges, log10_e_grid_edges): return pd_astro -def get_pd_bkg_E_nu_sin_dec_nu(pd_atmo, pd_astro, log10_e_grid_edges): +def get_pd_bkg_E_nu_sin_dec_nu(pd_atmo: np.ndarray, pd_astro: np.ndarray, log10_e_grid_edges: np.ndarray) -> np.ndarray: """Constructs the total background flux probability density p_bkg(E_nu|sin(dec_nu)) in unit 1/GeV. Parameters ---------- - pd_atmo : (n_sin_dec, n_e_grid)-shaped 2D numpy ndarray - The numpy ndarray holding the probability density values + pd_atmo + The (n_sin_dec, n_e_grid)-shaped 2D numpy ndarray holding the probability density values p(E_nu|sin(dec_nu)) in 1/GeV of the atmospheric flux. - pd_astro : (n_sin_dec, n_e_grid)-shaped 2D numpy ndarray - The numpy ndarray holding the probability density values + pd_astro + The (n_sin_dec, n_e_grid)-shaped 2D numpy ndarray holding the probability density values p(E_nu|sin(dec_nu)) in 1/GeV of the astrophysical flux. - log10_e_grid_edges : (n_e_grid+1,)-shaped numpy ndarray - The numpy ndarray holding the log10 values of the energy grid bin edges + log10_e_grid_edges + The (n_e_grid+1,)-shaped 1D numpy ndarray holding the log10 values of the energy grid bin edges in GeV. Returns ------- - pd_bkg : (n_sin_dec, n_e_grid)-shaped 2D numpy ndarray - The numpy ndarray holding total background probability density values + pd_bkg + The (n_sin_dec, n_e_grid)-shaped 2D numpy ndarray holding total background probability density values p_bkg(E_nu|sin(dec_nu)) in unit 1/GeV. """ pd_bkg = pd_atmo + pd_astro diff --git a/skyllh/analyses/i3/publicdata_ps/detsigyield.py b/skyllh/analyses/i3/publicdata_ps/detsigyield.py index 8a1ba25ade..f2284bb6e1 100644 --- a/skyllh/analyses/i3/publicdata_ps/detsigyield.py +++ b/skyllh/analyses/i3/publicdata_ps/detsigyield.py @@ -11,16 +11,20 @@ from skyllh.core.binning import ( BinningDefinition, ) +from skyllh.core.dataset import Dataset, DatasetData from skyllh.core.flux_model import ( FactorizedFluxModel, ) from skyllh.core.livetime import ( Livetime, ) +from skyllh.core.parameters import ParameterGrid +from skyllh.core.progressbar import ProgressBar from skyllh.core.py import ( classname, issequence, ) +from skyllh.core.source_hypo_grouping import SourceHypoGroup from skyllh.i3.detsigyield import ( SingleParamFluxPointLikeSourceI3DetSigYield, SingleParamFluxPointLikeSourceI3DetSigYieldBuilder, @@ -48,10 +52,10 @@ class PDSingleParamFluxPointLikeSourceI3DetSigYieldBuilder( def __init__( self, - param_grid, - spline_order_sinDec=2, - spline_order_param=2, - ncpu=None, + param_grid: ParameterGrid, + spline_order_sinDec: int = 2, + spline_order_param: int = 2, + ncpu: int | None = None, **kwargs, ): """Creates a new IceCube detector signal yield builder instance for @@ -61,18 +65,16 @@ def __init__( Parameters ---------- - param_grid : instance of ParameterGrid + param_grid The instance of ParameterGrid which defines the grid of parameter values. - spline_order_sinDec : int + spline_order_sinDec The order of the spline function for the logarithmic values of the detector signal yield along the sin(dec) axis. - The default is 2. - spline_order_param : int + spline_order_param The order of the spline function for the logarithmic values of the detector signal yield along the parameter axis. - The default is 2. - ncpu : int | None + ncpu The number of CPUs to utilize. If set to ``None``, global setting will take place. """ @@ -85,11 +87,11 @@ def __init__( **kwargs, ) - def assert_types_of_construct_detsigyield_arguments(self, shgs, **kwargs): + def assert_types_of_construct_detsigyield_arguments(self, dataset, data, shgs, ppbar, **kwargs): """Checks the correct types of the arguments for the ``construct_detsigyield`` method. """ - super().assert_types_of_construct_detsigyield_arguments(shgs=shgs, **kwargs) + super().assert_types_of_construct_detsigyield_arguments(dataset, data, shgs, ppbar, **kwargs) if not issequence(shgs): shgs = [shgs] @@ -101,27 +103,29 @@ def assert_types_of_construct_detsigyield_arguments(self, shgs, **kwargs): f'Its current type is {classname(shg.fluxmodel)}!' ) - def construct_detsigyield(self, dataset, data, shg, ppbar=None): + def construct_detsigyield( + self, dataset: Dataset, data: DatasetData, shg: SourceHypoGroup, ppbar: ProgressBar | None = None + ) -> SingleParamFluxPointLikeSourceI3DetSigYield: """Constructs a detector signal yield 2-dimensional log spline function for the given flux model with varying parameter values. Parameters ---------- - dataset : instance of Dataset + dataset The Dataset instance holding the sin(dec) binning definition. - data : instance of DatasetData + data The instance of DatasetData holding the monte-carlo event data. This implementation loads the effective area from the provided public data and hence does not need monte-carlo data. - shg : instance of SourceHypoGroup + shg The instance of SourceHypoGroup (i.e. sources and flux model) for which the detector signal yield should get constructed. - ppbar : ProgressBar instance | None + ppbar The instance of ProgressBar of the optional parent progress bar. Returns ------- - detsigyield : instance of SingleParamFluxPointLikeSourceI3DetSigYield + detsigyield The DetSigYield instance for a point-like source with a flux model of a single parameter. """ @@ -133,6 +137,7 @@ def construct_detsigyield(self, dataset, data, shg, ppbar=None): ) # Get integrated live-time in days. + assert data.livetime is not None livetime_days = Livetime.get_integrated_livetime(data.livetime) to_internal_time_unit_factor = self._cfg.to_internal_time_unit(time_unit=units.day) @@ -153,27 +158,31 @@ def construct_detsigyield(self, dataset, data, shg, ppbar=None): # Calculate the detector signal yield in sin_dec vs gamma. def _create_hist( - energy_bin_edges_lower, - energy_bin_edges_upper, - aeff, - fluxmodel, - to_internal_flux_unit_factor, - ): + energy_bin_edges_lower: np.ndarray, + energy_bin_edges_upper: np.ndarray, + aeff: np.ndarray, + fluxmodel: FactorizedFluxModel, + to_internal_flux_unit_factor: float, + ) -> np.ndarray: """Creates a histogram of the detector signal yield for the given sin(dec) binning. Parameters ---------- - energy_bin_edges_lower : 1d ndarray + energy_bin_edges_lower The array holding the lower bin edges in E_nu/GeV. - energy_bin_edges_upper : 1d ndarray + energy_bin_edges_upper The array holding the upper bin edges in E_nu/GeV. - aeff : (n_bins_sin_dec, n_bins_log_energy)-shaped 2d ndarray - The effective area binned data array. + aeff + The (n_bins_sin_dec, n_bins_log_energy)-shaped 2d ndarray holding the effective area binned data array. + fluxmodel + The flux model for which the detector signal yield should get calculated. + to_internal_flux_unit_factor + The factor to convert the flux model unit into the internal flux unit. Returns ------- - h : instance of ndarray + h The (n_bins_sin_dec,)-shaped 1d numpy ndarray containing the detector signal yield values for the different sin_dec bins and the given flux model. diff --git a/skyllh/analyses/i3/publicdata_ps/mcbkg_ps.py b/skyllh/analyses/i3/publicdata_ps/mcbkg_ps.py index f528a28af6..34c0bbb563 100644 --- a/skyllh/analyses/i3/publicdata_ps/mcbkg_ps.py +++ b/skyllh/analyses/i3/publicdata_ps/mcbkg_ps.py @@ -5,6 +5,7 @@ """ import pickle +from collections.abc import Sequence import numpy as np @@ -30,12 +31,11 @@ from skyllh.analyses.i3.publicdata_ps.utils import ( create_energy_cut_spline, ) -from skyllh.core.analysis import ( - SingleSourceMultiDatasetLLHRatioAnalysis as Analysis, -) +from skyllh.core.analysis import SingleSourceMultiDatasetLLHRatioAnalysis from skyllh.core.config import ( Config, ) +from skyllh.core.dataset import Dataset from skyllh.core.event_selection import ( SpatialBoxEventSelectionMethod, ) @@ -110,81 +110,83 @@ def TXS_location(): + """Returns the right-ascention and declination of the blazar TXS 0506+056 + in radians. + """ src_ra = np.radians(77.358) src_dec = np.radians(5.693) return (src_ra, src_dec) def create_analysis( - datasets, - source, - refplflux_Phi0=1, - refplflux_E0=1e3, - refplflux_gamma=2, - ns_seed=100, - ns_min=0, - ns_max=1e3, - gamma_seed=3, - gamma_min=1, - gamma_max=5, - minimizer_impl='LBFGS', - cut_sindec=None, - spl_smooth=None, - compress_data=False, - keep_data_fields=None, - evt_sel_delta_angle_deg=10, - efficiency_mode=None, - tl=None, - ppbar=None, - logger_name=None, -): + datasets: list[Dataset], + source: PointLikeSource, + refplflux_Phi0: float = 1.0, + refplflux_E0: float = 1e3, + refplflux_gamma: float = 2.0, + ns_seed: float = 100.0, + ns_min: float = 0.0, + ns_max: float = 1e3, + gamma_seed: float | None = 3.0, + gamma_min: float = 1.0, + gamma_max: float = 5.0, + minimizer_impl: str = 'LBFGS', + cut_sindec: Sequence[float] | np.ndarray | None = None, + spl_smooth: list[float] | None = None, + compress_data: bool = False, + keep_data_fields: list[str] | None = None, + evt_sel_delta_angle_deg: float = 10.0, + efficiency_mode: str | None = None, + tl: TimeLord | None = None, + ppbar: ProgressBar | None = None, + logger_name: str | None = None, +) -> SingleSourceMultiDatasetLLHRatioAnalysis: """Creates the Analysis instance for this particular analysis. Parameters ---------- - datasets : list of Dataset instances + datasets The list of Dataset instances, which should be used in the analysis. - source : PointLikeSource instance + source The PointLikeSource instance defining the point source position. - refplflux_Phi0 : float + refplflux_Phi0 The flux normalization to use for the reference power law flux model. - refplflux_E0 : float + refplflux_E0 The reference energy to use for the reference power law flux model. - refplflux_gamma : float + refplflux_gamma The spectral index to use for the reference power law flux model. - ns_seed : float + ns_seed Value to seed the minimizer with for the ns fit. - ns_min : float + ns_min Lower bound for ns fit. - ns_max : float + ns_max Upper bound for ns fit. - gamma_seed : float | None + gamma_seed Value to seed the minimizer with for the gamma fit. If set to None, the refplflux_gamma value will be set as gamma_seed. - gamma_min : float + gamma_min Lower bound for gamma fit. - gamma_max : float + gamma_max Upper bound for gamma fit. - minimizer_impl : str + minimizer_impl Minimizer implementation to be used. Supported options are ``"LBFGS"`` (L-BFG-S minimizer used from the :mod:`scipy.optimize` module), or ``"minuit"`` (Minuit minimizer used by the :mod:`iminuit` module). - Default: "LBFGS". - cut_sindec : list of float | None + cut_sindec sin(dec) values at which the energy cut in the southern sky should start. If None, np.sin(np.radians([-2, 0, -3, 0, 0])) is used. - spl_smooth : list of float + spl_smooth Smoothing parameters for the 1D spline for the energy cut. If None, [0., 0.005, 0.05, 0.2, 0.3] is used. - compress_data : bool + compress_data Flag if the data should get converted from float64 into float32. - keep_data_fields : list of str | None + keep_data_fields List of additional data field names that should get kept when loading the data. - evt_sel_delta_angle_deg : float + evt_sel_delta_angle_deg The delta angle in degrees for the event selection optimization methods. - efficiency_mode : str | None + efficiency_mode The efficiency mode the data should get loaded with. Possible values are: @@ -199,17 +201,17 @@ def create_analysis( The default value is ``'time'``. If set to ``None``, the default value will be used. - tl : TimeLord instance | None + tl The TimeLord instance to use to time the creation of the analysis. - ppbar : ProgressBar instance | None + ppbar The instance of ProgressBar for the optional parent progress bar. - logger_name : str | None + logger_name The name of the logger to be used. If set to ``None``, ``__name__`` will be used. Returns ------- - ana : instance of SingleSourceMultiDatasetLLHRatioAnalysis + ana The Analysis instance for this analysis. """ if logger_name is None: @@ -224,6 +226,12 @@ def create_analysis( else: raise NameError(f'Minimizer implementation `{minimizer_impl}` is not supported Please use `LBFGS` or `minuit`.') + dtc_dict = None + dtc_except_fields = None + if compress_data is True: + dtc_dict = {np.dtype(np.float64): np.dtype(np.float32)} + dtc_except_fields = ['mcweight', 'time'] + # Define the flux model. fluxmodel = SteadyPointlikeFFM( Phi0=refplflux_Phi0, energy_profile=PowerLawEnergyFluxProfile(E0=refplflux_E0, gamma=refplflux_gamma) @@ -233,6 +241,8 @@ def create_analysis( param_ns = Parameter(name='ns', initial=ns_seed, valmin=ns_min, valmax=ns_max) # Define the fit parameter gamma. + if gamma_seed is None: + gamma_seed = refplflux_gamma param_gamma = Parameter(name='gamma', initial=gamma_seed, valmin=gamma_min, valmax=gamma_max) # Define the detector signal efficiency implementation method for the @@ -274,7 +284,7 @@ def create_analysis( minimizer = Minimizer(LBFGSMinimizerImpl()) # Create the Analysis instance. - ana = Analysis( + ana = SingleSourceMultiDatasetLLHRatioAnalysis( shg_mgr=shg_mgr, pmm=pmm, test_statistic=test_statistic, @@ -304,13 +314,18 @@ def create_analysis( for ds_idx, ds in enumerate(datasets): # Load the data of the data set. data = ds.load_and_prepare_data( - keep_fields=keep_data_fields, compress=compress_data, efficiency_mode=efficiency_mode, tl=tl + keep_fields=keep_data_fields, + dtc_dict=dtc_dict, + dtc_except_fields=dtc_except_fields, + efficiency_mode=efficiency_mode, + tl=tl, ) + assert data.exp is not None sin_dec_binning = ds.get_binning_definition('sin_dec') # Create the spatial PDF ratio instance for this dataset. - spatial_sigpdf = RayleighPSFPointSourceSignalSpatialPDF(dec_range=np.arcsin(sin_dec_binning.range)) + spatial_sigpdf = RayleighPSFPointSourceSignalSpatialPDF(dec_range=tuple(np.arcsin(sin_dec_binning.range))) spatial_bkgpdf = DataBackgroundI3SpatialPDF(data_exp=data.exp, sin_dec_binning=sin_dec_binning) spatial_pdfratio = SigOverBkgPDFRatio(sig_pdf=spatial_sigpdf, bkg_pdf=spatial_bkgpdf) diff --git a/skyllh/analyses/i3/publicdata_ps/pdfratio.py b/skyllh/analyses/i3/publicdata_ps/pdfratio.py index 290d9a36c2..2f93102348 100644 --- a/skyllh/analyses/i3/publicdata_ps/pdfratio.py +++ b/skyllh/analyses/i3/publicdata_ps/pdfratio.py @@ -1,32 +1,44 @@ import numpy as np +from skyllh.analyses.i3.publicdata_ps.backgroundpdf import PDDataBackgroundI3EnergyPDF, PDMCBackgroundI3EnergyPDF +from skyllh.analyses.i3.publicdata_ps.signalpdf import PDSignalEnergyPDFSet from skyllh.core.logging import ( get_logger, ) from skyllh.core.parameters import ( ParameterModelMapper, ) +from skyllh.core.pdf import PDF from skyllh.core.pdfratio import ( SigSetOverBkgPDFRatio, ) from skyllh.core.py import ( module_class_method_name, ) +from skyllh.core.timing import TimeLord +from skyllh.core.trialdata import TrialDataManager class PDSigSetOverBkgPDFRatio(SigSetOverBkgPDFRatio): - def __init__(self, sig_pdf_set, bkg_pdf, **kwargs): + """This class provides the signal-over-background PDF ratio for the public + data, using a signal energy PDF set defined for discrete gamma values. + """ + + def __init__( + self, + sig_pdf_set: PDSignalEnergyPDFSet, + bkg_pdf: PDMCBackgroundI3EnergyPDF | PDDataBackgroundI3EnergyPDF, + **kwargs, + ): """Creates a PDFRatio instance for the public data. It takes a signal PDF set for different discrete gamma values. Parameters ---------- - sig_pdf_set : instance of PDSignalEnergyPDFSet - The PDSignalEnergyPDFSet instance holding the set of signal energy - PDFs. - bkg_pdf : instance of PDDataBackgroundI3EnergyPDF - The PDDataBackgroundI3EnergyPDF instance holding the background - energy PDF. + sig_pdf_set + Set of signal energy PDFs. + bkg_pdf + Background energy PDF. """ self._logger = get_logger(module_class_method_name(self, '__init__')) @@ -43,23 +55,23 @@ def __init__(self, sig_pdf_set, bkg_pdf, **kwargs): # method was called). self._cache_tdm_trial_data_state_id = None self._cache_fitparams_hash = None - self._cache_ratio = None - self._cache_grads = None + self._cache_ratio: np.ndarray | None = None + self._cache_grads: np.ndarray | None = None - def _is_cached(self, tdm, fitparams_hash): + def _is_cached(self, tdm: TrialDataManager, fitparams_hash: int) -> bool: """Checks if the ratio and gradients for the given hash of local fit parameters are already cached. Parameters ---------- - tdm : instance of TrialDataManager + tdm The instance of TrialDataManager holding the trial data events. - fitparams_hach : int + fitparams_hash The hash value of the local fit parameter values. Returns ------- - check : bool + check ``True`` if the ratio and gradient values are already cached, ``False`` otherwise. """ @@ -70,19 +82,19 @@ def _is_cached(self, tdm, fitparams_hash): and (self._cache_grads is not None) ) - def _get_hash_of_local_sig_fit_param_values(self, src_params_recarray): + def _get_hash_of_local_sig_fit_param_values(self, src_params_recarray: np.ndarray) -> int: """Gets the hash of the values of the local signal fit parameters from the given ``src_params_recarray``. Parameters ---------- - src_params_recarray : instance of ndarray + src_params_recarray The (N_sources,)-shaped structured numpy ndarray holding the local parameter names and values of the sources. Returns ------- - hash : int + hash The hash of the (N_fitparams, N_sources)-shaped tuple of tuples holding the values of the local signal fit parameters. """ @@ -130,6 +142,7 @@ def _get_ratio_values(self, tdm, eventdata, gridparams_recarray, n_values): ratio[m_values] = sig_pd[m_values] + assert isinstance(self.bkg_pdf, PDF) (bkg_pd, _) = self.bkg_pdf.get_pd(tdm=tdm) (bkg_pd,) = tdm.broadcast_selected_events_arrays_to_values_arrays((bkg_pd,)) @@ -159,7 +172,7 @@ def _calculate_ratio_and_grads(self, tdm, src_params_recarray, fitparams_hash): events and sources given the fit parameters using the interpolation method for the fit parameter. It caches the results. """ - (ratio, grads) = self._interpolmethod(tdm=tdm, eventdata=None, params_recarray=src_params_recarray) + (ratio, grads) = self._interpolmethod(tdm=tdm, eventdata=np.empty(0), params_recarray=src_params_recarray) # Cache the ratio and gradient values. self._cache_tdm_trial_data_state_id = tdm.trial_data_state_id @@ -167,27 +180,29 @@ def _calculate_ratio_and_grads(self, tdm, src_params_recarray, fitparams_hash): self._cache_ratio = ratio self._cache_grads = grads - def get_ratio(self, tdm, src_params_recarray, tl=None): + def get_ratio( + self, tdm: TrialDataManager, src_params_recarray: np.ndarray, tl: TimeLord | None = None + ) -> np.ndarray: """Calculates the PDF ratio values for all events and sources. Parameters ---------- - tdm : instance of TrialDataManager + tdm The instance of TrialDataManager holding the trial data events for which the PDF ratio values should get calculated. - src_params_recarray : instance of numpy structured ndarray | None + src_params_recarray The (N_sources,)-shaped numpy structured ndarray holding the parameter names and values of the sources. See the documentation of the :meth:`skyllh.core.parameters.ParameterModelMapper.create_src_params_recarray` method for more information. - tl : instance of TimeLord | None + tl The optional instance of TimeLord that should be used to measure timing information. Returns ------- - ratios : instance of ndarray + ratios The (N_values,)-shaped 1d numpy ndarray of float holding the PDF ratio value for each trial event and source. """ @@ -195,38 +210,46 @@ def get_ratio(self, tdm, src_params_recarray, tl=None): # Check if the ratio value is already cached. if self._is_cached(tdm=tdm, fitparams_hash=fitparams_hash): + assert self._cache_ratio is not None return self._cache_ratio self._calculate_ratio_and_grads(tdm=tdm, src_params_recarray=src_params_recarray, fitparams_hash=fitparams_hash) + assert self._cache_ratio is not None return self._cache_ratio - def get_gradient(self, tdm, src_params_recarray, fitparam_id, tl=None): + def get_gradient( + self, + tdm: TrialDataManager, + src_params_recarray: np.ndarray, + fitparam_id: int, + tl: TimeLord | None = None, + ) -> np.ndarray: """Retrieves the PDF ratio gradient for the global fit parameter ``fitparam_id`` for each trial data event and source, given the given set of parameters ``src_params_recarray`` for each source. Parameters ---------- - tdm : instance of TrialDataManager + tdm The instance of TrialDataManager holding the trial data events for which the PDF ratio gradient values should get calculated. - src_params_recarray : instance of numpy structured ndarray | None + src_params_recarray The (N_sources,)-shaped numpy structured ndarray holding the parameter names and values of the sources. See the documentation of the :meth:`skyllh.core.parameters.ParameterModelMapper.create_src_params_recarray` method for more information. - fitparam_id : int + fitparam_id The name of the fit parameter for which the gradient should get calculated. - tl : instance of TimeLord | None + tl The optional TimeLord instance that should be used to measure timing information. Returns ------- - grad : instance of ndarray + grad The (N_values,)-shaped numpy ndarray holding the gradient values for all sources and trial events w.r.t. the given global fit parameter. @@ -245,6 +268,7 @@ def get_gradient(self, tdm, src_params_recarray, fitparam_id, tl=None): # Loop through the parameters of the signal PDF set and match them with # the global fit parameter. + assert self._cache_grads is not None for pidx, pname in enumerate(self._sig_pdf_set.param_grid_set.params_name_list): if pname not in src_params_recarray.dtype.fields: continue diff --git a/skyllh/analyses/i3/publicdata_ps/scripts/mceq_atm_bkg.py b/skyllh/analyses/i3/publicdata_ps/scripts/mceq_atm_bkg.py index c2ba7f9879..d8adfce04c 100644 --- a/skyllh/analyses/i3/publicdata_ps/scripts/mceq_atm_bkg.py +++ b/skyllh/analyses/i3/publicdata_ps/scripts/mceq_atm_bkg.py @@ -2,10 +2,10 @@ import os.path import pickle -import crflux.models as pm -import mceq_config as config +import crflux.models as pm # pyright: ignore[reportMissingImports] +import mceq_config as config # pyright: ignore[reportMissingImports] import numpy as np -from MCEq.core import ( +from MCEq.core import ( # pyright: ignore[reportMissingImports] MCEqRun, ) diff --git a/skyllh/analyses/i3/publicdata_ps/signal_generator.py b/skyllh/analyses/i3/publicdata_ps/signal_generator.py index 1e122108af..4b2a16d75d 100644 --- a/skyllh/analyses/i3/publicdata_ps/signal_generator.py +++ b/skyllh/analyses/i3/publicdata_ps/signal_generator.py @@ -1,3 +1,5 @@ +from typing import Any, cast + import numpy as np from scipy import ( interpolate, @@ -12,6 +14,7 @@ from skyllh.analyses.i3.publicdata_ps.utils import ( psi_to_dec_and_ra, ) +from skyllh.core.dataset import Dataset from skyllh.core.flux_model import ( TimeFluxProfile, ) @@ -28,12 +31,15 @@ issequence, module_classname, ) +from skyllh.core.random import RandomStateService +from skyllh.core.services import SrcDetSigYieldWeightsService from skyllh.core.signal_generation import ( HasEnergyRange, ) from skyllh.core.signal_generator import ( SignalGenerator, ) +from skyllh.core.source_hypo_grouping import SourceHypoGroupManager from skyllh.core.storage import ( DataFieldRecordArray, ) @@ -52,13 +58,13 @@ class PDDatasetSignalGenerator( def __init__( self, - shg_mgr, - ds, - ds_idx, - energy_cut_spline=None, - cut_sindec=None, - energy_range=None, - sm=None, + shg_mgr: SourceHypoGroupManager, + ds: Dataset, + ds_idx: int, + energy_cut_spline: interpolate.UnivariateSpline | None = None, + cut_sindec: float | None = None, + energy_range: tuple[float, float] | None = None, + sm: PDSmearingMatrix | None = None, **kwargs, ): """Creates a new instance of the signal generator for generating @@ -66,24 +72,24 @@ def __init__( Parameters ---------- - shg_mgr : instance of SourceHypoGroupManager + shg_mgr The instance of SourceHypoGroupManager defining the source hypothesis groups. - ds : instance of Dataset + ds The instance of Dataset for which signal events should get generated. - ds_idx : int + ds_idx The index of the dataset. - energy_cut_spline : scipy.interpolate.UnivariateSpline + energy_cut_spline A spline of E(sin_dec) that defines the declination dependent energy cut in the IceCube southern sky. - cut_sindec : float + cut_sindec The sine of the declination to start applying the energy cut. The cut will be applied from this declination down. - energy_range : 2-element tuple of float | None + energy_range The energy range in which signal events should be generated. If set to None, the full energy range (1e2 - 1e9 GeV) is used. - sm : instance of PDSmearingMatrix | None + sm A pre-loaded smearing matrix to reuse. If ``None``, the matrix is loaded from the dataset's auxiliary data files. """ @@ -110,7 +116,7 @@ def __init__( self.energy_range = energy_range @property - def energy_range(self): + def energy_range(self) -> tuple[float, float]: """The configured true-energy range for signal generation in GeV. It is a 2-element tuple ``(E_min, E_max)`` in GeV. If no explicit @@ -222,17 +228,17 @@ def _create_source_dependent_data_structures(self): self._energy_range_correction_factors = self._calculate_energy_range_correction_factors() @staticmethod - def _eval_spline(x, spl): + def _eval_spline(x, spl) -> np.ndarray: """Evaluates the given spline at the given coordinates.""" x = np.asarray(x) if np.any(x < 0) or np.any(x > 1): raise ValueError(f'{x} is outside of the valid spline range. The valid range is [0,1].') - values = interpolate.splev(x, spl, ext=3) + values = np.asarray(interpolate.splev(x, spl, ext=3)) return values - def _create_inv_cdf_spline(self, src_idx, fluxmodel, log_e_min, log_e_max): + def _create_inv_cdf_spline(self, src_idx: int, fluxmodel, log_e_min, log_e_max): """Creates a spline for the inverse cumulative distribution function of the detectable true energy probability distribution. """ @@ -276,7 +282,15 @@ def _create_inv_cdf_spline(self, src_idx, fluxmodel, log_e_min, log_e_max): # Build a spline for the inverse CDF. return interpolate.splrep(cum_per_bin, bin_centers, k=1, s=0) - def _draw_signal_events_for_source(self, rss, src_dec, src_ra, dec_idx, log10_true_e_inv_cdf_spl, n_events): + def _draw_signal_events_for_source( + self, + rss: RandomStateService, + src_dec: float, + src_ra: float, + dec_idx: int, + log10_true_e_inv_cdf_spl, + n_events: int, + ) -> DataFieldRecordArray: """Generates `n_events` signal events for the given source location given the given inverse cumulative density function for the log10(E_true/GeV) distribution. @@ -286,24 +300,24 @@ def _draw_signal_events_for_source(self, rss, src_dec, src_ra, dec_idx, log10_tr Parameters ---------- - rss : instance of RandomStateService + rss The instance of RandomStateService to use for drawing random numbers. - src_dec : float + src_dec The declination of the source in radians. - src_ra : float + src_ra The right-ascention of the source in radians. - dec_idx : int + dec_idx The SM's declination bin index of the source's declination. - log10_true_e_inv_cdf_spl : instance of scipy.interpolate.splrep + log10_true_e_inv_cdf_spl The linear spline interpolation representation of the inverse cummulative density function of the log10(E_true/GeV) distribution. - n_events : int + n_events The number of events to generate. Returns ------- - events : instance of DataFieldRecordArray of size `n_events` + events The instance of DataFieldRecordArray of length `n_events` holding the event data. It contains the following data fields: @@ -392,28 +406,30 @@ def change_shg_mgr(self, shg_mgr): self._create_source_dependent_data_structures() @staticmethod - def create_energy_filter_mask(events, spline, cut_sindec, logger): + def create_energy_filter_mask( + events: DataFieldRecordArray, spline: interpolate.UnivariateSpline, cut_sindec: float | None, logger + ) -> np.ndarray: """Creates a mask for cutting all events below ``cut_sindec`` that have an energy smaller than the energy spline at their declination. Parameters ---------- - events : instance of DataFieldRecordArray + events The instance of DataFieldRecordArray holding the generated signal events. - spline : instance of scipy.interpolate.UnivariateSpline + spline A spline of E(sin_dec) that defines the declination dependent energy cut in the IceCube southern sky. - cut_sindec : float + cut_sindec The sine of the declination to start applying the energy cut. The cut will be applied from this declination down. - logger : instance of logging.Logger + logger The Logger instance. Returns ------- - filter_mask : instance of numpy ndarray + filter_mask The (len(events),)-shaped numpy ndarray with the mask of the events to cut. """ @@ -425,23 +441,25 @@ def create_energy_filter_mask(events, spline, cut_sindec, logger): return filter_mask - def generate_signal_events_for_source(self, rss, src_idx, n_events): + def generate_signal_events_for_source( + self, rss: RandomStateService, src_idx: int, n_events: int + ) -> DataFieldRecordArray | None: """Generates ``n_events`` signal events for the given source location and flux model. Parameters ---------- - rss : instance of RandomStateService + rss The instance of RandomStateService providing the random number generator state. - src_idx : int + src_idx The index of the source. - n_events : int + n_events Number of signal events to be generated. Returns ------- - events : instance of DataFieldRecordArray + events The numpy record array holding the event data. It contains the following data fields: @@ -499,46 +517,45 @@ def generate_signal_events_for_source(self, rss, src_idx, n_events): def generate_signal_events( self, - rss, - mean, - poisson=True, - src_detsigyield_weights_service=None, - ): + rss: RandomStateService, + mean: float, + poisson: bool = True, + src_detsigyield_weights_service: SrcDetSigYieldWeightsService | None = None, + ) -> tuple[int, dict[int, DataFieldRecordArray]]: """Generates ``mean`` number of signal events. Parameters ---------- - rss : instance of RandomStateService + rss The instance of RandomStateService providing the random number generator state. - mean : int | float + mean The mean number of signal events. If the ``poisson`` argument is set to True, the actual number of generated signal events will be drawn from a Poisson distribution with this given mean value of signal events. - poisson : bool + poisson If set to True, the actual number of generated signal events will be drawn from a Poisson distribution with the given mean value of signal events. If set to False, the argument ``mean`` specifies the actual number of generated signal events. - src_detsigyield_weights_service : instance of SrcDetSigYieldWeightsService + src_detsigyield_weights_service The instance of SrcDetSigYieldWeightsService providing the weighting of the sources within the detector. Returns ------- - n_signal : int + n_signal The number of generated signal events. - signal_events_dict : dict of DataFieldRecordArray + signal_events_dict The dictionary holding the DataFieldRecordArray instances with the generated signal events. Each key of this dictionary represents the dataset index for which the signal events have been generated. """ - if poisson: - mean = rss.random.poisson(float_cast(mean, 'The `mean` argument must be castable to type of float!')) + _mean = rss.random.poisson(float(mean)) if poisson else mean - n_events = int_cast(mean, 'The `mean` argument must be castable to type of int!') + n_events = int_cast(_mean, 'The `mean` argument must be castable to type of int!') if src_detsigyield_weights_service is None: raise ValueError( @@ -548,6 +565,7 @@ def generate_signal_events( (a_jk, _) = src_detsigyield_weights_service.get_weights() + assert a_jk is not None a_k = np.copy(a_jk[self.ds_idx]) a_k /= np.sum(a_k) @@ -595,7 +613,9 @@ def _calculate_flux_weight(self, src_idx, log_e_min, log_e_max): if not np.any(m): return 0.0 - flux_integral = fluxmodel.energy_profile.get_integral(E1=10 ** overlap_low[m], E2=10 ** overlap_high[m]) + flux_integral = cast(Any, fluxmodel).energy_profile.get_integral( + E1=10 ** overlap_low[m], E2=10 ** overlap_high[m] + ) aeff_for_dec = self._get_cached_aeff_for_source(src_idx=src_idx, src_dec=src.dec)[m] @@ -694,30 +714,30 @@ class TimeDependentPDDatasetSignalGenerator( def __init__( self, - shg_mgr, - ds, - ds_idx, - livetime, - time_flux_profile, + shg_mgr: SourceHypoGroupManager, + ds: Dataset, + ds_idx: int, + livetime: Livetime, + time_flux_profile: TimeFluxProfile, energy_cut_spline=None, - cut_sindec=None, + cut_sindec: float | None = None, **kwargs, ): """ Parameters ---------- - shg_mgr : instance of SourceHypoGroupManager + shg_mgr The instance of SourceHypoGroupManager that defines the list of source hypothesis groups, i.e. the list of sources. - ds : instance of Dataset + ds The instance of Dataset for which signal events should get generated. - ds_idx : int + ds_idx The index of the dataset. - livetime : instance of Livetime + livetime The instance of Livetime providing the live-time information of the dataset. - time_flux_profile : instance of TimeFluxProfile + time_flux_profile The instance of TimeFluxProfile providing the time profile of the source(s). @@ -726,10 +746,10 @@ def __init__( At this time the some time profile will be used for all sources! - energy_cut_spline : scipy.interpolate.UnivariateSpline + energy_cut_spline A spline of E(sin_dec) that defines the declination dependent energy cut in the IceCube southern sky. - cut_sindec : float + cut_sindec The sine of the declination to start applying the energy cut. The cut will be applied from this declination down. """ @@ -764,39 +784,38 @@ def livetime(self, lt): def generate_signal_events( self, - rss, - mean, - poisson=True, - src_detsigyield_weights_service=None, - **kwargs, - ): + rss: RandomStateService, + mean: float, + poisson: bool = True, + src_detsigyield_weights_service: SrcDetSigYieldWeightsService | None = None, + ) -> tuple[int, dict[int, DataFieldRecordArray]]: """Generates ``mean`` number of signal events with times. Parameters ---------- - rss : instance of RandomStateService + rss The instance of RandomStateService providing the random number generator state. - mean : int | float + mean The mean number of signal events. If the ``poisson`` argument is set to True, the actual number of generated signal events will be drawn from a Poisson distribution with this given mean value of signal events. - poisson : bool + poisson If set to True, the actual number of generated signal events will be drawn from a Poisson distribution with the given mean value of signal events. If set to False, the argument ``mean`` specifies the actual number of generated signal events. - src_detsigyield_weights_service : instance of SrcDetSigYieldWeightsService + src_detsigyield_weights_service The instance of SrcDetSigYieldWeightsService providing the weighting of the sources within the detector. Returns ------- - n_signal : int + n_signal The number of generated signal events. - signal_events_dict : dict of DataFieldRecordArray + signal_events_dict The dictionary holding the DataFieldRecordArray instances with the generated signal events. Each key of this dictionary represents the dataset index for which the signal events have been generated. @@ -806,7 +825,6 @@ def generate_signal_events( mean=mean, poisson=poisson, src_detsigyield_weights_service=src_detsigyield_weights_service, - **kwargs, ) # Create a scipy.stats.rv_continuous instance for the time flux profile. diff --git a/skyllh/analyses/i3/publicdata_ps/signalpdf.py b/skyllh/analyses/i3/publicdata_ps/signalpdf.py index 08e1ca113d..8a0c94483e 100644 --- a/skyllh/analyses/i3/publicdata_ps/signalpdf.py +++ b/skyllh/analyses/i3/publicdata_ps/signalpdf.py @@ -12,6 +12,7 @@ from skyllh.core.binning import ( get_bincenters_from_binedges, ) +from skyllh.core.dataset import Dataset from skyllh.core.flux_model import ( FactorizedFluxModel, ) @@ -32,13 +33,13 @@ PDFAxis, PDFSet, ) +from skyllh.core.progressbar import ProgressBar from skyllh.core.py import ( classname, module_classname, ) -from skyllh.core.timing import ( - TaskTimer, -) +from skyllh.core.timing import TaskTimer, TimeLord +from skyllh.core.trialdata import TrialDataManager from skyllh.i3.dataset import ( I3Dataset, ) @@ -52,7 +53,7 @@ class PDSignalEnergyPDF( def __init__( self, - f_e_spl, + f_e_spl: FctSpline1D, **kwargs, ): """Creates a new signal energy PDF instance for a particular spectral @@ -60,7 +61,7 @@ def __init__( Parameters ---------- - f_e_spl : instance of FctSpline1D + f_e_spl The instance of FctSpline1D representing the spline of the energy PDF. """ @@ -82,26 +83,28 @@ def __init__( # Add the PDF axes. self.add_axis(PDFAxis(name='log_energy', vmin=self.log10_reco_e_min, vmax=self.log10_reco_e_max)) - def assert_is_valid_for_trial_data(self, tdm, tl=None): - pass + def assert_is_valid_for_trial_data(self, tdm, tl=None, **kwargs): + """Checks if this PDF is valid for the given trial data. This PDF is + always valid by construction, hence this method does nothing. + """ - def get_pd_by_log10_reco_e(self, log10_reco_e, tl=None): + def get_pd_by_log10_reco_e(self, log10_reco_e: np.ndarray, tl: TimeLord | None = None) -> np.ndarray: """Calculates the probability density for the given log10(E_reco/GeV) values using the spline representation of the PDF. Parameters ---------- - log10_reco_e : instance of ndarray + log10_reco_e The (n_log10_reco_e,)-shaped numpy ndarray holding the log10(E_reco/GeV) values for which the energy PDF should get evaluated. - tl : instance of TimeLord | None + tl The optional instance of TimeLord that should be used to measure timing information. Returns ------- - pd : instance of numpy ndarray + pd The (n_log10_reco_e,)-shaped numpy ndarray with the probability density for each energy value. """ @@ -115,36 +118,42 @@ def get_pd_by_log10_reco_e(self, log10_reco_e, tl=None): return pd - def get_pd(self, tdm, params_recarray=None, tl=None): + def get_pd( + self, + tdm: TrialDataManager, + params_recarray: np.ndarray | None = None, + tl: TimeLord | None = None, + ) -> tuple[np.ndarray, dict]: """Calculates the probability density for all given trial data events and sources. Parameters ---------- - tdm : instance of TrialDataManager + tdm The instance of TrialDataManager holding the trial data events for which the probability density should be looked up. The following data fields must be present: - log_energy : float + log_energy The base-10 logarithm of the reconstructed energy. - params_recarray : None + params_recarray Unused interface argument. - tl : instance of TimeLord | None + tl The optional instance of TimeLord that should be used to measure timing information. Returns ------- - pd : instance of ndarray + pd The (N_values,)-shaped numpy ndarray holding the probability density for each trial data event and source. - grads : dict + grads The dictionary holding the gradient values for each global fit parameter. By definition this PDF does not depend on any fit parameters, hence, this is an empty dictionary. """ + assert tdm.src_evt_idxs is not None evt_idxs = tdm.src_evt_idxs[1] log10_reco_e = np.take(tdm['log_energy'], evt_idxs) @@ -156,7 +165,7 @@ def get_pd(self, tdm, params_recarray=None, tl=None): return (pd, grads) -class PDSignalEnergyPDFSet( +class PDSignalEnergyPDFSet( # pyright: ignore[reportIncompatibleMethodOverride] PDFSet, IsSignalPDF, PDF, @@ -169,35 +178,35 @@ class PDSignalEnergyPDFSet( def __init__( self, - ds, - src_dec, - fluxmodel, - param_grid_set, - ncpu=None, - ppbar=None, - sm=None, + ds: Dataset, + src_dec: float, + fluxmodel: FactorizedFluxModel, + param_grid_set: ParameterGrid | ParameterGridSet, + ncpu: int | None = None, + ppbar: ProgressBar | None = None, + sm: PDSmearingMatrix | None = None, **kwargs, ): """Creates a new PDSignalEnergyPDFSet instance for the public data. Parameters ---------- - ds : instance of Dataset + ds The instance of Dataset that defines the dataset of the public data. - src_dec : float + src_dec The declination of the source in radians. - fluxmodel : instance of FactorizedFluxModel + fluxmodel The instance of FactorizedFluxModel that defines the source's flux model. - param_grid_set : instance of ParameterGrid | instance of ParameterGridSet + param_grid_set The parameter grid set defining the grids of the parameters this energy PDF set depends on. - ncpu : int | None + ncpu The number of CPUs to utilize. Global setting will take place if not specified, i.e. set to None. - ppbar : instance of ProgressBar | None + ppbar The instance of ProgressBar for the optional parent progress bar. - sm : instance of PDSmearingMatrix | None + sm A pre-loaded smearing matrix to reuse. If ``None``, the matrix is loaded from the dataset's auxiliary data files. """ diff --git a/skyllh/analyses/i3/publicdata_ps/smearing_matrix.py b/skyllh/analyses/i3/publicdata_ps/smearing_matrix.py index 0024fd82d0..d47ec2a92e 100644 --- a/skyllh/analyses/i3/publicdata_ps/smearing_matrix.py +++ b/skyllh/analyses/i3/publicdata_ps/smearing_matrix.py @@ -1,48 +1,61 @@ import numpy as np +from skyllh.core.random import RandomStateService from skyllh.core.storage import create_FileLoader -def load_smearing_histogram(pathfilenames): +def load_smearing_histogram( + pathfilenames: str | list[str], +) -> tuple[ + np.ndarray, + np.ndarray, + np.ndarray, + np.ndarray, + np.ndarray, + np.ndarray, + np.ndarray, + np.ndarray, + np.ndarray, +]: """Loads the 5D smearing histogram from the given data file. Parameters ---------- - pathfilenames : str | list of str - The file name of the data file. + pathfilenames + The file name(s) of the data file(s). Returns ------- - histogram : 5d ndarray + histogram The 5d histogram array holding the probability values of the smearing matrix. The axes are (true_e, true_dec, reco_e, psi, ang_err). - true_e_bin_edges : 1d ndarray + true_e_bin_edges The ndarray holding the bin edges of the true energy axis. - true_dec_bin_edges : 1d ndarray + true_dec_bin_edges The ndarray holding the bin edges of the true declination axis in radians. - reco_e_lower_edges : 3d ndarray + reco_e_lower_edges The 3d ndarray holding the lower bin edges of the reco energy axis. For each pair of true_e and true_dec different reco energy bin edges are provided. The shape is (n_true_e, n_true_dec, n_reco_e). - reco_e_upper_edges : 3d ndarray + reco_e_upper_edges The 3d ndarray holding the upper bin edges of the reco energy axis. For each pair of true_e and true_dec different reco energy bin edges are provided. The shape is (n_true_e, n_true_dec, n_reco_e). - psi_lower_edges : 4d ndarray + psi_lower_edges The 4d ndarray holding the lower bin edges of the psi axis in radians. The shape is (n_true_e, n_true_dec, n_reco_e, n_psi). - psi_upper_edges : 4d ndarray + psi_upper_edges The 4d ndarray holding the upper bin edges of the psi axis in radians. The shape is (n_true_e, n_true_dec, n_reco_e, n_psi). - ang_err_lower_edges : 5d ndarray + ang_err_lower_edges The 5d ndarray holding the lower bin edges of the angular error axis in radians. The shape is (n_true_e, n_true_dec, n_reco_e, n_psi, n_ang_err). - ang_err_upper_edges : 5d ndarray + ang_err_upper_edges The 5d ndarray holding the upper bin edges of the angular error axis in radians. The shape is (n_true_e, n_true_dec, n_reco_e, n_psi, n_ang_err). @@ -292,18 +305,18 @@ def pdf(self): return pdf - def get_true_dec_idx(self, true_dec): + def get_true_dec_idx(self, true_dec: float) -> int: """Returns the true declination index for the given true declination value. Parameters ---------- - true_dec : float + true_dec The true declination value in radians. Returns ------- - true_dec_idx : int + true_dec_idx The index of the declination bin for the given declination value. """ if (true_dec < self.true_dec_bin_edges[0]) or (true_dec > self.true_dec_bin_edges[-1]): @@ -311,16 +324,16 @@ def get_true_dec_idx(self, true_dec): true_dec_idx = np.digitize(true_dec, self.true_dec_bin_edges) - 1 - return true_dec_idx + return int(true_dec_idx) - def get_log10_true_e_idx(self, log10_true_e, upper_edge=False): + def get_log10_true_e_idx(self, log10_true_e: float, upper_edge: bool = False) -> int: """Returns the true-energy bin index or edge index for a log10 value. Parameters ---------- - log10_true_e : float + log10_true_e The log10 value of the true energy. - upper_edge : bool + upper_edge If False (default), return the bin index whose lower edge is less than or equal to ``log10_true_e`` and whose upper edge is strictly greater than ``log10_true_e``. @@ -330,7 +343,7 @@ def get_log10_true_e_idx(self, log10_true_e, upper_edge=False): Returns ------- - log10_true_e_idx : int + log10_true_e_idx If ``upper_edge`` is False, the index of the true log10 energy bin for the given log10 true energy value. If ``upper_edge`` is True, the index of the corresponding upper @@ -346,28 +359,28 @@ def get_log10_true_e_idx(self, log10_true_e, upper_edge=False): ) if upper_edge: - return np.digitize(log10_true_e, self._true_e_bin_edges, right=True) + return int(np.digitize(log10_true_e, self._true_e_bin_edges, right=True)) log10_true_e_idx = np.digitize(log10_true_e, self._true_e_bin_edges) - 1 - return log10_true_e_idx + return int(log10_true_e_idx) - def get_reco_e_idx(self, true_e_idx, true_dec_idx, reco_e): + def get_reco_e_idx(self, true_e_idx: int, true_dec_idx: int, reco_e: float) -> int | None: """Returns the bin index for the given reco energy value given the given true energy and true declination bin indices. Parameters ---------- - true_e_idx : int + true_e_idx The index of the true energy bin. - true_dec_idx : int + true_dec_idx The index of the true declination bin. - reco_e : float + reco_e The reco energy value for which the bin index should get returned. Returns ------- - reco_e_idx : int | None + reco_e_idx The index of the reco energy bin the given reco energy value falls into. It returns None if the value is out of range. """ @@ -383,25 +396,25 @@ def get_reco_e_idx(self, true_e_idx, true_dec_idx, reco_e): return reco_e_idx - def get_psi_idx(self, true_e_idx, true_dec_idx, reco_e_idx, psi): + def get_psi_idx(self, true_e_idx: int, true_dec_idx: int, reco_e_idx: int, psi: float) -> int | None: """Returns the bin index for the given psi value given the true energy, true declination and reco energy bin indices. Parameters ---------- - true_e_idx : int + true_e_idx The index of the true energy bin. - true_dec_idx : int + true_dec_idx The index of the true declination bin. - reco_e_idx : int + reco_e_idx The index of the reco energy bin. - psi : float + psi The psi value in radians for which the bin index should get returned. Returns ------- - psi_idx : int | None + psi_idx The index of the psi bin the given psi value falls into. It returns None if the value is out of range. """ @@ -417,27 +430,29 @@ def get_psi_idx(self, true_e_idx, true_dec_idx, reco_e_idx, psi): return psi_idx - def get_ang_err_idx(self, true_e_idx, true_dec_idx, reco_e_idx, psi_idx, ang_err): + def get_ang_err_idx( + self, true_e_idx: int, true_dec_idx: int, reco_e_idx: int, psi_idx: int, ang_err: float + ) -> int | None: """Returns the bin index for the given angular error value given the true energy, true declination, reco energy, and psi bin indices. Parameters ---------- - true_e_idx : int + true_e_idx The index of the true energy bin. - true_dec_idx : int + true_dec_idx The index of the true declination bin. - reco_e_idx : int + reco_e_idx The index of the reco energy bin. - psi_idx : int + psi_idx The index of the psi bin. - ang_err : float + ang_err The angular error value in radians for which the bin index should get returned. Returns ------- - ang_err_idx : int | None + ang_err_idx The index of the angular error bin the given angular error value falls into. It returns None if the value is out of range. """ @@ -453,20 +468,20 @@ def get_ang_err_idx(self, true_e_idx, true_dec_idx, reco_e_idx, psi_idx, ang_err return ang_err_idx - def get_true_log_e_range_with_valid_log_e_pdfs(self, dec_idx): + def get_true_log_e_range_with_valid_log_e_pdfs(self, dec_idx: int) -> tuple[float, float]: """Determines the true log energy range for which log_e PDFs are available for the given declination bin. Parameters ---------- - dec_idx : int + dec_idx The declination bin index. Returns ------- - min_log_true_e : float + min_log_true_e The minimum true log energy value. - max_log_true_e : float + max_log_true_e The maximum true log energy value. """ m = np.sum((self.reco_e_upper_edges[:, dec_idx] - self.reco_e_lower_edges[:, dec_idx] > 0), axis=1) != 0 @@ -475,7 +490,9 @@ def get_true_log_e_range_with_valid_log_e_pdfs(self, dec_idx): return (min_log_true_e, max_log_true_e) - def get_log_e_pdf(self, log_true_e_idx, dec_idx): + def get_log_e_pdf( + self, log_true_e_idx: int, dec_idx: int + ) -> tuple[np.ndarray | None, np.ndarray | None, np.ndarray | None, np.ndarray | None]: """Retrieves the log_e PDF from the given true energy bin index and source bin index. Returns (None, None, None, None) if any of the bin indices are less then @@ -483,20 +500,20 @@ def get_log_e_pdf(self, log_true_e_idx, dec_idx): Parameters ---------- - log_true_e_idx : int + log_true_e_idx The index of the true energy bin. - dec_idx : int + dec_idx The index of the declination bin. Returns ------- - pdf : 1d ndarray + pdf The log_e pdf values. - lower_bin_edges : 1d ndarray + lower_bin_edges The lower bin edges of the energy pdf histogram. - upper_bin_edges : 1d ndarray + upper_bin_edges The upper bin edges of the energy pdf histogram. - bin_widths : 1d ndarray + bin_widths The bin widths of the energy pdf histogram. """ if log_true_e_idx < 0 or dec_idx < 0: @@ -518,7 +535,9 @@ def get_log_e_pdf(self, log_true_e_idx, dec_idx): return (pdf, lower_bin_edges, upper_bin_edges, bin_widths) - def get_psi_pdf(self, log_true_e_idx, dec_idx, log_e_idx): + def get_psi_pdf( + self, log_true_e_idx: int, dec_idx: int, log_e_idx: int + ) -> tuple[np.ndarray | None, np.ndarray | None, np.ndarray | None, np.ndarray | None]: """Retrieves the psi PDF from the given true energy bin index, the source bin index, and the log_e bin index. Returns (None, None, None, None) if any of the bin indices are less then @@ -526,22 +545,22 @@ def get_psi_pdf(self, log_true_e_idx, dec_idx, log_e_idx): Parameters ---------- - log_true_e_idx : int + log_true_e_idx The index of the true energy bin. - dec_idx : int + dec_idx The index of the declination bin. - log_e_idx : int + log_e_idx The index of the log_e bin. Returns ------- - pdf : 1d ndarray + pdf The psi pdf values. - lower_bin_edges : 1d ndarray + lower_bin_edges The lower bin edges of the psi pdf histogram. - upper_bin_edges : 1d ndarray + upper_bin_edges The upper bin edges of the psi pdf histogram. - bin_widths : 1d ndarray + bin_widths The bin widths of the psi pdf histogram. """ if log_true_e_idx < 0 or dec_idx < 0 or log_e_idx < 0: @@ -563,7 +582,9 @@ def get_psi_pdf(self, log_true_e_idx, dec_idx, log_e_idx): return (pdf, lower_bin_edges, upper_bin_edges, bin_widths) - def get_ang_err_pdf(self, log_true_e_idx, dec_idx, log_e_idx, psi_idx): + def get_ang_err_pdf( + self, log_true_e_idx: int, dec_idx: int, log_e_idx: int, psi_idx: int + ) -> tuple[np.ndarray | None, np.ndarray | None, np.ndarray | None, np.ndarray | None]: """Retrieves the angular error PDF from the given true energy bin index, the source bin index, the log_e bin index, and the psi bin index. Returns (None, None, None, None) if any of the bin indices are less then @@ -571,24 +592,24 @@ def get_ang_err_pdf(self, log_true_e_idx, dec_idx, log_e_idx, psi_idx): Parameters ---------- - log_true_e_idx : int + log_true_e_idx The index of the true energy bin. - dec_idx : int + dec_idx The index of the declination bin. - log_e_idx : int + log_e_idx The index of the log_e bin. - psi_idx : int + psi_idx The index of the psi bin. Returns ------- - pdf : 1d ndarray + pdf The ang_err pdf values. - lower_bin_edges : 1d ndarray + lower_bin_edges The lower bin edges of the ang_err pdf histogram. - upper_bin_edges : 1d ndarray + upper_bin_edges The upper bin edges of the ang_err pdf histogram. - bin_widths : 1d ndarray + bin_widths The bin widths of the ang_err pdf histogram. """ if log_true_e_idx < 0 or dec_idx < 0 or log_e_idx < 0 or psi_idx < 0: @@ -617,26 +638,28 @@ def get_ang_err_pdf(self, log_true_e_idx, dec_idx, log_e_idx, psi_idx): return (pdf, lower_bin_edges, upper_bin_edges, bin_widths) - def sample_log_e(self, rss, dec_idx, log_true_e_idxs): + def sample_log_e( + self, rss: RandomStateService, dec_idx: int, log_true_e_idxs: np.ndarray + ) -> tuple[np.ndarray, np.ndarray]: """Samples log energy values for the given source declination and true energy bins. Parameters ---------- - rss : instance of RandomStateService + rss The RandomStateService which should be used for drawing random numbers from. - dec_idx : int + dec_idx The index of the source declination bin. - log_true_e_idxs : 1d ndarray of int + log_true_e_idxs The bin indices of the true energy bins. Returns ------- - log_e_idx : 1d ndarray of int + log_e_idx The bin indices of the log_e pdf corresponding to the sampled log_e values. - log_e : 1d ndarray of float + log_e The sampled log_e values. """ n_evt = len(log_true_e_idxs) @@ -649,7 +672,7 @@ def sample_log_e(self, rss, dec_idx, log_true_e_idxs): b_size = np.count_nonzero(m) (pdf, low_bin_edges, up_bin_edges, bin_widths) = self.get_log_e_pdf(b_log_true_e_idx, dec_idx) - if pdf is None: + if pdf is None or low_bin_edges is None or up_bin_edges is None or bin_widths is None: log_e_idx[m] = -1 log_e[m] = np.nan continue @@ -662,28 +685,30 @@ def sample_log_e(self, rss, dec_idx, log_true_e_idxs): return (log_e_idx, log_e) - def sample_psi(self, rss, dec_idx, log_true_e_idxs, log_e_idxs): + def sample_psi( + self, rss: RandomStateService, dec_idx: int, log_true_e_idxs: np.ndarray, log_e_idxs: np.ndarray + ) -> tuple[np.ndarray, np.ndarray]: """Samples psi values for the given source declination, true energy bins, and log_e bins. Parameters ---------- - rss : instance of RandomStateService + rss The RandomStateService which should be used for drawing random numbers from. - dec_idx : int + dec_idx The index of the source declination bin. - log_true_e_idxs : 1d ndarray of int + log_true_e_idxs The bin indices of the true energy bins. - log_e_idxs : 1d ndarray of int + log_e_idxs The bin indices of the log_e bins. Returns ------- - psi_idx : 1d ndarray of int + psi_idx The bin indices of the psi pdf corresponding to the sampled psi values. - psi : 1d ndarray of float + psi The sampled psi values in radians. """ if len(log_true_e_idxs) != len(log_e_idxs): @@ -704,7 +729,7 @@ def sample_psi(self, rss, dec_idx, log_true_e_idxs, log_e_idxs): b_log_true_e_idx, dec_idx, bb_log_e_idx ) - if pdf is None: + if pdf is None or low_bin_edges is None or up_bin_edges is None or bin_widths is None: psi_idx[mm] = -1 psi[mm] = np.nan continue @@ -717,30 +742,37 @@ def sample_psi(self, rss, dec_idx, log_true_e_idxs, log_e_idxs): return (psi_idx, psi) - def sample_ang_err(self, rss, dec_idx, log_true_e_idxs, log_e_idxs, psi_idxs): + def sample_ang_err( + self, + rss: RandomStateService, + dec_idx: int, + log_true_e_idxs: np.ndarray, + log_e_idxs: np.ndarray, + psi_idxs: np.ndarray, + ) -> tuple[np.ndarray, np.ndarray]: """Samples ang_err values for the given source declination, true energy bins, log_e bins, and psi bins. Parameters ---------- - rss : instance of RandomStateService + rss The RandomStateService which should be used for drawing random numbers from. - dec_idx : int + dec_idx The index of the source declination bin. - log_true_e_idxs : 1d ndarray of int + log_true_e_idxs The bin indices of the true energy bins. - log_e_idxs : 1d ndarray of int + log_e_idxs The bin indices of the log_e bins. - psi_idxs : 1d ndarray of int + psi_idxs The bin indices of the psi bins. Returns ------- - ang_err_idx : 1d ndarray of int + ang_err_idx The bin indices of the angular error pdf corresponding to the sampled angular error values. - ang_err : 1d ndarray of float + ang_err The sampled angular error values in radians. """ if (len(log_true_e_idxs) != len(log_e_idxs)) and (len(log_e_idxs) != len(psi_idxs)): @@ -764,7 +796,7 @@ def sample_ang_err(self, rss, dec_idx, log_true_e_idxs, log_e_idxs, psi_idxs): b_log_true_e_idx, dec_idx, bb_log_e_idx, bbb_psi_idx ) - if pdf is None: + if pdf is None or low_bin_edges is None or up_bin_edges is None or bin_widths is None: ang_err_idx[mmm] = -1 ang_err[mmm] = np.nan continue diff --git a/skyllh/analyses/i3/publicdata_ps/time_dependent_ps.py b/skyllh/analyses/i3/publicdata_ps/time_dependent_ps.py index 28f74f8b55..8230f8f883 100644 --- a/skyllh/analyses/i3/publicdata_ps/time_dependent_ps.py +++ b/skyllh/analyses/i3/publicdata_ps/time_dependent_ps.py @@ -2,6 +2,8 @@ dataset. """ +from typing import cast + import numpy as np import skyllh @@ -27,6 +29,9 @@ clip_grl_start_times, create_energy_cut_spline, ) +from skyllh.core.analysis import ( + SingleSourceMultiDatasetLLHRatioAnalysis, +) from skyllh.core.analysis import ( SingleSourceMultiDatasetLLHRatioAnalysis as Analysis, ) @@ -39,6 +44,7 @@ from skyllh.core.config import ( Config, ) +from skyllh.core.dataset import Dataset from skyllh.core.event_selection import ( SpatialBoxEventSelectionMethod, ) @@ -51,6 +57,7 @@ PowerLawEnergyFluxProfile, SteadyPointlikeFFM, ) +from skyllh.core.llhratio import MultiDatasetTCLLHRatio, ZeroSigH0SingleDatasetTCLLHRatio from skyllh.core.logging import ( get_logger, ) @@ -72,6 +79,7 @@ Parameter, ParameterModelMapper, ) +from skyllh.core.pdf import PDF from skyllh.core.pdfratio import ( SigOverBkgPDFRatio, ) @@ -101,6 +109,9 @@ from skyllh.core.source_model import ( PointLikeSource, ) +from skyllh.core.storage import ( + DataFieldRecordArray, +) from skyllh.core.test_statistic import ( WilksTestStatistic, ) @@ -126,6 +137,7 @@ from skyllh.i3.config import ( add_icecube_specific_analysis_required_data_fields, ) +from skyllh.i3.dataset import I3DatasetData from skyllh.i3.livetime import ( I3Livetime, ) @@ -138,31 +150,29 @@ def create_signal_time_pdf( - cfg, - grl, - gauss=None, - box=None, -): + cfg: Config, + grl: np.ndarray, + gauss: dict | None = None, + box: dict | None = None, +) -> PDF: """Creates the signal time PDF, either a gaussian or a box shaped PDF. Parameters ---------- - cfg : instance of Config + cfg The instance of Config holding the local configuration. - grl : instance of numpy structured ndarray + grl The structured numpy ndarray holding the good-run-list data. - gauss : dict | None + gauss None or dictionary with {"mu": float, "sigma": float}. - box : dict | None + box None or dictionary with {"start": float, "stop": float}. Returns ------- - pdf : instance of PDF + pdf The created time PDF instance. """ - if (gauss is None) and (box is None): - raise TypeError('Either gauss or box have to be specified as time pdf.') livetime = I3Livetime.from_grl_data(grl_data=grl) @@ -170,6 +180,8 @@ def create_signal_time_pdf( time_flux_profile = GaussianTimeFluxProfile(t0=gauss['mu'], sigma_t=gauss['sigma'], cfg=cfg) elif box is not None: time_flux_profile = BoxTimeFluxProfile.from_start_and_stop_time(start=box['start'], stop=box['stop'], cfg=cfg) + else: + raise TypeError('Either gauss or box have to be specified as time pdf.') pdf = SignalTimePDF( cfg=cfg, @@ -181,19 +193,19 @@ def create_signal_time_pdf( def change_signal_time_pdf_of_llhratio_function( - ana, - gauss=None, - box=None, + ana: SingleSourceMultiDatasetLLHRatioAnalysis, + gauss: dict | None = None, + box: dict | None = None, ): """Changes the signal time PDF of the log-likelihood ratio function. Parameters ---------- - ana : instance of SingleSourceMultiDatasetLLHRatioAnalysis + ana The analysis instance. - gauss : dict | None + gauss None or dictionary with {"mu": float, "sigma": float}. - box : dict | None + box None or dictionary with {"start": float, "stop": float}. """ cfg = ana.cfg @@ -201,48 +213,52 @@ def change_signal_time_pdf_of_llhratio_function( time_sigpdf = create_signal_time_pdf(cfg=cfg, grl=grl, gauss=gauss, box=box) - pdfratio = ana.llhratio.llhratio_list[0].pdfratio + pdfratio = cast( + ZeroSigH0SingleDatasetTCLLHRatio, cast(MultiDatasetTCLLHRatio, ana.llhratio).llhratio_list[0] + ).pdfratio # pdfratio is an instance of PDFRatioProduct. # The first item is the PDF ratio product of the spatial and energy PDF # ratios. The second item is the time PDF ratio. - pdfratio.pdfratio2.sig_pdf = time_sigpdf + pdfratio.pdfratio2.sig_pdf = time_sigpdf # pyright: ignore[reportAttributeAccessIssue] # TODO: Change detector signal yield with flare livetime in sample # (1 / grl_norm in pdf), rebuild the histograms if it is changed. def get_energy_spatial_signal_over_background( - ana, - fitparam_values, - tl=None, -): + ana: SingleSourceMultiDatasetLLHRatioAnalysis, + fitparam_values: np.ndarray, + tl: TimeLord | None = None, +) -> np.ndarray: """Returns the signal over background ratio for (spatial_signal * energy_signal) / (spatial_background * energy_background). Parameters ---------- - ana : instance of SingleSourceMultiDatasetLLHRatioAnalysis + ana The analysis instance. - fitparam_values : instance of ndarray + fitparam_values The (N_fitparams,)-shaped numpy ndarray holding the values of the global fit parameters, e.g. ns and gamma. - tl : instance of TimeLord | None + tl The optional instance of TimeLord for measuring timing behavior. Returns ------- - ratio : 1d ndarray + ratio Product of spatial and energy signal over background pdfs. """ tdm = ana.tdm_list[0] - pdfratio = ana.llhratio.llhratio_list[0].pdfratio + pdfratio = cast( + ZeroSigH0SingleDatasetTCLLHRatio, cast(MultiDatasetTCLLHRatio, ana.llhratio).llhratio_list[0] + ).pdfratio # pdfratio is an instance of PDFRatioProduct. # The first item is the PDF ratio product of the spatial and energy PDF # ratios. The second item is the time PDF ratio. - pdfratio = pdfratio.pdfratio1 + pdfratio = pdfratio.pdfratio1 # pyright: ignore[reportAttributeAccessIssue] src_params_recarray = ana.pmm.create_src_params_recarray(gflp_values=fitparam_values) @@ -252,16 +268,16 @@ def get_energy_spatial_signal_over_background( def change_fluxmodel_gamma( - ana, - gamma, + ana: SingleSourceMultiDatasetLLHRatioAnalysis, + gamma: float, ): """Sets the given gamma value to the flux model of the single source. Parameters ---------- - ana : instance of SingleSourceMultiDatasetLLHRatioAnalysis + ana The analysis that should be used. - gamma : float + gamma Spectral index for the flux model. """ ana.shg_mgr.shg_list[0].fluxmodel.set_params({'gamma': gamma}) @@ -269,16 +285,16 @@ def change_fluxmodel_gamma( def change_time_flux_profile_params( - ana, - params, + ana: SingleSourceMultiDatasetLLHRatioAnalysis, + params: dict, ): """Changes the parameters of the source's time flux profile. Parameters ---------- - ana : instance of SingleSourceMultiDatasetLLHRatioAnalysis + ana The analysis that should be used. - params : dict + params The dictionary with the parameter names and values to be set. """ # Note: In the future the primary storage place for the time flux profile @@ -287,31 +303,31 @@ def change_time_flux_profile_params( def calculate_TS( - ana, - em_results, - rss, -): + ana: SingleSourceMultiDatasetLLHRatioAnalysis, + em_results: np.ndarray, + rss: RandomStateService, +) -> tuple[float | None, dict | None, np.ndarray | None]: """Calculate the best TS value from the expectation maximization gamma scan results. Parameters ---------- - ana : instance of SingleSourceMultiDatasetLLHRatioAnalysis + ana The analysis that should be used. - em_results : instance of structured ndarray + em_results The numpy structured ndarray holding the EM results (from the gamma scan). - rss : instance of RandomStateService + rss The instance of RandomStateService that should be used to generate random numbers from. Returns ------- - max_TS : float + max_TS The maximal TS value of all maximized time hypotheses. - best_em_result : instance of numpy structured ndarray + best_em_result The row of ``em_results`` that corresponds to the best fit. - best_fitparam_values : instance of numpy ndarray + best_fitparam_values The instance of numpy ndarray holding the fit parameter values of the overall best fit result. """ @@ -334,42 +350,42 @@ def calculate_TS( def run_gamma_scan_for_single_flare( - ana, - remove_time=None, - gamma_min=1, - gamma_max=5, - n_gamma=51, - ppbar=None, -): + ana: SingleSourceMultiDatasetLLHRatioAnalysis, + remove_time: float | None = None, + gamma_min: float = 1, + gamma_max: float = 5, + n_gamma: int = 51, + ppbar: ProgressBar | None = None, +) -> np.ndarray: """Runs ``em_fit`` for different gamma values in the signal energy PDF. Parameters ---------- - ana : instance of SingleSourceMultiDatasetLLHRatioAnalysis + ana The analysis that should be used. - remove_time : float + remove_time Time information of event that should be removed. - gamma_min : float + gamma_min Lower bound for gamma scan. - gamma_max : float + gamma_max Upper bound for gamma scan. - n_gamma : int + n_gamma Number of steps for gamma scan. - ppbar : instance of ProgressBar | None + ppbar The optional parent instance of ProgressBar. Returns ------- - em_results : instance of numpy structured ndarray + em_results The numpy structured ndarray with fields - gamma : float + gamma The spectral index value. - mu : float + mu The determined mean value of the gauss curve. - sigma : float + sigma The determined standard deviation of the gauss curve. - ns_em : float + ns_em The scaling factor of the flare. """ em_results_dt = [ @@ -404,28 +420,28 @@ def run_gamma_scan_for_single_flare( def unblind_single_flare( - ana, - remove_time=None, -): + ana: SingleSourceMultiDatasetLLHRatioAnalysis, + remove_time: float | None = None, +) -> tuple[float | None, dict | None, np.ndarray | None]: """Run EM for a single flare on unblinded data. Similar to the original analysis, remove the alert event. Parameters ---------- - ana : instance of SingleSourceMultiDatasetLLHRatioAnalysis + ana The analysis that should be used. - remove_time : float + remove_time Time of the event that should be removed. In the case of the TXS analysis: ``remove_time=TXS_0506_PLUS056_ALERT_TIME``. Returns ------- - max_TS : float + max_TS The maximal TS value of all maximized time hypotheses. - best_em_result : instance of numpy structured ndarray + best_em_result The EM result from the gamma scan corresponding to the best fit. - best_fitparam_values : instance of numpy ndarray + best_fitparam_values The instance of numpy ndarray holding the fit parameter values of the overall best fit result. """ @@ -441,83 +457,83 @@ def unblind_single_flare( def do_trial_with_em( - ana, - rss, - mean_n_sig=0, - gamma_src=2, - gamma_min=1, - gamma_max=5, - n_gamma=21, - gauss=None, - box=None, - tl=None, - ppbar=None, -): + ana: SingleSourceMultiDatasetLLHRatioAnalysis, + rss: RandomStateService, + mean_n_sig: float = 0, + gamma_src: float = 2, + gamma_min: float = 1, + gamma_max: float = 5, + n_gamma: int = 21, + gauss: dict | None = None, + box: dict | None = None, + tl: TimeLord | None = None, + ppbar: ProgressBar | None = None, +) -> np.ndarray: """Performs a trial using the expectation maximization algorithm. It runs a gamma scan and does the EM for each gamma value. Parameters ---------- - ana : instance of SingleSourceMultiDatasetLLHRatioAnalysis + ana The analysis instance that should be used to perform the trial. - rss : instance of RandomStateService + rss The instance of RandomStateService that should be used to generate random numbers. - mean_n_sig : float + mean_n_sig The mean number of signal events that should be generated. - gamma_src : float + gamma_src The spectral index of the source. - gamma_min : float + gamma_min Lower bound of the gamma scan. - gamma_max : float + gamma_max Upper bound of the gamma scan. - n_gamma : int + n_gamma Number of steps of the gamma scan. - gauss : dict | None + gauss Properties of the Gaussian time PDF. None or dictionary with {"mu": float, "sigma": float}. - box : dict | None + box Properties of the box time PDF. None or dictionary with {"start": float, "stop": float}. - tl : instance of TimeLord | None + tl The optional instance of TimeLord to measure timing information. - ppbar : instance of ProgressBar | None + ppbar The optional parent instance of ProgressBar. Returns ------- - trial : instance of structured ndarray + trial The numpy structured ndarray of length 1 with the following fields: - seed : numpy.int64 + seed The seed value used to generate the trial. - mean_n_sig : numpy.float64 + mean_n_sig The mean number of signal events of the trial. - n_sig : numpy.int64 + n_sig The actual number of signal events in the trial. - gamma_src : numpy.float64 + gamma_src The spectral index of the source. - mu_sig : numpy.float64 + mu_sig The mean value of the Gaussian time PDF of the source. - sigma_sig : numpy.float64 + sigma_sig The sigma value of the Gaussian time PDF of the source. - start_sig : numpy.float64 + start_sig The start time of the box time PDF of the source. - stop_sig : numpy.float64 + stop_sig The stop time of the box time PDF of the source. - ts : numpy.float64 + ts The test-statistic value of the trial. - ns_fit : numpy.float64 + ns_fit The fitted number of signal events. - ns_em : numpy.float64 + ns_em The scaling factor of the flare. - gamma_fit : numpy.float64 + gamma_fit The fitted spectral index of the trial. - gamma_em : numpy.float64 + gamma_em The spectral index of the best EM trial. - mu_fit : numpy.float64 + mu_fit The fitted mean value of the Gaussian time PDF. - sigma_fit : numpy.float64 + sigma_fit The fitted sigma value of the Gaussian time PDF. """ trial_dt = [ @@ -541,13 +557,19 @@ def do_trial_with_em( trial = np.empty((1,), dtype=trial_dt) (n_sig, n_events_list, events_list) = ana.generate_pseudo_data(rss=rss, mean_n_sig=mean_n_sig, tl=tl) - ana.initialize_trial(events_list, n_events_list) + ana.initialize_trial( + cast(list[DataFieldRecordArray], events_list), + cast(list[int | None], n_events_list), + ) em_results = run_gamma_scan_for_single_flare( ana=ana, gamma_min=gamma_min, gamma_max=gamma_max, n_gamma=n_gamma, ppbar=ppbar ) (max_ts, best_em_result, best_fitparams) = calculate_TS(ana=ana, em_results=em_results, rss=rss) + assert max_ts is not None + assert best_em_result is not None + assert best_fitparams is not None trial[0] = ( rss.seed, @@ -571,91 +593,91 @@ def do_trial_with_em( def do_trials_with_em( - ana, - n=1000, - ncpu=None, - seed=1, - mean_n_sig=0, - gamma_src=2, - gamma_min=1, - gamma_max=4, - n_gamma=21, - gauss=None, - box=None, - tl=None, - ppbar=None, -): + ana: SingleSourceMultiDatasetLLHRatioAnalysis, + n: int = 1000, + ncpu: int | None = None, + seed: int = 1, + mean_n_sig: float = 0, + gamma_src: float = 2, + gamma_min: float = 1, + gamma_max: float = 4, + n_gamma: int = 21, + gauss: dict | None = None, + box: dict | None = None, + tl: TimeLord | None = None, + ppbar: ProgressBar | None = None, +) -> np.ndarray: """Performs ``n_trials`` trials using the expectation maximization algorithm. For each trial it runs a gamma scan and does the EM for each gamma value. Parameters ---------- - ana : instance of SingleSourceMultiDatasetLLHRatioAnalysis + ana The analysis instance that should be used to perform the trials. - n : int + n The number of trials to generate. - ncpu : int | None + ncpu The number of CPUs to use to generate the trials. If set to ``None`` the configured default value will be used. - mean_n_sig : float + mean_n_sig The mean number of signal events that should be generated. - gamma_src : float + gamma_src The spectral index of the source. - gamma_min : float + gamma_min Lower bound of the gamma scan. - gamma_max : float + gamma_max Upper bound of the gamma scan. - n_gamma : int + n_gamma Number of steps of the gamma scan. - seed : int + seed The seed for the random number generator. - gauss : dict | None + gauss Properties of the Gaussian time PDF. None or dictionary with {"mu": float, "sigma": float}. - box : dict | None + box Properties of the box time PDF. None or dictionary with {"start": float, "stop": float}. - tl : instance of TimeLord | None + tl The optional instance of TimeLord to measure timing information. - ppbar : instance of ProgressBar | None + ppbar The optional parent instance of ProgressBar. Returns ------- - trials : instance of numpy structured ndarray + trials The numpy structured ndarray of length ``n_trials`` with the results for each trial. The array has the following fields: - seed : numpy.int64 + seed The seed value used to generate the trial. - mean_n_sig : numpy.float64 + mean_n_sig The mean number of signal events of the trial. - n_sig : numpy.int64 + n_sig The actual number of signal events in the trial. - gamma_src : numpy.float64 + gamma_src The spectral index of the source. - mu_sig : numpy.float64 + mu_sig The mean value of the Gaussian time PDF of the source. - sigma_sig : numpy.float64 + sigma_sig The sigma value of the Gaussian time PDF of the source. - start_sig : numpy.float64 + start_sig The start time of the box time PDF of the source. - stop_sig : numpy.float64 + stop_sig The stop time of the box time PDF of the source. - ts : numpy.float64 + ts The test-statistic value of the trial. - ns_fit : numpy.float64 + ns_fit The fitted number of signal events. - ns_em : numpy.float64 + ns_em The scaling factor of the flare. - gamma_fit : numpy.float64 + gamma_fit The fitted spectral index of the trial. - gamma_em : numpy.float64 + gamma_em The spectral index of the best EM trial. - mu_fit : numpy.float64 + mu_fit The fitted mean value of the Gaussian time PDF. - sigma_fit : numpy.float64 + sigma_fit The fitted sigma value of the Gaussian time PDF. """ rss = RandomStateService(seed=seed) @@ -696,105 +718,106 @@ def do_trials_with_em( trials = np.empty((n,), dtype=result.dtype) trials[i] = result[0] + assert trials is not None return trials def create_analysis( - cfg, - datasets, - source, + cfg: Config, + datasets: list[Dataset], + source: PointLikeSource, box=None, gauss=None, - refplflux_Phi0=1, - refplflux_E0=1e3, - refplflux_gamma=2.0, - ns_seed=10.0, - ns_min=0.0, - ns_max=1e3, - gamma_seed=3.0, - gamma_min=1.0, - gamma_max=5.0, - kde_smoothing=False, + refplflux_Phi0: float = 1, + refplflux_E0: float = 1e3, + refplflux_gamma: float = 2.0, + ns_seed: float = 10.0, + ns_min: float = 0.0, + ns_max: float = 1e3, + gamma_seed: float | None = 3.0, + gamma_min: float = 1.0, + gamma_max: float = 5.0, + kde_smoothing: bool = False, minimizer_impl='LBFGS', - compress_data=False, - keep_data_fields=None, - evt_sel_delta_angle_deg=10, - construct_bkg_generator=True, - construct_sig_generator=True, - tl=None, - ppbar=None, - logger_name=None, -): + compress_data: bool = False, + keep_data_fields: list[str] | None = None, + evt_sel_delta_angle_deg: float = 10, + construct_bkg_generator: bool = True, + construct_sig_generator: bool = True, + tl: TimeLord | None = None, + ppbar: ProgressBar | None = None, + logger_name: str | None = None, +) -> SingleSourceMultiDatasetLLHRatioAnalysis: """Creates the Analysis instance for this particular analysis. Parameters ---------- - cfg : instance of Config + cfg The instance of Config holding the local configuration. - datasets : list of Dataset instances + datasets The list of Dataset instances, which should be used in the analysis. - source : PointLikeSource instance + source The PointLikeSource instance defining the point source position. - box : None or dictionary with start, stop + box None if no box shaped time pdf, else dictionary of the format ``{'start': float, 'stop': float}``. - gauss : None or dictionary with mu, sigma + gauss None if no gaussian time pdf, else dictionary of the format ``{'mu': float, 'sigma': float}``. - refplflux_Phi0 : float + refplflux_Phi0 The flux normalization to use for the reference power law flux model. - refplflux_E0 : float + refplflux_E0 The reference energy to use for the reference power law flux model. - refplflux_gamma : float + refplflux_gamma The spectral index to use for the reference power law flux model. - ns_seed : float + ns_seed Value to seed the minimizer with for the ns fit. - ns_min : float + ns_min Lower bound for ns fit. - ns_max : float + ns_max Upper bound for ns fit. - gamma_seed : float | None + gamma_seed Value to seed the minimizer with for the gamma fit. If set to None, the refplflux_gamma value will be set as gamma_seed. - gamma_min : float + gamma_min Lower bound for gamma fit. - gamma_max : float + gamma_max Upper bound for gamma fit. - kde_smoothing : bool + kde_smoothing Deprecated: use of ``kde_smoothing=True`` is deprecated and will be removed in a future version. Apply a KDE-based smoothing to the data-driven background pdf. - Default: False. - minimizer_impl : str | "LBFGS" + Default + minimizer_impl Minimizer implementation to be used. Supported options are "LBFGS" (L-BFG-S minimizer used from the :mod:`scipy.optimize` module), or "minuit" (Minuit minimizer used by the :mod:`iminuit` module). - Default: "LBFGS". - compress_data : bool + Default + compress_data Flag if the data should get converted from float64 into float32. - keep_data_fields : list of str | None + keep_data_fields List of additional data field names that should get kept when loading the data. - evt_sel_delta_angle_deg : float + evt_sel_delta_angle_deg The delta angle in degrees for the event selection optimization methods. - construct_bkg_generator : bool + construct_bkg_generator Flag if the background generator should be constructed (``True``) or not (``False``). - construct_sig_generator : bool + construct_sig_generator Flag if the signal generator should be constructed (``True``) or not (``False``). - tl : TimeLord instance | None + tl The TimeLord instance to use to time the creation of the analysis. - ppbar : ProgressBar instance | None + ppbar The instance of ProgressBar for the optional parent progress bar. - logger_name : str | None + logger_name The name of the logger to be used. If set to ``None``, ``__name__`` will be used. Returns ------- - ana : instance of SingleSourceMultiDatasetLLHRatioAnalysis + ana The Analysis instance for this analysis. """ add_icecube_specific_analysis_required_data_fields(cfg) @@ -850,6 +873,8 @@ def create_analysis( param_ns = Parameter(name='ns', initial=ns_seed, valmin=ns_min, valmax=ns_max) # Define the fit parameter gamma. + if gamma_seed is None: + gamma_seed = refplflux_gamma if gamma_max > 4.0: logger.warning( 'You are allowing `gamma` values larger than 4.0. ' @@ -907,20 +932,26 @@ def create_analysis( # Add the data sets to the analysis. pbar = ProgressBar(len(datasets), parent=ppbar).start() for ds_idx, ds in enumerate(datasets): - data = ds.load_and_prepare_data( - keep_fields=keep_data_fields, dtc_dict=dtc_dict, dtc_except_fields=dtc_except_fields, tl=tl + data = cast( + I3DatasetData, + ds.load_and_prepare_data( + keep_fields=keep_data_fields, dtc_dict=dtc_dict, dtc_except_fields=dtc_except_fields, tl=tl + ), ) + assert data.exp is not None # Some runs might overlap slightly. So we need to clip those runs. - clip_grl_start_times(grl_data=data.grl) + clip_grl_start_times(grl_data=cast(np.ndarray, data.grl)) - livetime = I3Livetime.from_grl_data(grl_data=data.grl) + livetime = I3Livetime.from_grl_data(grl_data=cast(np.ndarray, data.grl)) sin_dec_binning = ds.get_binning_definition('sin_dec') log_energy_binning = ds.get_binning_definition('log_energy') # Create the spatial PDF ratio instance for this dataset. - spatial_sigpdf = RayleighPSFPointSourceSignalSpatialPDF(cfg=cfg, dec_range=np.arcsin(sin_dec_binning.range)) + spatial_sigpdf = RayleighPSFPointSourceSignalSpatialPDF( + cfg=cfg, dec_range=tuple(np.arcsin(sin_dec_binning.range)) + ) spatial_bkgpdf = DataBackgroundI3SpatialPDF(cfg=cfg, data_exp=data.exp, sin_dec_binning=sin_dec_binning) spatial_pdfratio = SigOverBkgPDFRatio(cfg=cfg, sig_pdf=spatial_sigpdf, bkg_pdf=spatial_bkgpdf) @@ -955,7 +986,7 @@ def create_analysis( start=livetime.time_start, stop=livetime.time_stop, cfg=cfg ), ) - time_sigpdf = create_signal_time_pdf(cfg=cfg, grl=data.grl, gauss=gauss, box=box) + time_sigpdf = create_signal_time_pdf(cfg=cfg, grl=cast(np.ndarray, data.grl), gauss=gauss, box=box) time_pdfratio = SigOverBkgPDFRatio( cfg=cfg, sig_pdf=time_sigpdf, @@ -987,6 +1018,7 @@ def create_analysis( cumulative_thr=ds.get_aux_data('cumulative_threshold'), ) + assert time_flux_profile is not None sig_generator = TimeDependentPDDatasetSignalGenerator( cfg=cfg, shg_mgr=shg_mgr, diff --git a/skyllh/analyses/i3/publicdata_ps/time_integrated_ps.py b/skyllh/analyses/i3/publicdata_ps/time_integrated_ps.py index c05dd0c9c9..7e04121b01 100644 --- a/skyllh/analyses/i3/publicdata_ps/time_integrated_ps.py +++ b/skyllh/analyses/i3/publicdata_ps/time_integrated_ps.py @@ -27,6 +27,9 @@ from skyllh.analyses.i3.publicdata_ps.utils import ( create_energy_cut_spline, ) +from skyllh.core.analysis import ( + SingleSourceMultiDatasetLLHRatioAnalysis, +) from skyllh.core.analysis import ( SingleSourceMultiDatasetLLHRatioAnalysis as Analysis, ) @@ -36,6 +39,7 @@ from skyllh.core.config import ( Config, ) +from skyllh.core.dataset import Dataset from skyllh.core.event_selection import ( SpatialBoxEventSelectionMethod, ) @@ -118,102 +122,102 @@ def create_analysis( - cfg, - datasets, - source, - refplflux_Phi0=1, - refplflux_E0=1e3, - refplflux_gamma=2.0, - refplflux_Ec=np.inf, - energy_range=None, - ns_seed=10.0, - ns_min=0.0, - ns_max=1e3, - gamma_seed=3.0, - gamma_min=1.0, - gamma_max=4.0, - kde_smoothing=False, - minimizer_impl='LBFGS', - minimizer_max_rep=100, - compress_data=False, - keep_data_fields=None, - evt_sel_delta_angle_deg=10, - construct_sig_generator=True, - tl=None, - ppbar=None, - logger_name=None, -): + cfg: Config, + datasets: list[Dataset], + source: PointLikeSource, + refplflux_Phi0: float = 1, + refplflux_E0: float = 1e3, + refplflux_gamma: float = 2.0, + refplflux_Ec: float = np.inf, + energy_range: tuple[float, float] | None = None, + ns_seed: float = 10.0, + ns_min: float = 0.0, + ns_max: float = 1e3, + gamma_seed: float | None = 3.0, + gamma_min: float = 1.0, + gamma_max: float = 4.0, + kde_smoothing: bool = False, + minimizer_impl: str = 'LBFGS', + minimizer_max_rep: int = 100, + compress_data: bool = False, + keep_data_fields: list[str] | None = None, + evt_sel_delta_angle_deg: float = 10, + construct_sig_generator: bool = True, + tl: TimeLord | None = None, + ppbar: ProgressBar | None = None, + logger_name: str | None = None, +) -> SingleSourceMultiDatasetLLHRatioAnalysis: """Creates the Analysis instance for this particular analysis. Parameters ---------- - cfg : instance of Config + cfg The instance of Config holding the local configuration. - datasets : list of Dataset instances + datasets The list of Dataset instances, which should be used in the analysis. - source : PointLikeSource instance + source The PointLikeSource instance defining the point source position. - refplflux_Phi0 : float + refplflux_Phi0 The flux normalization to use for the reference power law flux model. - refplflux_E0 : float + refplflux_E0 The reference energy to use for the reference power law flux model. - refplflux_gamma : float + refplflux_gamma The spectral index to use for the reference power law flux model. - refplflux_Ec: float, + refplflux_Ec The cutoff energy for the cutoff power law flux model. - energy_range: 2-element tuple of float (low energy, high energy) | None + energy_range The energy range for signal generation. Both low and high energies are given in GeV. If set to ``None``, the entire energy range of the dataset is used. - ns_seed : float + ns_seed Value to seed the minimizer with for the ns fit. - ns_min : float + ns_min Lower bound for ns fit. - ns_max : float + ns_max Upper bound for ns fit. - gamma_seed : float | None + gamma_seed Value to seed the minimizer with for the gamma fit. If set to None, the refplflux_gamma value will be set as gamma_seed. - gamma_min : float + gamma_min Lower bound for gamma fit. - gamma_max : float + gamma_max Upper bound for gamma fit. - kde_smoothing : bool + kde_smoothing Deprecated: use of ``kde_smoothing=True`` is deprecated and will be removed in a future version. Apply a KDE-based smoothing to the data-driven background pdf. - Default: False. - minimizer_impl : str + Default + minimizer_impl Minimizer implementation to be used. Supported options are ``"LBFGS"`` (L-BFG-S minimizer used from the :mod:`scipy.optimize` module), or ``"minuit"`` (Minuit minimizer used by the :mod:`iminuit` module). Default: "LBFGS". - minimizer_max_rep : int + minimizer_max_rep In case the minimization process did not converge at the first time this option specifies the maximum number of repetitions with different initials. Default is 100. - compress_data : bool + compress_data Flag if the data should get converted from float64 into float32. - keep_data_fields : list of str | None + keep_data_fields List of additional data field names that should get kept when loading the data. - evt_sel_delta_angle_deg : float + evt_sel_delta_angle_deg The delta angle in degrees for the event selection optimization methods. - construct_sig_generator : bool + construct_sig_generator Flag if the signal generator should be constructed (``True``) or not (``False``). - tl : TimeLord instance | None + tl The TimeLord instance to use to time the creation of the analysis. - ppbar : ProgressBar instance | None + ppbar The instance of ProgressBar for the optional parent progress bar. - logger_name : str | None + logger_name The name of the logger to be used. If set to ``None``, ``__name__`` will be used. Returns ------- - ana : instance of SingleSourceMultiDatasetLLHRatioAnalysis + ana The Analysis instance for this analysis. """ add_icecube_specific_analysis_required_data_fields(cfg) @@ -268,6 +272,8 @@ def create_analysis( param_ns = Parameter(name='ns', initial=ns_seed, valmin=ns_min, valmax=ns_max) # Define the fit parameter gamma. + if gamma_seed is None: + gamma_seed = refplflux_gamma if gamma_max > 4.0: logger.warning( 'You are allowing `gamma` values larger than 4.0. ' @@ -338,12 +344,15 @@ def create_analysis( data = ds.load_and_prepare_data( keep_fields=keep_data_fields, dtc_dict=dtc_dict, dtc_except_fields=dtc_except_fields, tl=tl ) + assert data.exp is not None sin_dec_binning = ds.get_binning_definition('sin_dec') log_energy_binning = ds.get_binning_definition('log_energy') # Create the spatial PDF ratio instance for this dataset. - spatial_sigpdf = RayleighPSFPointSourceSignalSpatialPDF(cfg=cfg, dec_range=np.arcsin(sin_dec_binning.range)) + spatial_sigpdf = RayleighPSFPointSourceSignalSpatialPDF( + cfg=cfg, dec_range=tuple(np.arcsin(sin_dec_binning.range)) + ) spatial_bkgpdf = DataBackgroundI3SpatialPDF(cfg=cfg, data_exp=data.exp, sin_dec_binning=sin_dec_binning) spatial_pdfratio = SigOverBkgPDFRatio(cfg=cfg, sig_pdf=spatial_sigpdf, bkg_pdf=spatial_bkgpdf) diff --git a/skyllh/analyses/i3/publicdata_ps/time_integrated_ps_function_energy_spectrum.py b/skyllh/analyses/i3/publicdata_ps/time_integrated_ps_function_energy_spectrum.py index 874167373b..2fdc490fe7 100644 --- a/skyllh/analyses/i3/publicdata_ps/time_integrated_ps_function_energy_spectrum.py +++ b/skyllh/analyses/i3/publicdata_ps/time_integrated_ps_function_energy_spectrum.py @@ -3,6 +3,8 @@ energy event PDF. """ +from typing import cast + import numpy as np from scipy.interpolate import BSpline, splrep @@ -27,6 +29,9 @@ from skyllh.analyses.i3.publicdata_ps.utils import ( create_energy_cut_spline, ) +from skyllh.core.analysis import ( + SingleSourceMultiDatasetLLHRatioAnalysis, +) from skyllh.core.analysis import ( SingleSourceMultiDatasetLLHRatioAnalysis as Analysis, ) @@ -36,11 +41,13 @@ from skyllh.core.config import ( Config, ) +from skyllh.core.dataset import Dataset from skyllh.core.event_selection import ( SpatialBoxEventSelectionMethod, ) from skyllh.core.flux_model import ( EpeakFunctionEnergyProfile, + FactorizedFluxModel, SteadyPointlikeFFM, ) from skyllh.core.logging import ( @@ -83,9 +90,11 @@ SourceHypoGroup, SourceHypoGroupManager, ) +from skyllh.core.source_model import PointLikeSource from skyllh.core.test_statistic import ( WilksTestStatistic, ) +from skyllh.core.timing import TimeLord from skyllh.core.trialdata import ( TrialDataManager, ) @@ -108,37 +117,41 @@ cfg = Config() -def set_epeak(analysis, e_peak): +def set_epeak(analysis: SingleSourceMultiDatasetLLHRatioAnalysis, e_peak: float) -> None: """Change the peak energy. The shape stays the same but the spectrum is moved to higher/lower energies. Parameters ---------- - analysis : instance of SingleSourceMultiDatasetLLHRatioAnalysis + analysis Analysis instance with the defined flux model and signal generator - e_peak : float + e_peak Peak energy of the flux model (this defines the reference flux) """ - analysis.shg_mgr.get_fluxmodel_by_src_idx(0).energy_profile.e_peak = e_peak + cast( + EpeakFunctionEnergyProfile, + cast(FactorizedFluxModel, analysis.shg_mgr.get_fluxmodel_by_src_idx(0)).energy_profile, + ).e_peak = e_peak + assert analysis.sig_generator is not None analysis.sig_generator.change_shg_mgr(analysis.shg_mgr) -def flux_from_ns(analysis, e_peak, ns): +def flux_from_ns(analysis: SingleSourceMultiDatasetLLHRatioAnalysis, e_peak: float, ns: float) -> float | np.ndarray: """Get the flux at e_peak for a certain flux model (defined by e_peak) for a mean number of signal neutrinos ns Parameters ---------- - analysis : instance of SingleSourceMultiDatasetLLHRatioAnalysis + analysis Analysis instance with the defined flux model and signal generator - e_peak : float + e_peak Peak energy of the flux model (this defines the reference flux) - ns : float + ns Mean number of detected signal neutrinos Returns ------- - flux : float + flux Flux (dN / dE) in (GeV cm^2 s)^-1 at peak energy. """ # set the fluxmodel to e_peak @@ -146,24 +159,31 @@ def flux_from_ns(analysis, e_peak, ns): scaling_factor = analysis.calculate_fluxmodel_scaling_factor() * ns - return analysis.shg_mgr.get_fluxmodel_by_src_idx(0).energy_profile(E=10**e_peak).squeeze() * scaling_factor + return ( + float( + cast(FactorizedFluxModel, analysis.shg_mgr.get_fluxmodel_by_src_idx(0)) + .energy_profile(E=10**e_peak) + .squeeze() + ) + * scaling_factor + ) -def ns_from_flux(analysis, e_peak, flux): +def ns_from_flux(analysis: SingleSourceMultiDatasetLLHRatioAnalysis, e_peak: float, flux: float) -> float: """Get the mean number of signal neutrinos ns for a certain flux model (defined by e_peak) for a flux at e_peak (1/GeV/cm2/s). Parameters ---------- - analysis : instance of SingleSourceMultiDatasetLLHRatioAnalysis + analysis Analysis instance with the defined flux model and signal generator - e_peak : float + e_peak Peak energy of the flux model (this defines the reference flux) - flux : float + flux Flux in 1/(GeV cm2 s) Returns ------- - ns : float + ns Mean number of signal neutrinos """ @@ -171,101 +191,102 @@ def ns_from_flux(analysis, e_peak, flux): set_epeak(analysis, e_peak) # reference flux at e_peak - reference_flux = analysis.shg_mgr.get_fluxmodel_by_src_idx(0).energy_profile(E=10**e_peak).squeeze() + reference_flux = ( + cast(FactorizedFluxModel, analysis.shg_mgr.get_fluxmodel_by_src_idx(0)).energy_profile(E=10**e_peak).squeeze() + ) scaling_factor = flux / reference_flux scaling_factor_norm = analysis.calculate_fluxmodel_scaling_factor() - return scaling_factor / scaling_factor_norm + return float(scaling_factor / scaling_factor_norm) def create_analysis( - cfg, - datasets, - source, - source_energies, - source_energy_spectrum, - refplflux_Phi0=1, - ns_seed=10.0, - ns_min=0.0, - ns_max=1e3, - e_peak_signal=5, - e_peak_seed=3, - e_peak_min=1.06, - e_peak_max=10.06, - kde_smoothing=False, - minimizer_impl='minuit', - compress_data=False, - keep_data_fields=None, - evt_sel_delta_angle_deg=10, - construct_sig_generator=True, - tl=None, - ppbar=None, - logger_name=None, -): + cfg: Config, + datasets: list[Dataset], + source: PointLikeSource, + source_energies: np.ndarray, + source_energy_spectrum: np.ndarray, + refplflux_Phi0: float = 1, + ns_seed: float = 10.0, + ns_min: float = 0.0, + ns_max: float = 1e3, + e_peak_signal: float = 5, + e_peak_seed: float = 3, + e_peak_min: float = 1.06, + e_peak_max: float = 10.06, + kde_smoothing: bool = False, + minimizer_impl: str = 'minuit', + compress_data: bool = False, + keep_data_fields: list[str] | None = None, + evt_sel_delta_angle_deg: float = 10, + construct_sig_generator: bool = True, + tl: TimeLord | None = None, + ppbar: ProgressBar | None = None, + logger_name: str | None = None, +) -> SingleSourceMultiDatasetLLHRatioAnalysis: """Creates the Analysis instance for this particular analysis. Parameters ---------- - cfg : instance of Config + cfg The instance of Config holding the local configuration. - datasets : list of Dataset instances + datasets The list of Dataset instances, which should be used in the analysis. - source : PointLikeSource instance + source The PointLikeSource instance defining the point source position. - source_energies : numpy array + source_energies Energies in GeV for which source_energy_spectrum is given - source_energy_spectrum : numpy array + source_energy_spectrum The energy spectrum in GeV / cm^2 / s - refplflux_Phi0 : float + refplflux_Phi0 The flux normalization to use for the reference power law flux model. - ns_seed : float + ns_seed Value to seed the minimizer with for the ns fit. - ns_min : float + ns_min Lower bound for ns fit. - ns_max : float + ns_max Upper bound for ns fit. - e_peak_signal : float + e_peak_signal Default energy peak value for the signal generator. - e_peak_seed : float + e_peak_seed Seed value for minimizer for fitting energy peak. - e_peak_min : float + e_peak_min Lower bound for energy peak fit. - e_peak_max : float + e_peak_max Upper bound for energy peak fit, - kde_smoothing : bool + kde_smoothing Deprecated: use of ``kde_smoothing=True`` is deprecated and will be removed in a future version. Apply a KDE-based smoothing to the data-driven background pdf. - Default: False. - minimizer_impl : str + minimizer_impl Minimizer implementation to be used. Supported options are ``"LBFGS"`` (L-BFG-S minimizer used from the :mod:`scipy.optimize` module), or ``"minuit"`` (Minuit minimizer used by the :mod:`iminuit` module). - Default: "LBFGS". - compress_data : bool + Default + compress_data Flag if the data should get converted from float64 into float32. - keep_data_fields : list of str | None + keep_data_fields List of additional data field names that should get kept when loading the data. - evt_sel_delta_angle_deg : float + evt_sel_delta_angle_deg The delta angle in degrees for the event selection optimization methods. - construct_sig_generator : bool + construct_sig_generator Flag if the signal generator should be constructed (``True``) or not (``False``). - tl : TimeLord instance | None + tl The TimeLord instance to use to time the creation of the analysis. - ppbar : ProgressBar instance | None + ppbar The instance of ProgressBar for the optional parent progress bar. - logger_name : str | None + logger_name The name of the logger to be used. If set to ``None``, ``__name__`` will be used. Returns ------- - ana : instance of SingleSourceMultiDatasetLLHRatioAnalysis + ana The Analysis instance for this analysis. """ @@ -296,7 +317,7 @@ def create_analysis( energy_spectrum_spline = splrep(source_energies, source_energy_spectrum / source_energies / source_energies, k=1) - spline_eval = BSpline(*energy_spectrum_spline) + spline_eval = BSpline(*energy_spectrum_spline) # pyright: ignore[reportArgumentType] e_peak = np.log10(source_energies[np.argmax(source_energy_spectrum)]) @@ -382,12 +403,15 @@ def create_analysis( # compress=compress_data, tl=tl, ) + assert data.exp is not None sin_dec_binning = ds.get_binning_definition('sin_dec') log_energy_binning = ds.get_binning_definition('log_energy') # Create the spatial PDF ratio instance for this dataset. - spatial_sigpdf = RayleighPSFPointSourceSignalSpatialPDF(cfg=cfg, dec_range=np.arcsin(sin_dec_binning.range)) + spatial_sigpdf = RayleighPSFPointSourceSignalSpatialPDF( + cfg=cfg, dec_range=tuple(np.arcsin(sin_dec_binning.range)) + ) spatial_bkgpdf = DataBackgroundI3SpatialPDF(cfg=cfg, data_exp=data.exp, sin_dec_binning=sin_dec_binning) spatial_pdfratio = SigOverBkgPDFRatio(cfg=cfg, sig_pdf=spatial_sigpdf, bkg_pdf=spatial_bkgpdf) diff --git a/skyllh/analyses/i3/publicdata_ps/utils.py b/skyllh/analyses/i3/publicdata_ps/utils.py index 2a416b1efe..120e626d2b 100644 --- a/skyllh/analyses/i3/publicdata_ps/utils.py +++ b/skyllh/analyses/i3/publicdata_ps/utils.py @@ -4,7 +4,10 @@ from skyllh.core.binning import ( get_bincenters_from_binedges, ) +from skyllh.core.dataset import Dataset from skyllh.core.flux_model import EnergyFluxProfile +from skyllh.core.random import RandomStateService +from skyllh.core.storage import DataFieldRecordArray class FctSpline1D: @@ -14,17 +17,17 @@ class from scipy. The evaluate the spline, use the ``__call__`` method. """ - def __init__(self, f, x_binedges, norm=False, **kwargs): + def __init__(self, f: np.ndarray, x_binedges: np.ndarray, norm: bool = False, **kwargs): """Creates a new 1D function spline using the PchipInterpolator class from scipy. Parameters ---------- - f : (n_x,)-shaped 1D numpy ndarray - The numpy ndarray holding the function values at the bin centers. - x_binedges : (n_x+1,)-shaped 1D numpy ndarray - The numpy ndarray holding the bin edges of the x-axis. - norm : bool + f + The (n_x,)-shaped 1D numpy ndarray holding the function values at the bin centers. + x_binedges + The (n_x+1,)-shaped 1D numpy ndarray holding the bin edges of the x-axis. + norm Whether to precalculate and save normalization internally. """ super().__init__(**kwargs) @@ -44,22 +47,22 @@ class from scipy. # We choose not to extrapolate out-of-range values. self.norm = float(self.spl_f.integrate(x[0], x[-1])) - def __call__(self, x, oor_value=0): + def __call__(self, x: np.ndarray, oor_value: float = 0) -> np.ndarray: """Evaluates the spline at the given x values. For x-values outside the spline's range, the oor_value is returned. Parameters ---------- - x : (n_x,)-shaped 1D numpy ndarray - The numpy ndarray holding the x values at which the spline should + x + The (n_x,)-shaped 1D numpy ndarray holding the x values at which the spline should get evaluated. - oor_value : float + oor_value The value for out-of-range (oor) coordinates. Returns ------- - f : (n_x,)-shaped 1D numpy ndarray - The numpy ndarray holding the evaluated values of the spline. + f + The (n_x,)-shaped 1D numpy ndarray holding the evaluated values of the spline. """ f = self.spl_f(x) f = np.where(np.isnan(f), oor_value, f) @@ -81,18 +84,18 @@ class from scipy. The evaluate the spline, use the ``__call__`` method. """ - def __init__(self, f, x_binedges, y_binedges, **kwargs): + def __init__(self, f: np.ndarray, x_binedges: np.ndarray, y_binedges: np.ndarray, **kwargs): """Creates a new 2D function spline using the RectBivariateSpline class from scipy. Parameters ---------- - f : (n_x, n_y)-shaped 2D numpy ndarray - The numpy ndarray holding the function values at the bin centers. - x_binedges : (n_x+1,)-shaped 1D numpy ndarray - The numpy ndarray holding the bin edges of the x-axis. - y_binedges : (n_y+1,)-shaped 1D numpy ndarray - The numpy ndarray holding the bin edges of the y-axis. + f + The (n_x, n_y)-shaped 2D numpy ndarray holding the function values at the bin centers. + x_binedges + The (n_x+1,)-shaped 1D numpy ndarray holding the bin edges of the x-axis. + y_binedges + The (n_y+1,)-shaped 1D numpy ndarray holding the bin edges of the y-axis. """ super().__init__(**kwargs) @@ -125,6 +128,9 @@ class from scipy. self._prepare_quadrature() def _prepare_quadrature(self, n=128): + """Pre-computes the Gauss-Legendre quadrature nodes and weights used for + renormalizing the spline evaluation. + """ gx, gw = np.polynomial.legendre.leggauss(n) self._qx = 0.5 * (self.x_max - self.x_min) * gx + 0.5 * (self.x_max + self.x_min) self._qw = 0.5 * (self.x_max - self.x_min) * gw @@ -133,10 +139,16 @@ def _prepare_quadrature(self, n=128): @staticmethod def _pow10(arr): + """Computes ``10 ** arr`` using an optimized exp-based version that is + about 3x faster than ``np.power(10, arr)``. + """ # Alternative optimized version of np.power(10, arr), ~3x faster. return np.exp(FctSpline2D._LOG10 * arr) def _mask_oor_axes(self, x, y): + """Returns the masks selecting the ``x`` and ``y`` values that are out of + range of the spline's x- and y-axis, respectively. + """ m_x = (x < self.x_min) | (x > self.x_max) m_y = (y < self.y_min) | (y > self.y_max) return m_x, m_y @@ -165,7 +177,7 @@ def _eval_grid(self, x, y): f_sorted = f_sorted[:, np.argsort(iy)] return f_sorted - def _renorm_per_y_grid(self, f2d, y, *, in_user_order=True): + def _renorm_per_y_grid(self, f2d, y, *, in_user_order: bool = True): """Renormalize columns so ∫_x f(x, y) dx = 1 (grid=True).""" y = np.asarray(y) # For renorm we can evaluate on (qx, y) with grid=True (expects sorted y). @@ -193,31 +205,33 @@ def _renorm_per_y_pairs(self, x, y, f): f /= Z[inv] return f - def __call__(self, x, y, oor_value=0, grid=False, renorm=True): + def __call__( + self, x: np.ndarray, y: np.ndarray, oor_value: float = 0, grid: bool = False, renorm: bool = True + ) -> np.ndarray: """Evaluates the spline at the given coordinates. For coordinates outside the spline's range, the oor_value is returned. Parameters ---------- - x : (n_x,)-shaped 1D numpy ndarray - The numpy ndarray holding the x values at which the spline should + x + The (n_x,)-shaped 1D numpy ndarray holding the x values at which the spline should get evaluated. - y : (n_y,)-shaped 1D numpy ndarray - The numpy ndarray holding the y values at which the spline should + y + The (n_y,)-shaped 1D numpy ndarray holding the y values at which the spline should get evaluated. - oor_value : float | 0 + oor_value The value for out-of-range (oor) coordinates. - grid : bool | False + grid Whether the interpolation should return a 2D numpy array or a 1D sequence of values. - renorm : bool | True + renorm Whether to renormalize the histogram along the x axis for each y-value. Useful when constructing the background energy PDF. Returns ------- - f : numpy ndarray - The numpy ndarray holding the evaluated values of the spline. + f + The (n_x,)-shaped 1D numpy ndarray holding the evaluated values of the spline. """ x = np.asarray(x) y = np.asarray(y) @@ -251,19 +265,19 @@ def __call__(self, x, y, oor_value=0, grid=False, renorm=True): return f2d -def clip_grl_start_times(grl_data): +def clip_grl_start_times(grl_data: np.ndarray): """Make sure that the start time of a run is not smaller than the stop time of the previous run. Parameters ---------- - grl_data : instance of numpy structured ndarray + grl_data The numpy structured ndarray of length N_runs, with the following fields: - start : float + start The start time of the run. - stop : float + stop The stop time of the run. """ start = grl_data['start'] @@ -275,26 +289,28 @@ def clip_grl_start_times(grl_data): grl_data['start'][1:] = new_start -def psi_to_dec_and_ra(rss, src_dec, src_ra, psi): +def psi_to_dec_and_ra( + rss: RandomStateService, src_dec: float, src_ra: float, psi: np.ndarray +) -> tuple[np.ndarray, np.ndarray]: """Generates random declinations and right-ascension coordinates for the given source location and opening angle `psi`. Parameters ---------- - rss : instance of RandomStateService + rss The instance of RandomStateService to use for drawing random numbers. - src_dec : float + src_dec The declination of the source in radians. - src_ra : float + src_ra The right-ascension of the source in radians. - psi : 1d ndarray of float + psi The opening-angle values in radians. Returns ------- - dec : 1d ndarray of float + dec The declination values. - ra : 1d ndarray of float + ra The right-ascension values. """ @@ -332,7 +348,7 @@ def psi_to_dec_and_ra(rss, src_dec, src_ra, psi): return (dec, ra) -def create_energy_cut_spline(ds, exp_data, spl_smooth, cumulative_thr=0): +def create_energy_cut_spline(ds: Dataset, exp_data: DataFieldRecordArray, spl_smooth: float, cumulative_thr: float = 0): """Create the spline for the declination-dependent energy cut that the signal generator needs for injection in the southern sky. Cut bins which do not exceed the defined `cumulative_thr` threshold @@ -341,19 +357,19 @@ def create_energy_cut_spline(ds, exp_data, spl_smooth, cumulative_thr=0): Parameters ---------- - ds : instance of Dataset + ds The instance of Dataset for which the spline should be calculated. - exp_data : instance of DataFieldRecordArray + exp_data The array containing the experimental data for dataset `ds`. - spl_smooth : float + spl_smooth - cumulative_thr : float + cumulative_thr Defaults to 0 that corresponds to no cut. Returns ------- - spline : instance of scipy.interpolate.UnivariateSpline + spline """ data_exp = exp_data.copy(keep_fields=['sin_dec', 'log_energy']) @@ -510,7 +526,7 @@ def _reco_energy_counts_per_ds(ds, sm, flux, dec, reco_e_edges, Phi0, energy_ran def compute_expected_reco_energy_counts( - datasets, flux, dec, Phi0, time_integral=False, livetimes=None, energy_range=None + datasets, flux, dec, Phi0, time_integral: bool = False, livetimes=None, energy_range=None ): """Returns the expected distribution of events in reconstructed energy for a given flux and declination. If a list of datasets is provided, it returns the sum of all datasets contributions. @@ -525,7 +541,7 @@ def compute_expected_reco_energy_counts( The declination in radians. Phi0 : float The flux normalization factor. - time_integral : bool + time_integral Whether to return the total expected counts (True) or the expected counts per second (False). Default is False (counts per second). livetimes : (len(datasets),)-iterable of float and None | None @@ -597,6 +613,7 @@ def compute_expected_reco_energy_counts( else: data = ds.load_data() livetime = data.livetime + assert livetime is not None counts_total += counts_per_sec * livetime * 24 * 3600 # Convert days to seconds else: counts_total += counts_per_sec diff --git a/skyllh/core/analysis.py b/skyllh/core/analysis.py index 25b634e00d..2484d288f5 100644 --- a/skyllh/core/analysis.py +++ b/skyllh/core/analysis.py @@ -1,6 +1,7 @@ """The analysis module provides classes for pre-defined analyses.""" import abc +from typing import cast import numpy as np from astropy import units @@ -22,11 +23,13 @@ from skyllh.core.llhratio import ( LLHRatio, MultiDatasetTCLLHRatio, + SingleDatasetTCLLHRatio, ZeroSigH0SingleDatasetTCLLHRatio, ) from skyllh.core.logging import ( get_logger, ) +from skyllh.core.minimizer import Minimizer from skyllh.core.multiproc import ( get_ncpu, parallelize, @@ -38,6 +41,7 @@ PDFRatio, SourceWeightedPDFRatio, ) +from skyllh.core.progressbar import ProgressBar from skyllh.core.py import ( classname, issequenceof, @@ -69,9 +73,7 @@ from skyllh.core.test_statistic import ( TestStatistic, ) -from skyllh.core.timing import ( - TaskTimer, -) +from skyllh.core.timing import TaskTimer, TimeLord from skyllh.core.trialdata import ( TrialDataManager, ) @@ -88,28 +90,36 @@ class Analysis( overall analysis interface how to setup and run an analysis. """ - def __init__(self, shg_mgr, pmm, test_statistic, bkg_generator_cls=None, sig_generator_cls=None, **kwargs): + def __init__( + self, + shg_mgr: SourceHypoGroupManager, + pmm: ParameterModelMapper, + test_statistic: TestStatistic, + bkg_generator_cls: type[MultiDatasetBackgroundGenerator] | None = None, + sig_generator_cls: type[MultiDatasetSignalGenerator] | None = None, + **kwargs, + ): """Constructor of the analysis base class. Parameters ---------- - shg_mgr : instance of SourceHypoGroupManager + shg_mgr The instance of SourceHypoGroupManager, which defines the groups of source hypotheses, their flux model, and their detector signal yield implementation method. - pmm : instance of ParameterModelMapper + pmm The ParameterModelMapper instance managing the global set of parameters and their relation to individual models, e.g. sources. - test_statistic : TestStatistic instance + test_statistic The TestStatistic instance that defines the test statistic function of the analysis. - bkg_generator_cls : class of MultiDatasetBackgroundGenerator | None + bkg_generator_cls The background generator class used to create the background generator instance for multiple datasets. If set to ``None``, the :class:`skyllh.core.background_generator.MultiDatasetBackgroundGenerator` class is used. - sig_generator_cls : class of MultiDatasetSignalGenerator | None + sig_generator_cls The signal generator class used to create the signal generator instance for multiple datasets. If set to ``None``, the @@ -136,7 +146,7 @@ class is used. self._bkg_generator_list = [] self._bkg_generator = None self._sig_generator_list = [] - self._sig_generator = None + self._sig_generator: SignalGenerator | None = None self.sig_gen_energy_range = None self.sig_gen_energy_range_is_set = False @@ -377,6 +387,9 @@ def _iter_energy_range_capable(self): yield generator def _apply_energy_range(self, value): + """Applies the given energy range to all energy-range-capable components + of the analysis. + """ for component in self._iter_energy_range_capable(): component.energy_range = value @@ -390,7 +403,7 @@ def _has_explicit_energy_range_cut(self): # explicit cut from the full-range default. Mirror the same pattern # used in MultiDatasetSignalGenerator.fluxmodel_scaling_factor. if hasattr(component, '_input_range'): - if component._input_range is not None: + if getattr(component, '_input_range', None) is not None: return True elif component.energy_range is not None: return True @@ -455,7 +468,7 @@ def total_livetime(self): def construct_services( self, - ppbar=None, + ppbar: ProgressBar | None = None, ): """Constructs the following services: @@ -465,7 +478,7 @@ def construct_services( Parameters ---------- - ppbar : instance of ProgressBar | None + ppbar The instance of ProgressBar of the optional parent progress bar. """ self.detsigyield_service = DetSigYieldService( @@ -485,36 +498,36 @@ def construct_services( def add_dataset( self, - dataset, - data, - tdm=None, - event_selection_method=None, - bkg_generator=None, - sig_generator=None, + dataset: Dataset, + data: DatasetData, + tdm: TrialDataManager | None = None, + event_selection_method: EventSelectionMethod | None = None, + bkg_generator: BackgroundGenerator | None = None, + sig_generator: SignalGenerator | None = None, ): """Adds the given dataset to the list of datasets for this analysis. Parameters ---------- - dataset : instance of Dataset + dataset The Dataset instance that should get added. - data : instance of DatasetData + data The DatasetData instance holding the original (prepared) data of the dataset. - tdm : instance of TrialDataManager | None + tdm The TrialDataManager instance managing the trial data and additional data fields of the data set. If set to None, it means that no additional data fields are defined. - event_selection_method : instance of EventSelectionMethod | None + event_selection_method The instance of EventSelectionMethod to use to select only signal-like events from the data. All other events will be treated as pure background events. This reduces the amount of log-likelihood-ratio function evaluations. If set to None, all events will be evaluated. - bkg_generator : instance of BackgroundGenerator | None + bkg_generator The optional instance of BackgroundGenerator, which should be used to generate background events for this particular dataset. - sig_generator : instance of SignalGenerator | None + sig_generator The optional instance of SignalGenerator, which should be used to generate signal events for this particular dataset. """ @@ -542,8 +555,8 @@ def add_dataset( if bkg_generator is not None and not isinstance(bkg_generator, BackgroundGenerator): raise TypeError( - 'The bkg_generator argument must be None or an instance of' - 'BackgroundGenerator! ' + 'The bkg_generator argument must be None or an instance of ' + 'BackgroundGenerator!' f'Its current type is {classname(bkg_generator)}!' ) @@ -564,7 +577,7 @@ def add_dataset( if sig_generator is not None and self.sig_gen_energy_range_is_set and isinstance(sig_generator, HasEnergyRange): sig_generator.energy_range = self.sig_gen_energy_range - def get_livetime(self, dataset_key=None, unit=None): + def get_livetime(self, dataset_key: int | str | None = None, unit: units.Unit | None = None): """Retrieves the numeric livetime of the given dataset in the specified unit. The dataset can be specified either through its index or its name. If no dataset is specified, the total livetime, i.e. the sum of the @@ -572,11 +585,11 @@ def get_livetime(self, dataset_key=None, unit=None): Parameters ---------- - dataset_key : int | str | None + dataset_key The index or name of the dataset for which the livetime should get retrieved. If set to ``None``, the total livetime of all datasets will be returned. - unit : instance of astropy.units.Unit | None + unit The time unit in which the livetime should be returned. If set to ``None``, ``astropy.units.day`` will be used. """ @@ -606,17 +619,17 @@ def get_livetime(self, dataset_key=None, unit=None): return livetime - def calculate_test_statistic(self, log_lambda, fitparam_values, **kwargs): + def calculate_test_statistic(self, log_lambda: float, fitparam_values: np.ndarray, **kwargs) -> float: """Calculates the test statistic value by calling the ``evaluate`` method of the TestStatistic class with the given log_lambda value and fit parameter values. Parameters ---------- - log_lambda : float + log_lambda The value of the log-likelihood ratio function. Usually, this is its maximum. - fitparam_values : instance of numpy ndarray + fitparam_values The (N_fitparam,)-shaped 1D ndarray holding the global fit parameter values of the log-likelihood ratio function for the given log_lambda value. @@ -626,7 +639,7 @@ def calculate_test_statistic(self, log_lambda, fitparam_values, **kwargs): Returns ------- - TS : float + TS The calculated test-statistic value. """ return self._test_statistic(pmm=self._pmm, log_lambda=log_lambda, fitparam_values=fitparam_values, **kwargs) @@ -666,18 +679,18 @@ def construct_signal_generator(self, **kwargs): self._apply_energy_range(self.sig_gen_energy_range) @abc.abstractmethod - def initialize_trial(self, events_list, n_events_list=None): + def initialize_trial(self, events_list: list[np.ndarray], n_events_list: list[int] | None = None): """This method is supposed to initialize the log-likelihood ratio function with a new set of given trial data. This is a low-level method. For convenient methods see the `unblind` and `do_trial` methods. Parameters ---------- - events_list : list of numpy record ndarray + events_list The list of data events to use for the log-likelihood function evaluation. The data arrays for the datasets must be in the same order than the added datasets. - n_events_list : list of int | None + n_events_list The list of the number of events of each data set. These numbers can be larger than the number of events given by the `events_list` argument in cases where an event selection method was already used. @@ -686,27 +699,27 @@ def initialize_trial(self, events_list, n_events_list=None): """ @abc.abstractmethod - def unblind(self, minimizer_rss, tl=None): + def unblind(self, minimizer_rss: RandomStateService, tl: TimeLord | None = None) -> tuple[float, dict, dict]: """This method is supposed to run the analysis on the experimental data, i.e. unblinds the data. Parameters ---------- - minimizer_rss : instance of RandomStateService + minimizer_rss The instance of RandomStateService that should be used by the minimizer to generate new random initial fit parameter values. - tl : instance of TimeLord | None + tl The optional instance of TimeLord that should be used to time the maximization of the LLH ratio function. Returns ------- - TS : float + TS The test-statistic value. - global_params_dict : dict + global_params_dict The dictionary holding the global parameter names and their best fit values. It includes fixed and floating parameters. - status : dict + status The status dictionary with information about the performed minimization process of the analysis. """ @@ -714,74 +727,74 @@ def unblind(self, minimizer_rss, tl=None): @abc.abstractmethod def do_trial_with_given_pseudo_data( self, - seed, - mean_n_sig, - n_sig, - n_events_list, - events_list, - minimizer_rss, - minimizer_status_dict=None, - tl=None, + seed: int, + mean_n_sig: float, + n_sig: int, + n_events_list: list[int], + events_list: list[DataFieldRecordArray], + minimizer_rss: RandomStateService, + minimizer_status_dict: dict | None = None, + tl: TimeLord | None = None, **kwargs, - ): + ) -> np.ndarray: """This method is supposed to perform an analysis trial on a given pseudo data. Parameters ---------- - seed : int + seed The seed value, which was used to generate the pseudo data. It will be stored in the returned result array. - mean_n_sig : float + mean_n_sig The mean number of signal events the pseudo data was generated with. - n_sig : int + n_sig The total number of actual signal events in the pseudo data. - n_events_list : list of int + n_events_list The total number of events for each data set of the pseudo data. - events_list : list of instance of DataFieldRecordArray + events_list The list of instance of DataFieldRecordArray containing the pseudo data events for each data sample. The number of events for each data sample can be less than the number of events given by ``n_events_list`` if an event selection method was already utilized when generating background events. - minimizer_rss : instance of RandomStateService + minimizer_rss The instance of RandomStateService to use for generating random numbers for the minimizer, e.g. for new initial fit parameter values. - minimizer_status_dict : dict | None + minimizer_status_dict If a dictionary is provided, it will be updated with the minimizer status dictionary. - tl : instance of TimeLord | None + tl The optional instance of TimeLord that should be used to time individual tasks. Returns ------- - recarray : instance of numpy record ndarray + recarray The numpy record ndarray holding the result of the trial. It must contain the following data fields: - rss_seed : int + rss_seed The RandomStateService seed. - mean_n_sig : float + mean_n_sig The mean number of signal events. - n_sig : int + n_sig The actual number of injected signal events. - ts : float + ts The test-statistic value. [ : float ] Any additional parameters of the analysis. """ - def change_shg_mgr(self, shg_mgr, update_detsigyield_service=True): + def change_shg_mgr(self, shg_mgr: SourceHypoGroupManager, update_detsigyield_service: bool = True): """If the SourceHypoGroupManager instance changed, this method needs to be called to propagate the change to all components of the analysis. Parameters ---------- - shg_mgr : instance of SourceHypoGroupManager + shg_mgr The new instance of SourceHypoGroupManager. - update_detsigyield_service : bool + update_detsigyield_service The option whether to update SourceHypoGroupManager of `detsigyield_service` property. It can be set to false for runtime optimization when the detector signal yield does not change for the @@ -811,18 +824,18 @@ def change_shg_mgr(self, shg_mgr, update_detsigyield_service=True): def do_trial_with_given_bkg_and_sig_pseudo_data( self, - seed, - mean_n_sig, - n_sig, - n_bkg_events_list, - n_sig_events_list, - bkg_events_list, - sig_events_list, - minimizer_rss, - minimizer_status_dict=None, - tl=None, + seed: int, + mean_n_sig: float, + n_sig: int, + n_bkg_events_list: list[int], + n_sig_events_list: list[int], + bkg_events_list: list[DataFieldRecordArray | None], + sig_events_list: list[DataFieldRecordArray | None] | None, + minimizer_rss: RandomStateService, + minimizer_status_dict: dict | None = None, + tl: TimeLord | None = None, **kwargs, - ): + ) -> np.ndarray: """Performs an analysis trial on the given background and signal pseudo data. This method merges the background and signal pseudo events and calls the ``do_trial_with_given_pseudo_data`` method of this class. @@ -834,34 +847,34 @@ def do_trial_with_given_bkg_and_sig_pseudo_data( Parameters ---------- - seed : int + seed The seed value, which was used to generate the background and signal pseudo data. It will be stored in the returned result array. - mean_n_sig : float + mean_n_sig The mean number of signal events the pseudo data was generated with. - n_sig : int + n_sig The total number of actual signal events in the pseudo data. - n_bkg_events_list : list of int + n_bkg_events_list The total number of background events for each data set of the pseudo data. - n_sig_events_list : list of int + n_sig_events_list The total number of signal events for each data set of the pseudo data. - bkg_events_list : list of instance of DataFieldRecordArray + bkg_events_list The list of instance of DataFieldRecordArray containing the background pseudo data events for each data set. - sig_events_list : list of instance of DataFieldRecordArray | None + sig_events_list The list of instance of DataFieldRecordArray containing the signal pseudo data events for each data set. If a particular dataset has no signal events, the entry for that dataset can be ``None``. - minimizer_rss : instance of RandomStateService + minimizer_rss The instance of RandomStateService to use for generating random numbers for the minimizer, e.g. for new initial fit parameter values. - minimizer_status_dict : dict | None + minimizer_status_dict If a dictionary is provided, it will be updated with the minimizer status dictionary. - tl : instance of TimeLord | None + tl The instance of TimeLord that should be used to time individual tasks. **kwargs : dict @@ -871,7 +884,7 @@ def do_trial_with_given_bkg_and_sig_pseudo_data( Returns ------- - recarray : instance of numpy record ndarray + recarray The numpy record ndarray holding the result of the trial. See the documentation of the :meth:`~skyllh.core.analysis.Analysis.do_trial_with_given_pseudo_data` @@ -882,19 +895,26 @@ def do_trial_with_given_bkg_and_sig_pseudo_data( events_list = bkg_events_list # Add potential signal events to the background events. - for ds_idx in range(len(events_list)): - if sig_events_list[ds_idx] is not None: - if events_list[ds_idx] is None: - events_list[ds_idx] = sig_events_list[ds_idx] + if sig_events_list is not None: + for ds_idx in range(len(events_list)): + sig_events_ds = sig_events_list[ds_idx] + if sig_events_ds is None: + continue + bkg_events_ds = events_list[ds_idx] + if bkg_events_ds is None: + events_list[ds_idx] = sig_events_ds else: - events_list[ds_idx].append(sig_events_list[ds_idx]) + bkg_events_ds.append(sig_events_ds) + + for ds_idx in range(len(events_list)): + assert events_list[ds_idx] is not None recarray = self.do_trial_with_given_pseudo_data( seed=seed, mean_n_sig=mean_n_sig, n_sig=n_sig, n_events_list=n_events_list, - events_list=events_list, + events_list=cast(list[DataFieldRecordArray], events_list), minimizer_rss=minimizer_rss, minimizer_status_dict=minimizer_status_dict, tl=tl, @@ -905,35 +925,35 @@ def do_trial_with_given_bkg_and_sig_pseudo_data( def generate_background_events( self, - rss, - mean_n_bkg_list=None, - bkg_kwargs=None, - tl=None, - ): + rss: RandomStateService, + mean_n_bkg_list: list[float | None] | None = None, + bkg_kwargs: dict | None = None, + tl: TimeLord | None = None, + ) -> tuple[list[int], list[DataFieldRecordArray]]: """Generates background events utilizing the background generator. Parameters ---------- - rss : instance of RandomStateService + rss The instance of RandomStateService to use for generating random numbers. - mean_n_bkg_list : list of float | None + mean_n_bkg_list The mean number of background events that should be generated for each dataset. If set to None (the default), the background generation method needs to obtain this number itself. - bkg_kwargs : dict | None + bkg_kwargs Optional keyword arguments for the ``generate_background_events`` method of the background generator. - tl : instance of TimeLord | None + tl The optional instance of TimeLord that should be used to time individual tasks of this method. Returns ------- - n_events_list : list of int + n_events_list The list of the number of events that have been generated for each pseudo data set. - events_list : list of instance of DataFieldRecordArray + events_list The list of instance of DataFieldRecordArray containing the pseudo data events for each data sample. The number of events for each data set can be less than the number of events given by @@ -945,6 +965,7 @@ def generate_background_events( if self._bkg_generator is None: self.construct_background_generator() + assert self._bkg_generator is not None (n_events_list, events_list) = self._bkg_generator.generate_background_events( rss=rss, @@ -992,40 +1013,48 @@ def _assert_input_arguments_of_generate_signal_events(self, rss, n_events_list, f'Currently it is of length {len(events_list)}.' ) - def generate_signal_events(self, rss, mean_n_sig, sig_kwargs=None, n_events_list=None, events_list=None, tl=None): + def generate_signal_events( + self, + rss: RandomStateService, + mean_n_sig: float, + sig_kwargs: dict | None = None, + n_events_list: list[int] | None = None, + events_list: list[DataFieldRecordArray | None] | None = None, + tl: TimeLord | None = None, + ) -> tuple[int, list[int], list[DataFieldRecordArray | None]]: """Generates signal events utilizing the signal generator. Parameters ---------- - rss : instance of RandomStateService + rss The instance of RandomStateService to use for generating random numbers. - mean_n_sig : float + mean_n_sig The mean number of signal events that should be generated for the trial. The actual number of generated events will be drawn from a Poisson distribution with this given signal mean as mean. - sig_kwargs : dict | None + sig_kwargs Additional keyword arguments for the ``generate_signal_events`` method of the ``sig_generator_cls`` class. An usual keyword argument is ``poisson``. - n_events_list : list of int | None + n_events_list If given, it specifies the number of events of each data set already present and the number of signal events will be added. - events_list : list of instance of DataFieldRecordArray | None + events_list If given, it specifies the events of each data set already present and the signal events will be added. - tl : instance of TimeLord | None + tl The instance of TimeLord that should be used to time individual tasks of this method. Returns ------- - n_sig : int + n_sig The actual number of injected signal events. - n_events_list : list of int + n_events_list The list of the number of signal events that have been generated for each data set. - events_list : list of instance of DataFieldRecordArray + events_list The list of instance of DataFieldRecordArray containing the signal data events for each data set. An entry is None, if no signal events were generated for this particular data set. @@ -1037,7 +1066,7 @@ def generate_signal_events(self, rss, mean_n_sig, sig_kwargs=None, n_events_list n_events_list = [0] * self.n_datasets if events_list is None: - events_list = [None] * self.n_datasets + events_list = cast(list[DataFieldRecordArray | None], [None] * self.n_datasets) self._assert_input_arguments_of_generate_signal_events( rss=rss, n_events_list=n_events_list, events_list=events_list @@ -1052,6 +1081,7 @@ def generate_signal_events(self, rss, mean_n_sig, sig_kwargs=None, n_events_list if self._sig_generator is None: with TaskTimer(tl, 'Constructing signal generator.'): self.construct_signal_generator() + assert self._sig_generator is not None # Generate signal events with the given mean number of signal # events. @@ -1062,51 +1092,60 @@ def generate_signal_events(self, rss, mean_n_sig, sig_kwargs=None, n_events_list # Inject the signal events to the generated background data. for ds_idx, sig_events in ds_sig_events_dict.items(): n_events_list[ds_idx] += len(sig_events) - if events_list[ds_idx] is None: + existing = events_list[ds_idx] + if existing is None: events_list[ds_idx] = sig_events else: - events_list[ds_idx].append(sig_events) + existing.append(sig_events) return (n_sig, n_events_list, events_list) - def generate_pseudo_data(self, rss, mean_n_bkg_list=None, mean_n_sig=0, bkg_kwargs=None, sig_kwargs=None, tl=None): + def generate_pseudo_data( + self, + rss: RandomStateService, + mean_n_bkg_list: list[float | None] | None = None, + mean_n_sig: float = 0, + bkg_kwargs: dict | None = None, + sig_kwargs: dict | None = None, + tl: TimeLord | None = None, + ) -> tuple[int, list[int], list[DataFieldRecordArray | None]]: """Generates pseudo data with background and possible signal events for each data set using the background and signal generation methods of the analysis. Parameters ---------- - rss : instance of RandomStateService + rss The instance of RandomStateService to use for generating random numbers. - mean_n_bkg_list : list of float | None + mean_n_bkg_list The mean number of background events that should be generated for each dataset. If set to None (the default), the background generation method needs to obtain this number itself. - mean_n_sig : float + mean_n_sig The mean number of signal events that should be generated for the trial. The actual number of generated events will be drawn from a Poisson distribution with this given signal mean as mean. - bkg_kwargs : dict | None + bkg_kwargs Additional keyword arguments for the `generate_events` method of the background generation method class. An usual keyword argument is `poisson`. - sig_kwargs : dict | None + sig_kwargs Additional keyword arguments for the `generate_signal_events` method of the `SignalGenerator` class. An usual keyword argument is `poisson`. - tl : instance of TimeLord | None + tl The optional instance of TimeLord that should be used to time individual tasks of this method. Returns ------- - n_sig : int + n_sig The actual number of injected signal events. - n_events_list : list of int + n_events_list The list of the number of events that have been generated for each pseudo data set. - events_list : list of instance of DataFieldRecordArray + events_list The list of DataFieldRecordArray instances containing the pseudo data events for each data sample. The number of events for each data set can be less than the number of events given by @@ -1125,7 +1164,7 @@ def generate_pseudo_data(self, rss, mean_n_bkg_list=None, mean_n_sig=0, bkg_kwar mean_n_sig=mean_n_sig, sig_kwargs=sig_kwargs, n_events_list=n_events_list, - events_list=events_list, + events_list=cast(list[DataFieldRecordArray | None], events_list), tl=tl, ) @@ -1133,16 +1172,16 @@ def generate_pseudo_data(self, rss, mean_n_bkg_list=None, mean_n_sig=0, bkg_kwar def do_trial( self, - rss, - mean_n_bkg_list=None, - mean_n_sig=0, - bkg_kwargs=None, - sig_kwargs=None, - minimizer_rss=None, - minimizer_status_dict=None, - tl=None, + rss: RandomStateService, + mean_n_bkg_list: list[float | None] | None = None, + mean_n_sig: float = 0, + bkg_kwargs: dict | None = None, + sig_kwargs: dict | None = None, + minimizer_rss: RandomStateService | None = None, + minimizer_status_dict: dict | None = None, + tl: TimeLord | None = None, **kwargs, - ): + ) -> np.ndarray: """This method performs an analysis trial by generating a pseudo data sample with background events and possible signal events via the :meth:`generate_pseudo_data` method, and performs the analysis @@ -1151,33 +1190,33 @@ def do_trial( Parameters ---------- - rss : instance of RandomStateService + rss The instance of RandomStateService to use for generating random numbers. - mean_n_bkg_list : list of float | None + mean_n_bkg_list The mean number of background events that should be generated for each dataset. If set to None (the default), the background generation method needs to obtain this number itself. - mean_n_sig : float + mean_n_sig The mean number of signal events that should be generated for the trial. - bkg_kwargs : dict | None + bkg_kwargs Additional keyword arguments for the `generate_events` method of the background generation method class. An usual keyword argument is `poisson`. - sig_kwargs : dict | None + sig_kwargs Additional keyword arguments for the `generate_signal_events` method of the `SignalGenerator` class. An usual keyword argument is `poisson`. - minimizer_rss : instance of RandomStateService | None + minimizer_rss The instance of RandomStateService to use for generating random numbers for the minimizer, e.g. new initial fit parameter values. If set to ``None``, a rss with the same seed as ``rss`` will be initialized. - minimizer_status_dict : dict | None + minimizer_status_dict If a dictionary is provided, it will be updated with the minimizer status dictionary. - tl : instance of TimeLord | None + tl The optional instance of TimeLord that should be used to time individual tasks. **kwargs : dict @@ -1186,7 +1225,7 @@ def do_trial( Returns ------- - recarray : instance of numpy record ndarray + recarray The numpy record ndarray holding the result of the trial. See the documentation of the :py:meth:`~skyllh.core.analysis.Analysis.do_trial_with_given_pseudo_data` @@ -1205,12 +1244,13 @@ def do_trial( tl=tl, ) + assert rss.seed is not None recarray = self.do_trial_with_given_pseudo_data( seed=rss.seed, mean_n_sig=mean_n_sig, n_sig=n_sig, n_events_list=n_events_list, - events_list=events_list, + events_list=cast(list[DataFieldRecordArray], events_list), minimizer_rss=minimizer_rss, minimizer_status_dict=minimizer_status_dict, tl=tl, @@ -1219,24 +1259,32 @@ def do_trial( return recarray - def do_trials(self, rss, n, ncpu=None, tl=None, ppbar=None, **kwargs): + def do_trials( + self, + rss: RandomStateService, + n: int, + ncpu: int | None = None, + tl: TimeLord | None = None, + ppbar: ProgressBar | None = None, + **kwargs, + ) -> np.ndarray: """Executes the :meth:`do_trial` method ``n`` times with possible multi-processing. Parameters ---------- - rss : instance of RandomStateService + rss The RandomStateService instance to use for generating random numbers. - n : int + n Number of trials to generate using the `do_trial` method. - ncpu : int | None + ncpu The number of CPUs to use, i.e. the number of subprocesses to spawn. If set to None, the global setting will be used. - tl : instance of TimeLord | None + tl The optional instance of TimeLord that should be used to time individual tasks. - ppbar : instance of ProgressBar | None + ppbar The possible parent ProgressBar instance. **kwargs Additional keyword arguments are passed to the :meth:`do_trial` @@ -1245,7 +1293,7 @@ def do_trials(self, rss, n, ncpu=None, tl=None, ppbar=None, **kwargs): Returns ------- - recarray : numpy record ndarray + recarray The numpy record ndarray holding the result of all trials. See the documentation of the :py:meth:`~skyllh.core.analysis.Analysis.do_trial` method for the @@ -1291,28 +1339,36 @@ class LLHRatioAnalysis(Analysis, metaclass=abc.ABCMeta): before any random trial data can be generated. """ - def __init__(self, shg_mgr, pmm, test_statistic, bkg_generator_cls=None, sig_generator_cls=None, **kwargs): + def __init__( + self, + shg_mgr: SourceHypoGroupManager, + pmm: ParameterModelMapper, + test_statistic: TestStatistic, + bkg_generator_cls=None, + sig_generator_cls: type[SignalGenerator] | None = None, + **kwargs, + ): """Constructs a new instance of LLHRatioAnalysis. Parameters ---------- - shg_mgr : instance of SourceHypoGroupManager + shg_mgr The instance of SourceHypoGroupManager, which defines the groups of source hypotheses, their flux model, and their detector signal yield implementation method. - pmm : instance of ParameterModelMapper + pmm The ParameterModelMapper instance managing the global set of parameters and their relation to individual models, e.g. sources. - test_statistic : TestStatistic instance + test_statistic The TestStatistic instance that defines the test statistic function of the analysis. - bkg_generator_cls : class of BackgroundGeneratorBase | None + bkg_generator_cls The background generator class used to create the background generator instance. If set to ``None``, the :class:`skyllh.core.background_generator.BackgroundGenerator` class is used. - sig_generator_cls : class of SignalGenerator | None + sig_generator_cls The signal generator class used to create the signal generator instance. If set to None, the @@ -1324,7 +1380,7 @@ class is used. pmm=pmm, test_statistic=test_statistic, bkg_generator_cls=bkg_generator_cls, - sig_generator_cls=sig_generator_cls, + sig_generator_cls=cast(type[MultiDatasetSignalGenerator] | None, sig_generator_cls), **kwargs, ) @@ -1354,43 +1410,50 @@ def llhratio(self, obj): self._llhratio = obj @abc.abstractmethod - def construct_llhratio(self, minimizer, ppbar=None): + def construct_llhratio(self, minimizer, ppbar=None) -> LLHRatio: """This method is supposed to construct the LLH ratio function. Returns ------- - llhratio : instance of LLHRatio + llhratio The instance of LLHRatio that implements the log-likelihood-ratio function of this LLH ratio analysis. """ - def add_dataset( - self, dataset, data, pdfratio, tdm=None, event_selection_method=None, bkg_generator=None, sig_generator=None + def add_dataset( # pyright: ignore[reportIncompatibleMethodOverride] + self, + dataset: Dataset, + data: DatasetData, + pdfratio: PDFRatio, + tdm: TrialDataManager | None = None, + event_selection_method: EventSelectionMethod | None = None, + bkg_generator: BackgroundGenerator | None = None, + sig_generator: SignalGenerator | None = None, ): """Adds a dataset with its PDF ratio instances to the analysis. Parameters ---------- - dataset : instance of Dataset + dataset The instance of Dataset that should get added. - data : instance of DatasetData + data The instance of DatasetData holding the original (prepared) data of the dataset. - pdfratio : instance of PDFRatio + pdfratio The instance of PDFRatio for the to-be-added data set. - tdm : instance of TrialDataManager | None + tdm The TrialDataManager instance that manages the trial data and additional data fields for this data set. - event_selection_method : instance of EventSelectionMethod | None + event_selection_method The instance of EventSelectionMethod to use to select only signal-like events from the trial data. All other events will be treated as pure background events. This reduces the amount of log-likelihood-ratio function evaluations. If set to None, all events will be evaluated. - bkg_generator : instance of BackgroundGenerator | None + bkg_generator The optional instance of BackgroundGenerator, which should be used to generate background events for this particular dataset. - sig_generator : instance of SignalGenerator | None + sig_generator The optional instance of SignalGenerator, which should be used to generate signal events for this particular dataset. """ @@ -1410,15 +1473,15 @@ def add_dataset( self._pdfratio_list.append(pdfratio) - def change_shg_mgr(self, shg_mgr, update_detsigyield_service=True): + def change_shg_mgr(self, shg_mgr: SourceHypoGroupManager, update_detsigyield_service: bool = True): """If the SourceHypoGroupManager instance changed, this method needs to be called to propagate the change to all components of the analysis. Parameters ---------- - shg_mgr : instance of SourceHypoGroupManager + shg_mgr The new instance of SourceHypoGroupManager. - update_detsigyield_service : bool + update_detsigyield_service The option whether to update SourceHypoGroupManager of `detsigyield_service` property. It can be set to false for runtime optimization when the detector signal yield does not change for the @@ -1433,29 +1496,34 @@ def change_shg_mgr(self, shg_mgr, update_detsigyield_service=True): # Change the source hypo group manager of the LLH ratio function # instance. - self._llhratio.change_shg_mgr(shg_mgr=shg_mgr) + cast(MultiDatasetTCLLHRatio, self._llhratio).change_shg_mgr(shg_mgr=shg_mgr) - def initialize_trial(self, events_list, n_events_list=None, tl=None): + def initialize_trial( # pyright: ignore[reportIncompatibleMethodOverride] + self, + events_list: list[DataFieldRecordArray], + n_events_list: list[int | None] | None = None, + tl: TimeLord | None = None, + ): """This method initializes the log-likelihood ratio function with a new set of given trial data. This is a low-level method. For convenient methods see the ``unblind`` and ``do_trial`` methods. Parameters ---------- - events_list : list of DataFieldRecordArray instances + events_list The list of DataFieldRecordArray instances holding the data events to use for the log-likelihood function evaluation. The data arrays for the datasets must be in the same order than the added datasets. - n_events_list : list of int | None + n_events_list The list of the number of events of each data set. If set to None, the number of events is taken from the size of the given events arrays. - tl : instance of TimeLord | None + tl The optional instance of TimeLord that should be used for timing measurements. """ if n_events_list is None: - n_events_list = [None] * len(events_list) + n_events_list = cast(list[int | None], [None] * len(events_list)) for tdm, events, n_events, evt_sel_method in zip( self._tdm_list, events_list, n_events_list, self._event_selection_method_list, strict=True @@ -1470,28 +1538,28 @@ def initialize_trial(self, events_list, n_events_list=None, tl=None): tl=tl, ) - self._llhratio.initialize_for_new_trial(tl=tl) + self.llhratio.initialize_for_new_trial(tl=tl) - def unblind(self, minimizer_rss, tl=None): + def unblind(self, minimizer_rss: RandomStateService, tl: TimeLord | None = None) -> tuple[float, dict, dict]: """Evaluates the unscrambled data, i.e. unblinds the data. Parameters ---------- - minimizer_rss : instance of RandomStateService + minimizer_rss The instance of RandomStateService that should be used by the minimizer to generate new random initial fit parameter values. - tl : instance of TimeLord | None + tl The optional instance of TimeLord that should be used to time the maximization of the LLH ratio function. Returns ------- - TS : float + TS The test-statistic value. - global_params_dict : dict + global_params_dict The dictionary holding the global parameter names and their best fit values. It includes fixed and floating parameters. - status : dict + status The status dictionary with information about the performed minimization process of the negative of the log-likelihood ratio function. @@ -1499,7 +1567,7 @@ def unblind(self, minimizer_rss, tl=None): events_list = [data.exp for data in self._data_list] self.initialize_trial(events_list) - (log_lambda, fitparam_values, status) = self._llhratio.maximize(rss=minimizer_rss, tl=tl) + (log_lambda, fitparam_values, status) = self.llhratio.maximize(rss=minimizer_rss, tl=tl) TS = self.calculate_test_statistic(log_lambda=log_lambda, fitparam_values=fitparam_values) @@ -1507,67 +1575,67 @@ def unblind(self, minimizer_rss, tl=None): return (TS, global_params_dict, status) - def do_trial_with_given_pseudo_data( + def do_trial_with_given_pseudo_data( # pyright: ignore[reportIncompatibleMethodOverride] self, - seed, - mean_n_sig, - n_sig, - n_events_list, - events_list, - minimizer_rss, - minimizer_status_dict=None, - tl=None, - mean_n_sig_0=None, - ): + seed: int, + mean_n_sig: float, + n_sig: int, + n_events_list: list[int], + events_list: list[DataFieldRecordArray], + minimizer_rss: RandomStateService, + minimizer_status_dict: dict | None = None, + tl: TimeLord | None = None, + mean_n_sig_0: float | None = None, + ) -> np.ndarray: """Performs an analysis trial on the given pseudo data. Parameters ---------- - seed : int + seed The seed value, which was used to generate the pseudo data. It will be stored in the returned result array. - mean_n_sig : float + mean_n_sig The mean number of signal events the pseudo data was generated with. - n_sig : int + n_sig The total number of actual signal events in the pseudo data. - n_events_list : list of int + n_events_list The total number of events for each data set of the pseudo data. - events_list : list of instance of DataFieldRecordArray + events_list The list of instance of DataFieldRecordArray containing the pseudo data events for each data sample. The number of events for each data sample can be less than the number of events given by ``n_events_list`` if an event selection method was already utilized when generating background events. - minimizer_rss : instance of RandomStateService + minimizer_rss The instance of RandomStateService to use for generating random numbers for the minimizer, e.g. for new initial fit parameter values. - minimizer_status_dict : dict | None + minimizer_status_dict If a dictionary is provided, it will be updated with the minimizer status dictionary. - tl : instance of TimeLord | None + tl The optional instance of TimeLord that should be used to time individual tasks. - mean_n_sig_0 : float | None + mean_n_sig_0 The fixed mean number of signal events for the null-hypothesis, when using a ns-profile log-likelihood-ratio function. If set to None, this argument is interpreted as 0. Returns ------- - recarray : instance of numpy record ndarray + recarray The numpy record ndarray holding the result of the trial. It contains the following data fields: - seed : int + seed The seed value of the RandomStateService. - mean_n_sig : float + mean_n_sig The mean number of signal events. - n_sig : int + n_sig The actual number of injected signal events. - mean_n_sig_0 : float + mean_n_sig_0 The fixed mean number of signal events for the null-hypothesis. - ts : float + ts The test-statistic value. [ : float ] Any additional parameters of the LLH ratio function. @@ -1575,13 +1643,13 @@ def do_trial_with_given_pseudo_data( if mean_n_sig_0 is None: mean_n_sig_0 = 0 - self._llhratio.mean_n_sig_0 = mean_n_sig_0 + cast(MultiDatasetTCLLHRatio, self.llhratio).mean_n_sig_0 = mean_n_sig_0 with TaskTimer(tl, 'Initializing trial.'): - self.initialize_trial(events_list, n_events_list) + self.initialize_trial(events_list, cast(list[int | None], n_events_list)) with TaskTimer(tl, 'Maximizing LLH ratio function.'): - (log_lambda, fitparam_values, status) = self._llhratio.maximize(rss=minimizer_rss, tl=tl) + (log_lambda, fitparam_values, status) = self.llhratio.maximize(rss=minimizer_rss, tl=tl) if isinstance(minimizer_status_dict, dict): minimizer_status_dict.update(status) @@ -1621,29 +1689,37 @@ class SingleSourceMultiDatasetLLHRatioAnalysis(LLHRatioAnalysis): documentation of the :class:`~skyllh.core.analysis.LLHRatioAnalysis` class. """ - def __init__(self, shg_mgr, pmm, test_statistic, bkg_generator_cls=None, sig_generator_cls=None, **kwargs): + def __init__( + self, + shg_mgr: SourceHypoGroupManager, + pmm: ParameterModelMapper, + test_statistic: TestStatistic, + bkg_generator_cls=None, + sig_generator_cls: type[SignalGenerator] | None = None, + **kwargs, + ): """Creates a new time-integrated point-like source analysis assuming a single source. Parameters ---------- - shg_mgr : instance of SourceHypoGroupManager + shg_mgr The instance of SourceHypoGroupManager, which defines the groups of source hypotheses, their flux model, and their detector signal efficiency implementation method. - pmm : instance of ParameterModelMapper + pmm The ParameterModelMapper instance managing the global set of parameters and their relation to individual models, e.g. sources. - test_statistic : TestStatistic instance + test_statistic The TestStatistic instance that defines the test statistic function of the analysis. - bkg_generator_cls : class of BackgroundGeneratorBase | None + bkg_generator_cls The background generator class used to create the background generator instance. If set to ``None``, the :class:`skyllh.core.background_generator.BackgroundGenerator` class is used. - sig_generator_cls : SignalGenerator class | None + sig_generator_cls The signal generator class that should be used to create the signal generator instance for multiple datasets. If set to None, the :class:`~skyllh.core.signal_generator.MultiDatasetSignalGenerator` @@ -1658,7 +1734,7 @@ class is used. **kwargs, ) - def construct_llhratio(self, minimizer, ppbar=None): + def construct_llhratio(self, minimizer: Minimizer, ppbar: ProgressBar | None = None) -> MultiDatasetTCLLHRatio: """Constructs the log-likelihood (LLH) ratio function of the analysis. This setups all the necessary analysis objects like detector signal yields and dataset signal weights, constructs the log-likelihood ratio @@ -1666,15 +1742,15 @@ def construct_llhratio(self, minimizer, ppbar=None): Parameters ---------- - minimizer : instance of Minimizer + minimizer The instance of Minimizer that should be used to minimize the negative of the log-likelihood ratio function. - ppbar : instance of ProgressBar | None + ppbar The instance of ProgressBar of the optional parent progress bar. Returns ------- - llhratio : instance of MultiDatasetTCLLHRatio + llhratio The instance of MultiDatasetTCLLHRatio that implements the log-likelihood-ratio function of the analysis. """ @@ -1688,26 +1764,28 @@ def construct_llhratio(self, minimizer, ppbar=None): ] # Create the final multi-dataset log-likelihood ratio function. + assert self.src_detsigyield_weights_service is not None + assert self.ds_sig_weight_factors_service is not None llhratio = MultiDatasetTCLLHRatio( cfg=self._cfg, pmm=self._pmm, minimizer=minimizer, src_detsigyield_weights_service=self.src_detsigyield_weights_service, ds_sig_weight_factors_service=self.ds_sig_weight_factors_service, - llhratio_list=llhratio_list, + llhratio_list=cast(list[SingleDatasetTCLLHRatio], llhratio_list), ) return llhratio - def change_source(self, source, update_detsigyield_service=True): + def change_source(self, source: SourceModel, update_detsigyield_service: bool = True): """Changes the source of the analysis to the given source. It makes the necessary changes to all the objects of the analysis. Parameters ---------- - source : instance of SourceModel + source The instance of SourceModel describing the new source. - update_detsigyield_service : bool + update_detsigyield_service The option whether to update SourceHypoGroupManager of `detsigyield_service` property. It can be set to false for runtime optimization when the detector signal yield does not change for the @@ -1725,7 +1803,9 @@ def change_source(self, source, update_detsigyield_service=True): self.change_shg_mgr(shg_mgr=self._shg_mgr, update_detsigyield_service=update_detsigyield_service) - def calculate_fluxmodel_scaling_factor(self, fitparam_values=None, per_source=False): + def calculate_fluxmodel_scaling_factor( + self, fitparam_values: np.ndarray | None = None, per_source: bool = False + ) -> float | np.ndarray: """Calculates the factor the source's fluxmodel has to be scaled in order to obtain one signal event in the detector. @@ -1742,25 +1822,27 @@ def calculate_fluxmodel_scaling_factor(self, fitparam_values=None, per_source=Fa Parameters ---------- - fitparam_values : numpy ndarray | None + fitparam_values The (N_fitparam,)-shaped 1D ndarray holding the values of the global floating fit parameters at which the scaling factor should be evaluated. The order must match the order in which parameters were defined in the parameter model mapper. If ``None``, the scaling factor is evaluated at the reference parameter values defined in the flux model. - per_source : bool + per_source Whether to return the scaling factor for each source separately. Default is False. Returns ------- - float | (n_sources,)-shaped numpy ndarray + factor The factor(s) the source's fluxmodel needs to be scaled in order to obtain 1 signal event in the detector. """ if self._sig_generator is None: self.construct_signal_generator() - if not hasattr(self._sig_generator, 'fluxmodel_scaling_factor'): - raise RuntimeError( + assert self._sig_generator is not None + + if not isinstance(self._sig_generator, MultiDatasetSignalGenerator): + raise TypeError( 'The configured signal generator does not implement the fluxmodel_scaling_factor interface!' ) @@ -1781,18 +1863,20 @@ def calculate_fluxmodel_scaling_factor(self, fitparam_values=None, per_source=Fa src_params_recarray=src_params_recarray, ) - def mu2flux(self, mu, fitparam_values=None, per_source=False): + def mu2flux( + self, mu: float | np.ndarray, fitparam_values: np.ndarray | None = None, per_source: bool = False + ) -> float | np.ndarray: """Converts the given number of signal events in the detector to the corresponding flux model normalization. Parameters ---------- - mu : float + mu The number of signal events in the detector to convert. - fitparam_values : numpy ndarray | None + fitparam_values The (N_fitparam,)-shaped 1D ndarray holding the values of the global floating fit parameters at which the scaling factor should be evaluated. If ``None``, the reference flux model parameter values are used. - per_source : bool + per_source Whether to return the flux normalization for each source separately. Default is False. Returns @@ -1802,18 +1886,20 @@ def mu2flux(self, mu, fitparam_values=None, per_source=False): """ return self.calculate_fluxmodel_scaling_factor(fitparam_values=fitparam_values, per_source=per_source) * mu - def flux2mu(self, flux_norm, fitparam_values=None, per_source=False): + def flux2mu( + self, flux_norm: float | np.ndarray, fitparam_values: np.ndarray | None = None, per_source: bool = False + ) -> float | np.ndarray: """Converts the given flux model normalization to the corresponding number of signal events in the detector. Parameters ---------- - flux_norm : float + flux_norm The flux model normalization to convert. - fitparam_values : numpy ndarray | None + fitparam_values The (N_fitparam,)-shaped 1D ndarray holding the values of the global floating fit parameters at which the scaling factor should be evaluated. If ``None``, the reference flux model parameter values are used. - per_source : bool + per_source Whether to return the number of signal events for each source separately. Default is False. Returns @@ -1834,28 +1920,36 @@ class MultiSourceMultiDatasetLLHRatioAnalysis(LLHRatioAnalysis): documentation of the :class:`~skyllh.core.analysis.LLHRatioAnalysis` class. """ - def __init__(self, shg_mgr, pmm, test_statistic, bkg_generator_cls=None, sig_generator_cls=None, **kwargs): + def __init__( + self, + shg_mgr: SourceHypoGroupManager, + pmm: ParameterModelMapper, + test_statistic: TestStatistic, + bkg_generator_cls=None, + sig_generator_cls: type[SignalGenerator] | None = None, + **kwargs, + ): """Constructs a new instance of MultiDatasetLLHRatioAnalysis. Parameters ---------- - shg_mgr : instance of SourceHypoGroupManager + shg_mgr The instance of SourceHypoGroupManager, which defines the groups of source hypotheses, their flux model, and their detector signal yield implementation method. - pmm : instance of ParameterModelMapper + pmm The ParameterModelMapper instance managing the global set of parameters and their relation to individual models, e.g. sources. - test_statistic : TestStatistic instance + test_statistic The TestStatistic instance that defines the test statistic function of the analysis. - bkg_generator_cls : class of BackgroundGeneratorBase | None + bkg_generator_cls The background generator class used to create the background generator instance. If set to ``None``, the :class:`skyllh.core.background_generator.BackgroundGenerator` class is used. - sig_generator_cls : subclass of SignalGenerator| None + sig_generator_cls The signal generator class that should be used to create the signal generator instance handling multiple datasets. If set to None, the @@ -1871,7 +1965,7 @@ class is used. **kwargs, ) - def construct_llhratio(self, minimizer, ppbar=None): + def construct_llhratio(self, minimizer: Minimizer, ppbar: ProgressBar | None = None) -> MultiDatasetTCLLHRatio: """Constructs the log-likelihood (LLH) ratio function of the analysis. This setups all the necessary analysis objects like detector signal yields and dataset signal weights, constructs the log-likelihood ratio @@ -1879,20 +1973,22 @@ def construct_llhratio(self, minimizer, ppbar=None): Parameters ---------- - minimizer : instance of Minimizer + minimizer The instance of Minimizer that should be used to minimize the negative of the log-likelihood ratio function. - ppbar : instance of ProgressBar | None + ppbar The instance of ProgressBar of the optional parent progress bar. Returns ------- - llhratio : instance of MultiDatasetTCLLHRatio + llhratio The instance of MultiDatasetTCLLHRatio that implements the log-likelihood-ratio function of the analysis. """ # Create the list of log-likelihood ratio functions, one for each # dataset. + assert self.src_detsigyield_weights_service is not None + assert self.ds_sig_weight_factors_service is not None llhratio_list = [ ZeroSigH0SingleDatasetTCLLHRatio( cfg=self._cfg, @@ -1917,22 +2013,22 @@ def construct_llhratio(self, minimizer, ppbar=None): minimizer=minimizer, src_detsigyield_weights_service=self.src_detsigyield_weights_service, ds_sig_weight_factors_service=self.ds_sig_weight_factors_service, - llhratio_list=llhratio_list, + llhratio_list=cast(list[SingleDatasetTCLLHRatio], llhratio_list), ) return llhratio - def calculate_fluxmodel_scaling_factors(self, mean_ns, fitparam_values): + def calculate_fluxmodel_scaling_factors(self, mean_ns: float, fitparam_values: np.ndarray) -> np.ndarray: """Calculates the factors the source's fluxmodel has to be scaled in order to obtain the given mean number of signal events in the detector. Parameters ---------- - mean_ns : float + mean_ns The mean number of signal events in the detector for which the scaling factors will be calculated. - fitparam_values : instance of numpy ndarray + fitparam_values The (N_fitparam,)-shaped 1D ndarray holding the values of the global fit parameters, which should be used for the flux calculation. The order of the values must match the order the fit parameters were @@ -1940,7 +2036,7 @@ def calculate_fluxmodel_scaling_factors(self, mean_ns, fitparam_values): Returns ------- - factors : instance of numpy ndarray + factors The (N_sources,)-shaped numpy ndarray of float holding the factors the flux models of the sources need to be scaled in order to obtain the given mean number of signal events in the detector. @@ -1951,6 +2047,8 @@ def calculate_fluxmodel_scaling_factors(self, mean_ns, fitparam_values): # events in the detector, for the given reference flux model. mean_ns_ref = np.zeros((self._shg_mgr.n_sources,), dtype=np.float64) + assert self.detsigyield_service is not None + assert self.src_detsigyield_weights_service is not None for g, _ in enumerate(self._shg_mgr.shg_list): shg_src_mask = self._shg_mgr.get_src_mask_of_shg(shg_idx=g) diff --git a/skyllh/core/background_generation.py b/skyllh/core/background_generation.py index 4faf34eef2..ad21a29649 100644 --- a/skyllh/core/background_generation.py +++ b/skyllh/core/background_generation.py @@ -1,7 +1,12 @@ import abc +from collections.abc import Callable +from typing import TYPE_CHECKING, Any, cast import numpy as np +if TYPE_CHECKING: + from skyllh.core.storage import DataFieldRecordArray + from skyllh.core.config import ( HasConfig, ) @@ -11,6 +16,7 @@ from skyllh.core.datafields import ( DataFieldStages as DFS, ) +from skyllh.core.dataset import Dataset, DatasetData from skyllh.core.event_selection import ( AllEventSelectionMethod, EventSelectionMethod, @@ -24,15 +30,12 @@ func_has_n_args, issequenceof, ) -from skyllh.core.random import ( - RandomChoice, -) +from skyllh.core.random import RandomChoice, RandomStateService from skyllh.core.scrambling import ( DataScrambler, ) -from skyllh.core.timing import ( - TaskTimer, -) +from skyllh.core.source_hypo_grouping import SourceHypoGroupManager +from skyllh.core.timing import TaskTimer, TimeLord logger = get_logger(__name__) @@ -52,43 +55,43 @@ def __init__( """Constructs a new background generation method instance.""" super().__init__(**kwargs) - def change_shg_mgr(self, shg_mgr): + def change_shg_mgr(self, shg_mgr: SourceHypoGroupManager): """Notifies the background generation method about an updated SourceHypoGroupManager instance. Parameters ---------- - shg_mgr : instance of SourceHypoGroupManager + shg_mgr The new instance of SourceHypoGroupManager. """ @abc.abstractmethod def generate_events( self, - rss, - dataset, - data, - mean, - tl=None, + rss: RandomStateService, + dataset: Dataset, + data: DatasetData, + mean: float, + tl: TimeLord | None = None, **kwargs, - ): + ) -> 'tuple[int, DataFieldRecordArray]': """This method is supposed to generate a `mean` number of background events for the given dataset and its data. Parameters ---------- - rss : instance of RandomStateService + rss The instance of RandomStateService that should be used to generate random numbers from. - dataset : instance of Dataset + dataset The Dataset instance describing the dataset for which background events should get generated. - data : instance of DatasetData + data The DatasetData instance holding the data of the dataset for which background events should get generated. - mean : float + mean The mean number of background events to generate. - tl : instance of TimeLord | None + tl The optional instance of TimeLord that should be used to collect timing information about this method. **kwargs @@ -97,9 +100,9 @@ def generate_events( Returns ------- - n_bkg : int + n_bkg The number of generated background events. - bkg_events : instance of DataFieldRecordArray + bkg_events The instance of DataFieldRecordArray holding the generated background events. The number of events in this array might be less than ``n_bkg`` if an event selection method was used for @@ -121,18 +124,18 @@ class MCDataSamplingBkgGenMethod( def __init__( self, - get_event_prob_func, - get_mean_func=None, - data_scrambler=None, - keep_mc_data_fields=None, - pre_event_selection_method=None, + get_event_prob_func: Callable, + get_mean_func: Callable | None = None, + data_scrambler: DataScrambler | None = None, + keep_mc_data_fields: str | list[str] | None = None, + pre_event_selection_method: EventSelectionMethod | None = None, **kwargs, ): """Creates a new instance of the MCDataSamplingBkgGenMethod class. Parameters ---------- - get_event_prob_func : callable + get_event_prob_func The function to get the background probability of each monte-carlo event. The call signature of this function must be @@ -143,7 +146,7 @@ def __init__( needs to get generated. The ``events`` argument holds the actual set of events, for which the background event probabilities need to get calculated. - get_mean_func : callable | None + get_mean_func The function to get the mean number of background events. The call signature of this function must be @@ -157,17 +160,17 @@ def __init__( number of background events to generate needs to get specified through the ``generate_events`` method. However, if a pre event selection method is provided, this argument cannot be ``None``! - data_scrambler : instance of DataScrambler | None + data_scrambler If set to an instance of DataScrambler, the drawn monte-carlo background events will get scrambled. This can ensure more independent data trials. It is especially important when monte-carlo statistics are low. - keep_mc_data_fields : str | list of str | None + keep_mc_data_fields The MC data field names that should be kept in order to be able to calculate the background events rates by the functions ``get_event_prob_func`` and ``get_mean_func``. All other MC fields will get dropped due to computational efficiency reasons. - pre_event_selection_method : instance of EventSelectionMethod | None + pre_event_selection_method If set to an instance of EventSelectionMethod, this method will pre-select the MC events that will be used for later background event generation. Using this pre-selection a large portion of the @@ -188,12 +191,12 @@ def __init__( # Define cache members to cache the background probabilities for each # monte-carlo event. The probabilities change only if the data changes. - self._cache_data_id = None - self._cache_mc = None - self._cache_mc_event_bkg_prob = None - self._cache_mean = None - self._cache_mean_pre_selected = None - self._cache_random_choice = None + self._cache_data_id: int | None = None + self._cache_mc: DataFieldRecordArray | None = None + self._cache_mc_event_bkg_prob: np.ndarray | None = None + self._cache_mean: float | None = None + self._cache_mean_pre_selected: float | None = None + self._cache_random_choice: RandomChoice | None = None @property def get_event_prob_func(self): @@ -304,14 +307,14 @@ def pre_event_selection_method(self, method): self._pre_event_selection_method = method - def change_shg_mgr(self, shg_mgr): + def change_shg_mgr(self, shg_mgr: SourceHypoGroupManager): """Changes the instance of SourceHypoGroupManager of the pre-event-selection method. Also it invalidates the data cache of this background generation method. Parameters ---------- - shg_mgr : instance of SourceHypoGroupManager + shg_mgr The new instance of SourceHypoGroupManager. """ if self._pre_event_selection_method is not None: @@ -320,15 +323,15 @@ def change_shg_mgr(self, shg_mgr): # Invalidate the data cache. self._cache_data_id = None - def generate_events( + def generate_events( # pyright: ignore[reportIncompatibleMethodOverride] self, - rss, - dataset, - data, - mean=None, - poisson=True, - tl=None, - ): + rss: RandomStateService, + dataset: Dataset, + data: DatasetData, + mean: float | None = None, + poisson: bool = True, + tl: TimeLord | None = None, + ) -> 'tuple[int, DataFieldRecordArray]': """Generates a ``mean`` number of background events for the given dataset and its data. @@ -352,34 +355,34 @@ def generate_events( Parameters ---------- - rss : instance of RandomStateService + rss The instance of RandomStateService that should be used to generate random numbers from. - dataset : instance of Dataset + dataset The Dataset instance describing the dataset for which background events should get generated. - data : instance of DatasetData + data The DatasetData instance holding the data of the dataset for which background events should get generated. - mean : float | None + mean The mean number of background events to generate. Can be `None`. In that case the mean number of background events is obtained through the `get_mean_func` function. - poisson : bool + poisson If set to ``True`` (default), the actual number of generated background events will be drawn from a Poisson distribution with the given mean number of background events. If set to ``False``, the argument ``mean`` specifies the actual number of generated background events. - tl : instance of TimeLord | None + tl The optional instance of TimeLord that should be used to collect timing information about this method. Returns ------- - n_bkg : int + n_bkg The number of generated background events for the data set. - bkg_events : instance of DataFieldRecordArray + bkg_events The instance of DataFieldRecordArray holding the generated background events. The number of events can be less than `n_bkg` if an event selection method is used. @@ -409,29 +412,35 @@ def generate_events( + self._keep_mc_data_field_names ) ) + assert data.mc is not None self._cache_mc = data.mc.copy(keep_fields=keep_field_names) if self._get_mean_func is not None: with TaskTimer(tl, 'Calculate total MC background mean.'): - self._cache_mean = self._get_mean_func(dataset=dataset, data=data, events=self._cache_mc) + self._cache_mean = float( + cast(Any, self._get_mean_func(dataset=dataset, data=data, events=self._cache_mc)) + ) if self._pre_event_selection_method is not None: with TaskTimer(tl, 'Pre-select MC events.'): - (self._cache_mc, _) = self._pre_event_selection_method.select_events( + _sel_result = self._pre_event_selection_method.select_events( events=self._cache_mc, ret_original_evt_idxs=False, tl=tl ) + self._cache_mc = _sel_result[0] with TaskTimer(tl, 'Calculate selected MC background mean.'): - self._cache_mean_pre_selected = self._get_mean_func( - dataset=dataset, data=data, events=self._cache_mc + assert self._get_mean_func is not None + self._cache_mean_pre_selected = float( + cast(Any, self._get_mean_func(dataset=dataset, data=data, events=self._cache_mc)) ) with TaskTimer(tl, 'Calculate MC background event probability cache.'): - self._cache_mc_event_bkg_prob = self._get_event_prob_func( - dataset=dataset, data=data, events=self._cache_mc + self._cache_mc_event_bkg_prob = np.asarray( + self._get_event_prob_func(dataset=dataset, data=data, events=self._cache_mc) ) with TaskTimer(tl, 'Create RandomChoice for MC background events.'): + assert self._cache_mc_event_bkg_prob is not None self._cache_random_choice = RandomChoice( items=self._cache_mc.indices, probabilities=self._cache_mc_event_bkg_prob ) @@ -443,7 +452,7 @@ def generate_events( 'get_mean_func were specified! One of the two must be ' 'specified!' ) - mean = self._cache_mean + mean = float(self._cache_mean) else: mean = float_cast(mean, 'The mean number of background events must be cast-able to type float!') @@ -454,11 +463,12 @@ def generate_events( # Calculate the mean number of background events for the pre-selected # MC events. - if self._pre_event_selection_method is None: # noqa: SIM108 + if self._pre_event_selection_method is None: # No selection at all, use the total mean. - mean_pre_selected = mean + mean_pre_selected: float = mean else: - mean_pre_selected = self._cache_mean_pre_selected + assert self._cache_mean_pre_selected is not None + mean_pre_selected = float(self._cache_mean_pre_selected) # Calculate the actual number of background events for the selected # events. @@ -467,9 +477,11 @@ def generate_events( # Draw the actual background events from the selected events of the # monte-carlo data set. with TaskTimer(tl, 'Draw MC background indices.'): + assert self._cache_random_choice is not None bkg_event_indices = self._cache_random_choice(rss=rss, size=n_bkg_selected) with TaskTimer(tl, 'Select MC background events from indices.'): + assert self._cache_mc is not None bkg_events = self._cache_mc[bkg_event_indices] # Scramble the drawn MC events if requested. @@ -506,19 +518,19 @@ class CompositeMCDataSamplingBkgGenMethod( def __init__( self, - bkg_component_rate_calc_func_dict, - get_event_prob_func, - get_mean_func=None, - data_scrambler=None, - keep_mc_data_fields=None, - pre_event_selection_method=None, + bkg_component_rate_calc_func_dict: dict, + get_event_prob_func: Callable, + get_mean_func: Callable | None = None, + data_scrambler: DataScrambler | None = None, + keep_mc_data_fields: str | list[str] | None = None, + pre_event_selection_method: EventSelectionMethod | None = None, **kwargs, ): """Creates a new instance of CompositeMCDataSamplingBkgGenMethod. Parameters ---------- - bkg_component_rate_calc_func_dict : dict + bkg_component_rate_calc_func_dict The dictionary holding the name of the background component as key, e.g. "gp", and the background rate calculation function for that component as value. @@ -531,7 +543,7 @@ def __init__( needs to get generated. The ``events`` argument holds the actual set of events, for which the background rate needs to get calculated. - get_event_prob_func : callable + get_event_prob_func The function to get the background probability of each monte-carlo event. The call signature of this function must be @@ -542,7 +554,7 @@ def __init__( needs to get generated. The ``events`` argument holds the actual set of events, for which the background event probabilities need to get calculated. - get_mean_func : callable | None + get_mean_func The function to get the mean number of background events. The call signature of this function must be @@ -556,17 +568,17 @@ def __init__( number of background events to generate needs to get specified through the ``generate_events`` method. However, if a pre event selection method is provided, this argument cannot be ``None``! - data_scrambler : instance of DataScrambler | None + data_scrambler If set to an instance of DataScrambler, the monte-carlo events will get scrambled before drawing the background events. This can ensure more independent data trials. It is especially important when monte-carlo statistics are low. - keep_mc_data_fields : str | list of str | None + keep_mc_data_fields The MC data field names that should be kept in order to be able to calculate the background events rates by the functions ``get_event_prob_func`` and ``get_mean_func``. All other MC fields will get dropped due to computational efficiency reasons. - pre_event_selection_method : instance of EventSelectionMethod | None + pre_event_selection_method If set to an instance of EventSelectionMethod, this method will pre-select the MC events that will be used for later background event generation. Using this pre-selection a large portion of the @@ -612,13 +624,13 @@ def bkg_component_rate_calc_func_dict(self, d): def generate_events( self, - rss, - dataset, - data, - mean=None, - poisson=True, - tl=None, - ): + rss: RandomStateService, + dataset: Dataset, + data: DatasetData, + mean: float | None = None, + poisson: bool = True, + tl: TimeLord | None = None, + ) -> 'tuple[int, DataFieldRecordArray]': """Generates a ``mean`` number of background events for the given monte-carlo dataset and its data. @@ -638,34 +650,34 @@ def generate_events( Parameters ---------- - rss : instance of RandomStateService + rss The instance of RandomStateService that should be used to generate random numbers from. - dataset : instance of Dataset + dataset The Dataset instance describing the dataset for which background events should get generated. - data : instance of DatasetData + data The DatasetData instance holding the data of the dataset for which background events should get generated. - mean : float | None + mean The mean number of background events to generate. Can be `None`. In that case the mean number of background events is obtained through the `get_mean_func` function. - poisson : bool + poisson If set to ``True`` (default), the actual number of generated background events will be drawn from a Poisson distribution with the given mean number of background events. If set to ``False``, the argument ``mean`` specifies the actual number of generated background events. - tl : instance of TimeLord | None + tl The optional instance of TimeLord that should be used to collect timing information about this method. Returns ------- - n_bkg : int + n_bkg The number of generated background events for the data set. - bkg_events : instance of DataFieldRecordArray + bkg_events The instance of DataFieldRecordArray holding the generated background events. The number of events can be less than `n_bkg` if an event selection method is used. @@ -681,6 +693,7 @@ def generate_events( + self._keep_mc_data_field_names ) ) + assert data.mc is not None data_mc = data.mc.copy(keep_fields=keep_field_names) # Scramble the MC events if requested. @@ -703,18 +716,23 @@ def generate_events( mean_mc = self._get_mean_func(dataset=dataset, data=data, events=data_mc) # Pre-select MC events if a pre event selection method is set. + mean_mc_pre_selected: float | None = None if self._pre_event_selection_method is not None: with TaskTimer(tl, 'Pre-select MC events.'): - (data_mc, _) = self._pre_event_selection_method.select_events( + _sel_result = self._pre_event_selection_method.select_events( events=data_mc, ret_original_evt_idxs=False, tl=tl ) + data_mc = _sel_result[0] with TaskTimer(tl, 'Calculate selected MC background mean.'): - mean_mc_pre_selected = self._get_mean_func(dataset=dataset, data=data, events=data_mc) + assert self._get_mean_func is not None + mean_mc_pre_selected = float(cast(Any, self._get_mean_func(dataset=dataset, data=data, events=data_mc))) # Calculate the drawing probability of each selected MC event. with TaskTimer(tl, 'Calculate MC background event probability.'): - mc_event_bkg_prob = self._get_event_prob_func(dataset=dataset, data=data, events=data_mc) + mc_event_bkg_prob: np.ndarray = np.asarray( + self._get_event_prob_func(dataset=dataset, data=data, events=data_mc) + ) with TaskTimer(tl, 'Create RandomChoice for MC background events.'): random_choice = RandomChoice(items=data_mc.indices, probabilities=mc_event_bkg_prob) @@ -727,7 +745,7 @@ def generate_events( 'get_mean_func were specified! One of the two must be ' 'specified!' ) - mean = mean_mc + mean = float(cast(Any, mean_mc)) else: mean = float_cast(mean, 'The mean number of background events must be cast-able to type float!') @@ -738,10 +756,11 @@ def generate_events( # Calculate the mean number of background events for the pre-selected # MC events. - if self._pre_event_selection_method is None: # noqa: SIM108 + if self._pre_event_selection_method is None: # No selection at all, use the total mean. - mean_pre_selected = mean + mean_pre_selected: float = mean else: + assert mean_mc_pre_selected is not None mean_pre_selected = mean_mc_pre_selected # Calculate the actual number of background events for the selected diff --git a/skyllh/core/background_generator.py b/skyllh/core/background_generator.py index c8d1cdc2f6..3d04283d99 100644 --- a/skyllh/core/background_generator.py +++ b/skyllh/core/background_generator.py @@ -1,4 +1,8 @@ import abc +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from skyllh.core.storage import DataFieldRecordArray from skyllh.core.background_generation import ( BackgroundGenerationMethod, @@ -17,9 +21,8 @@ from skyllh.core.random import ( RandomStateService, ) -from skyllh.core.timing import ( - TaskTimer, -) +from skyllh.core.source_hypo_grouping import SourceHypoGroupManager +from skyllh.core.timing import TaskTimer, TimeLord class BackgroundGenerator( @@ -38,38 +41,38 @@ def __init__( Parameters ---------- - bkg_gen_method : instance of BackgroundGenerationMethod - The optional background event generation method, which should be - used to generate events. + **kwargs + Additional keyword arguments are passed to the constructor of the + base class, :class:`~skyllh.core.config.HasConfig`. """ super().__init__(**kwargs) - def change_shg_mgr(self, shg_mgr): + def change_shg_mgr(self, shg_mgr: SourceHypoGroupManager): """This method should be reimplemented when the background generator depends on the sources. Parameters ---------- - shg_mgr : instance of SourceHypoGroupManager + shg_mgr The new instance of SourceHypoGroupManager. """ @abc.abstractmethod def generate_background_events( self, - rss, - tl=None, + rss: RandomStateService, + tl: TimeLord | None = None, **kwargs, - ): + ) -> 'tuple[list[int], list[DataFieldRecordArray]]': """This method is supposed to generate a mean number of background events for the datasets of this background generator. Parameters ---------- - rss : instance of RandomStateService + rss The instance of RandomStateService that should be used to generate random numbers from. - tl : instance of TimeLord | None + tl The optional instance of TimeLord that should be used to collect timing information about this method. **kwargs @@ -82,9 +85,9 @@ def generate_background_events( Returns ------- - n_bkg_list : list of int + n_bkg_list The number of generated background events. - bkg_events_list : list of instance of DataFieldRecordArray + bkg_events_list The list of instance of DataFieldRecordArray holding the generated background events. The number of events can be less than stated in `n_bkg_list` if an event selection method is used. @@ -101,11 +104,28 @@ class DatasetBackgroundGenerator( def __init__( self, - dataset, - data, - bkg_gen_method, + dataset: Dataset, + data: DatasetData, + bkg_gen_method: BackgroundGenerationMethod | None, **kwargs, ): + """Constructs a new instance of DatasetBackgroundGenerator. + + Parameters + ---------- + dataset + The instance of Dataset for which background events should get + generated. + data + The instance of DatasetData holding the experimental and simulation + data of the dataset. + bkg_gen_method + The instance of BackgroundGenerationMethod which should be used to + generate background events. This can be ``None``. + **kwargs + Additional keyword arguments are passed to the constructor of the + base class, :class:`~skyllh.core.background_generator.BackgroundGenerator`. + """ super().__init__(**kwargs) self.dataset = dataset @@ -159,32 +179,33 @@ def bkg_gen_method(self, method): ) self._bkg_gen_method = method - def change_shg_mgr(self, shg_mgr): + def change_shg_mgr(self, shg_mgr: SourceHypoGroupManager): """Changes the SourceHypoGroupManager instance of the background generation method. Parameters ---------- - shg_mgr : instance of SourceHypoGroupManager + shg_mgr The new instance of SourceHypoGroupManager. """ + assert self._bkg_gen_method is not None self._bkg_gen_method.change_shg_mgr(shg_mgr=shg_mgr) def generate_background_events( self, - rss, - tl=None, + rss: RandomStateService, + tl: TimeLord | None = None, **kwargs, - ): + ) -> 'tuple[list[int], list[DataFieldRecordArray]]': """Generates a mean number of background events for the dataset of this background generator. Parameters ---------- - rss : instance of RandomStateService + rss The instance of RandomStateService that should be used to generate random numbers from. - tl : instance of TimeLord | None + tl The optional instance of TimeLord that should be used to collect timing information about this method. **kwargs @@ -197,15 +218,16 @@ def generate_background_events( Returns ------- - n_bkg_list : list of int + n_bkg_list The list of length 1 holding the number of generated background events for the dataset. - bkg_events_list : list of instance of DataFieldRecordArray + bkg_events_list The list of length 1 holding the instance of DataFieldRecordArray holding the generated background events. The number of events can be less than stated in ``n_bkg_list`` if an event selection method is used. """ + assert self._bkg_gen_method is not None (n_bkg, bkg_events) = self._bkg_gen_method.generate_events( rss=rss, dataset=self._dataset, data=self._data, tl=tl, **kwargs ) @@ -224,22 +246,22 @@ class MultiDatasetBackgroundGenerator( def __init__( self, - dataset_list, - data_list, - bkg_generator_list, + dataset_list: list[Dataset], + data_list: list[DatasetData], + bkg_generator_list: 'list[BackgroundGenerator]', **kwargs, ): """Constructs a new instance of MultiDatasetBackgroundGenerator. Parameters ---------- - dataset_list : list of instance of Dataset + dataset_list The list of Dataset instances for which background events should get generated for. - data_list : list of instance of DatasetData + data_list The list of DatasetData instances holding the actual data of each dataset. The order must match the order of ``dataset_list``. - bkg_generator_list : list of instance of BackgroundGenerator + bkg_generator_list The list of BackgroundGenerator instances, one for each dataset. The order must match the order of ``dataset_list``. """ @@ -303,39 +325,39 @@ def bkg_generator_list(self, generators): def change_shg_mgr( self, - shg_mgr, + shg_mgr: SourceHypoGroupManager, ): """Calls the ``change_shg_mgr`` method of each individual dataset background generator. Parameters ---------- - shg_mgr : instance of SourceHypoGroupManager + shg_mgr The new instance of SourceHypoGroupManager. """ for bkg_generator in self._bkg_generator_list: bkg_generator.change_shg_mgr(shg_mgr=shg_mgr) - def generate_background_events( + def generate_background_events( # pyright: ignore[reportIncompatibleMethodOverride] self, - rss, - mean_n_bkg_list=None, - tl=None, + rss: RandomStateService, + mean_n_bkg_list: list[float | None] | None = None, + tl: TimeLord | None = None, **kwargs, - ): + ) -> 'tuple[list[int], list[DataFieldRecordArray]]': """Generates a mean number of background events for each individual dataset of this multi-dataset background generator. Parameters ---------- - rss : instance of RandomStateService + rss The instance of RandomStateService that should be used to generate random numbers from. - mean_n_bkg_list : list of float | None + mean_n_bkg_list The mean number of background events that should be generated for each dataset. If set to None (the default), the individual background generator instance needs to obtain this number itself. - tl : instance of TimeLord | None + tl The optional instance of TimeLord that should be used to collect timing information about this method. **kwargs @@ -345,10 +367,10 @@ def generate_background_events( Returns ------- - n_bkg_events_list : list of int + n_bkg_events_list The list holding the number of generated background events for each dataset. - bkg_events_list : list of instance of DataFieldRecordArray + bkg_events_list The list holding the instance of DataFieldRecordArray holding the generated background events of each dataset. The number of events can be less than stated in ``n_bkg_list`` if an event selection @@ -359,9 +381,12 @@ def generate_background_events( f'The rss argument must be an instance of RandomStateService! Its current type is {classname(rss)}.' ) - if mean_n_bkg_list is None: - mean_n_bkg_list = [None] * len(self._bkg_generator_list) - if not issequenceof(mean_n_bkg_list, (type(None), float)): + _mean_n_bkg_list: list[float | None] + if mean_n_bkg_list is None: # noqa SIM108 + _mean_n_bkg_list = [None] * len(self._bkg_generator_list) + else: + _mean_n_bkg_list = mean_n_bkg_list + if not issequenceof(_mean_n_bkg_list, (type(None), float)): raise TypeError( 'The mean_n_bkg_list argument must be a sequence of None ' 'and/or floats! ' @@ -371,10 +396,10 @@ def generate_background_events( if kwargs is None: kwargs = {} - n_bkg_events_list = [] - bkg_events_list = [] + n_bkg_events_list: list[int] = [] + bkg_events_list: list[DataFieldRecordArray] = [] for ds, bkg_generator, mean_n_bkg in zip( - self._dataset_list, self._bkg_generator_list, mean_n_bkg_list, strict=True + self._dataset_list, self._bkg_generator_list, _mean_n_bkg_list, strict=True ): kwargs.update(mean=mean_n_bkg) with TaskTimer(tl, f'Generating background events for dataset "{ds.name}".'): diff --git a/skyllh/core/backgroundpdf.py b/skyllh/core/backgroundpdf.py index db04977957..904620efcd 100644 --- a/skyllh/core/backgroundpdf.py +++ b/skyllh/core/backgroundpdf.py @@ -4,6 +4,8 @@ import numpy as np +from skyllh.core.flux_model import TimeFluxProfile +from skyllh.core.livetime import Livetime from skyllh.core.pdf import ( IsBackgroundPDF, MultiDimGridPDF, @@ -12,6 +14,8 @@ from skyllh.core.py import ( classname, ) +from skyllh.core.timing import TimeLord +from skyllh.core.trialdata import TrialDataManager class BackgroundMultiDimGridPDF(MultiDimGridPDF, IsBackgroundPDF): @@ -34,30 +38,30 @@ def __init__(self, *args, **kwargs): class BackgroundTimePDF(TimePDF, IsBackgroundPDF): """This class provides a background time PDF class.""" - def __init__(self, livetime, time_flux_profile, **kwargs): + def __init__(self, livetime: Livetime, time_flux_profile: TimeFluxProfile, **kwargs): """Creates a new signal time PDF instance for a given time flux profile and detector live time. Parameters ---------- - livetime : instance of Livetime + livetime An instance of Livetime, which provides the detector live-time information. - time_flux_profile : instance of TimeFluxProfile + time_flux_profile The signal's time flux profile. """ super().__init__(pmm=None, livetime=livetime, time_flux_profile=time_flux_profile, **kwargs) self._pd = None - def initialize_for_new_trial(self, tdm, tl=None, **kwargs): + def initialize_for_new_trial(self, tdm: TrialDataManager, tl: TimeLord | None = None, **kwargs): """Initializes the background time PDF with new trial data. Because this PDF does not depend on any parameters, the probability density values can be pre-computed here. Parameters ---------- - tdm : instance of TrialDataManager + tdm The instance of TrialDataManager holding the trial event data for which to calculate the PDF value. The following data fields must exist: @@ -65,7 +69,7 @@ def initialize_for_new_trial(self, tdm, tl=None, **kwargs): ``'time'`` : float The MJD time of the event. - tl : instance of TimeLord | None + tl The optional TimeLord instance that should be used to measure timing information. """ @@ -79,11 +83,13 @@ def initialize_for_new_trial(self, tdm, tl=None, **kwargs): self._pd[on] = self._time_flux_profile(t=times[on]) / self._S - def get_pd(self, tdm, params_recarray=None, tl=None): + def get_pd( + self, tdm: TrialDataManager, params_recarray: np.ndarray | None = None, tl: TimeLord | None = None + ) -> tuple[np.ndarray, dict]: """ Parameters ---------- - tdm : instance of TrialDataManager + tdm The instance of TrialDataManager holding the trial event data for which to calculate the PDF value. The following data fields must exist: @@ -91,18 +97,18 @@ def get_pd(self, tdm, params_recarray=None, tl=None): ``'time'`` : float The MJD time of the event. - params_recarray : None + params_recarray Unused interface argument. - tl : instance of TimeLord | None + tl The optional TimeLord instance that should be used to measure timing information. Returns ------- - pd : instance of numpy ndarray + pd The (N_events,)-shaped numpy ndarray holding the background probability density value for each event. - grads : dict + grads The dictionary holding the gradients of the probability density w.r.t. each global fit parameter. The background PDF does not depend on any global fit parameter, diff --git a/skyllh/core/binning.py b/skyllh/core/binning.py index 92619b641b..68b24bd2e5 100644 --- a/skyllh/core/binning.py +++ b/skyllh/core/binning.py @@ -1,3 +1,5 @@ +from collections.abc import Sequence + import numpy as np from skyllh.core.py import ( @@ -5,34 +7,34 @@ ) -def get_bincenters_from_binedges(edges): +def get_bincenters_from_binedges(edges: np.ndarray) -> np.ndarray: """Calculates the bin center values from the given bin edge values. Parameters ---------- - edges : 1D numpy ndarray + edges The (n+1,)-shaped 1D ndarray holding the bin edge values. Returns ------- - bincenters : 1D numpy ndarray + bincenters The (n,)-shaped 1D ndarray holding the bin center values. """ return 0.5 * (edges[:-1] + edges[1:]) -def get_binedges_from_bincenters(centers): +def get_binedges_from_bincenters(centers: np.ndarray) -> np.ndarray: """Calculates the bin edges from the given bin center values. The bin center values must be evenly spaced. Parameters ---------- - centers : 1D numpy ndarray + centers The (n,)-shaped 1D ndarray holding the bin center values. Returns ------- - edges : 1D numpy ndarray + edges The (n+1,)-shaped 1D ndarray holding the bin edge values. """ d = np.diff(centers) @@ -47,7 +49,7 @@ def get_binedges_from_bincenters(centers): return edges -def get_bin_indices_from_lower_and_upper_binedges(le, ue, values): +def get_bin_indices_from_lower_and_upper_binedges(le: np.ndarray, ue: np.ndarray, values: np.ndarray) -> np.ndarray: """Returns the bin indices for the given values which must fall into bins defined by the given lower and upper bin edges. @@ -55,16 +57,16 @@ def get_bin_indices_from_lower_and_upper_binedges(le, ue, values): Parameters ---------- - le : (m,)-shaped 1D numpy ndarray + le The lower bin edges. - ue : (m,)-shaped 1D numpy ndarray + ue The upper bin edges. - values : (n,)-shaped 1D numpy ndarray + values The values for which to get the bin indices. Returns ------- - idxs : (n,)-shaped 1D numpy ndarray + idxs The bin indices of the given values. """ if len(le) != len(ue): @@ -95,15 +97,15 @@ class BinningDefinition: binning definitions for an analysis. """ - def __init__(self, name, binedges): + def __init__(self, name: str, binedges: np.ndarray | Sequence): """Creates a new binning definition object. Parameters ---------- - name : str + name The name of the binning definition. - binedges : sequence of float - The sequence of the bin edges, which should be used for the binning. + binedges + The bin edges, which should be used for the binning. """ self.name = name self.binedges = binedges @@ -114,7 +116,7 @@ def __str__(self): s += str(self._binedges) return s - def __eq__(self, other): + def __eq__(self, other: object) -> bool: """Checks if object ``other`` is equal to this BinningDefinition object.""" if not isinstance(other, BinningDefinition): raise TypeError( @@ -125,7 +127,7 @@ def __eq__(self, other): if self.name != other.name: return False - return np.all(self.binedges == other.binedges) + return bool(np.all(self.binedges == other.binedges)) @property def name(self): @@ -180,23 +182,23 @@ def range(self): """The tuple (lower_edge, upper_edge) of the binning.""" return (self.lower_edge, self.upper_edge) - def any_data_out_of_range(self, data): + def any_data_out_of_range(self, data: np.ndarray) -> bool: """Checks if any of the given data is outside the range of this binning definition. Parameters ---------- - data : instance of ndarray + data The 1D ndarray with the data values to check. Returns ------- - outofrange : bool + outofrange True if any data value is outside the binning range. False otherwise. """ outofrange = np.any((data < self.lower_edge) | (data > self.upper_edge)) - return outofrange + return bool(outofrange) def get_binwidth_from_value(self, value): """Returns the width of the bin the given value falls into.""" @@ -206,18 +208,18 @@ def get_binwidth_from_value(self, value): return bin_width - def get_out_of_range_data(self, data): + def get_out_of_range_data(self, data: np.ndarray) -> np.ndarray: """Returns the data values which are outside the range of this binning definition. Parameters ---------- - data : instance of numpy.ndarray + data The 1D ndarray with the data values to check. Returns ------- - oor_data : instance of numpy.ndarray + oor_data The 1D ndarray with data outside the range of this binning definition. """ @@ -226,21 +228,21 @@ def get_out_of_range_data(self, data): return oor_data - def get_subset(self, lower_edge, upper_edge): + def get_subset(self, lower_edge: float, upper_edge: float) -> 'BinningDefinition': """Creates a new BinningDefinition instance which contains only a subset of the bins of this BinningDefinition instance. The range of the subset is given by a lower and upper edge value. Parameters ---------- - lower_edge : float + lower_edge The lower edge value of the subset. - upper_edge : float + upper_edge The upper edge value of the subset. Returns ------- - binning : instance of BinningDefinition + binning The new instance of BinningDefinition holding the binning subset. """ @@ -274,6 +276,9 @@ class UsesBinning: """ def __init__(self, *args, **kwargs): + """Creates a new instance of UsesBinning and initializes the empty list + of binning definitions. + """ super().__init__(*args, **kwargs) # Define the list of binning definition objects and a name->list_index @@ -293,17 +298,17 @@ def binning_ndim(self): """(read-only) The number of dimensions that uses binning.""" return len(self._binnings) - def has_same_binning_as(self, obj): + def has_same_binning_as(self, obj: 'UsesBinning') -> bool: """Checks if this object has the same binning as the given object. Parameters ---------- - obj : instance of UsesBinning + obj The object that should be checked for same binning. Returns ------- - check : bool + check True if ``obj`` uses the same binning, False otherwise. """ if not isinstance(obj, UsesBinning): @@ -317,14 +322,14 @@ def has_same_binning_as(self, obj): return True - def add_binning(self, binning, name=None): + def add_binning(self, binning: 'BinningDefinition', name: str | None = None): """Adds the given binning definition to the list of binnings. Parameters ---------- - binning : instance of BinningDefinition + binning The binning definition to add. - name : str | (default) None + name The name of the binning. If not None and it's different to the name of the given binning definition, a copy of the BinningDefinition object is made and the new name is set. @@ -346,18 +351,18 @@ def add_binning(self, binning, name=None): self._binnings.append(binning) self._binning_name2idx[binning.name] = len(self._binnings) - 1 - def get_binning(self, name): + def get_binning(self, name: str | int) -> 'BinningDefinition': """Retrieves the binning definition of the given name. Parameters ---------- - name : str | int + name The name of the binning definition. A string specifies the name and an integer the dimension index. Returns ------- - binning : instance of BinningDefinition + binning The binning definition of the given name. """ if isinstance(name, str): diff --git a/skyllh/core/catalog.py b/skyllh/core/catalog.py index d05853f7d8..7d56d8cdc2 100644 --- a/skyllh/core/catalog.py +++ b/skyllh/core/catalog.py @@ -13,16 +13,16 @@ class SourceCatalog(SourceModelCollection): SourceModelCollection. A catalog has a name. """ - def __init__(self, name, sources=None, source_type=None, **kwargs): + def __init__(self, name: str, sources=None, source_type: type | None = None, **kwargs): """Creates a new source catalog. Parameters ---------- - name : str + name The name of the catalog. - sources : sequence of source_type | None + sources The sequence of sources this catalog should be initialized with. - source_type : type | None + source_type The type of the source class. If set to None (default), the default type defined by SourceCollection will be used. """ @@ -44,13 +44,13 @@ def __str__(self): s = f'"{self.name}" {super().__str__()}' return s - def as_SourceModelCollection(self): + def as_SourceModelCollection(self) -> SourceModelCollection: """Creates a SourceModelCollection object for this catalog and returns it. Returns ------- - source_model_collection : instance of SourceModelCollection + source_model_collection The created instance of SourceModelCollection. """ return SourceModelCollection(sources=self.sources, source_type=self.source_type) diff --git a/skyllh/core/config.py b/skyllh/core/config.py index f0279ea99f..d7dada177d 100644 --- a/skyllh/core/config.py +++ b/skyllh/core/config.py @@ -110,21 +110,21 @@ def __init__( def from_yaml( cls, pathfilename: str | None, - ): + ) -> 'Config': """Creates a new instance of Config holding the base configuration and updated by the configuration items contained in the yaml file using the :meth:`dict.update` method. Parameters ---------- - pathfilename: str | None + pathfilename Path and filename to the yaml file containing the to-be-updated configuration items. If set to ``None``, nothing is done. Returns ------- - cfg : instance of Config + cfg The instance of Config holding the base configuration and updated by the configuration given in the yaml file. """ @@ -150,19 +150,19 @@ def from_yaml( def from_dict( cls, user_dict: dict[str, Any], - ): + ) -> 'Config': """Creates a new instance of Config holding the base configuration and updated by the given configuration dictionary using the :meth:`dict.update` method. Parameters ---------- - user_dict: dict + user_dict The dictionary containing the to-be-updated configuration items. Returns ------- - cfg : instance of Config + cfg The instance of Config holding the base configuration and updated by the given configuration dictionary. """ @@ -179,12 +179,12 @@ def is_tracing_enabled(self): def disable_tracing( self, - ): + ) -> 'Config': """Disables the tracing mode of SkyLLH. Returns ------- - self : instance of Config + self The updated instance of Config. """ self['logging']['enable_tracing'] = False @@ -193,12 +193,12 @@ def disable_tracing( def enable_tracing( self, - ): + ) -> 'Config': """Enables the tracing mode of SkyLLH. Returns ------- - self : instance of Config + self The updated instance of Config. """ self['logging']['enable_tracing'] = True @@ -207,13 +207,13 @@ def enable_tracing( def get_wd( self, - ): + ) -> str: """Retrieves the absolute path to the working directory as configured in this configuration. Returns ------- - wd : str + wd The absolute path to the project's working directory. """ wd = os.path.abspath(self['project']['working_directory']) @@ -222,19 +222,19 @@ def get_wd( def set_enable_tracing( self, - flag, - ): + flag: bool, + ) -> 'Config': """Sets the setting for tracing. Parameters ---------- - flag : bool + flag The flag if tracing should be enabled (``True``) or disabled (``False``). Returns ------- - self : instance of Config + self The updated instance of Config. """ self['logging']['enable_tracing'] = flag @@ -243,32 +243,32 @@ def set_enable_tracing( def set_internal_units( self, - angle_unit=None, - energy_unit=None, - length_unit=None, - time_unit=None, - ): + angle_unit: units.UnitBase | None = None, + energy_unit: units.UnitBase | None = None, + length_unit: units.UnitBase | None = None, + time_unit: units.UnitBase | None = None, + ) -> 'Config': """Sets the units used internally to compute quantities. These units must match the units used in the monte-carlo files. Parameters ---------- - angle_unit : instance of astropy.units.UnitBase | None + angle_unit The internal unit that should be used for angles. If set to ``None``, the unit is not changed. - energy_unit : instance of astropy.units.UnitBase | None + energy_unit The internal unit that should be used for energy. If set to ``None``, the unit is not changed. - length_unit : instance of astropy.units.UnitBase | None + length_unit The internal unit that should be used for length. If set to ``None``, the unit is not changed. - time_unit : instance of astropy.units.UnitBase | None + time_unit The internal unit that should be used for time. If set to ``None``, the unit is not changed. Returns ------- - self : instance of Config + self The updated instance of Config. """ if angle_unit is not None: @@ -295,19 +295,19 @@ def set_internal_units( def set_ncpu( self, - ncpu, - ): + ncpu: int, + ) -> 'Config': """Sets the global setting for the number of CPUs to use, when parallelization is available. Parameters ---------- - ncpu : int + ncpu The number of CPUs. Returns ------- - self : instance of Config + self The updated instance of Config. """ self['multiproc']['ncpu'] = ncpu @@ -316,16 +316,14 @@ def set_ncpu( def set_wd( self, - path=None, - ): + path: str | None = None, + ) -> str: """Sets the project's working directory configuration variable and adds it to the Python path variable. Parameters ---------- - cfg : instance of Config - The instance of Config holding the local configuration. - path : str | None + path The path of the project's working directory. This can be a path relative to the path given by ``os.path.getcwd``, the current working directory of the program. @@ -334,7 +332,7 @@ def set_wd( Returns ------- - wd : str + wd The absolute path to the project's working directory. """ if path is None: @@ -343,7 +341,7 @@ def set_wd( if self['project']['working_directory'] in sys.path: sys.path.remove(self['project']['working_directory']) - wd = os.path.abspath(path) + wd = os.path.abspath(str(path)) self['project']['working_directory'] = wd sys.path.insert(0, wd) @@ -351,14 +349,14 @@ def set_wd( def to_internal_time_unit( self, - time_unit, + time_unit: units.UnitBase, ): """Calculates the conversion factor from the given time unit to the internal time unit specified by this local configuration. Parameters ---------- - time_unit : instance of astropy.units.UnitBase + time_unit The time unit from which to convert to the internal time unit. """ internal_time_unit = self['units']['internal']['time'] @@ -366,19 +364,19 @@ def to_internal_time_unit( return factor - def wd_filename(self, filename): + def wd_filename(self, filename: str) -> str: """Generates the fully qualified file name under the project's working directory of the given file. Parameters ---------- - filename : str + filename The name of the file for which to generate the working directory path file name. Returns ------- - pathfilename : str + pathfilename The generated fully qualified path file name of ``filename`` with the project's working directory prefixed. """ @@ -394,7 +392,7 @@ class HasConfig: def __init__( self, - cfg, + cfg: 'Config', *args, **kwargs, ): @@ -402,7 +400,7 @@ def __init__( Parameters ---------- - cfg : instance of Config + cfg The instance of Config holding the local configuration. """ super().__init__(*args, **kwargs) diff --git a/skyllh/core/datafields.py b/skyllh/core/datafields.py index ae7870b203..1d18c6cdd6 100644 --- a/skyllh/core/datafields.py +++ b/skyllh/core/datafields.py @@ -2,6 +2,8 @@ file is required at what stage. """ +from collections.abc import Sequence + class DataFieldStages: """This class provides the data field stage values, which are individual @@ -15,21 +17,21 @@ class DataFieldStages: @staticmethod def and_check( - stage, - stages, - ): + stage: int, + stages: int | Sequence[int], + ) -> bool: """Checks if the given stage matches all of the given stages. Parameters ---------- - stage : int + stage The stage value, which should get checked. - stages : int | sequence of int + stages The stage(s) to check for. Returns ------- - check : bool + check ``True`` if the given stage contains all of the given stages, ``False`` otherwise. """ @@ -40,21 +42,21 @@ def and_check( @staticmethod def or_check( - stage, - stages, - ): + stage: int, + stages: int | Sequence[int], + ) -> bool: """Checks if the given stage matches any of the given stages. Parameters ---------- - stage : int + stage The stage value, which should get checked. - stages : int | sequence of int + stages The stage(s) to check for. Returns ------- - check : bool + check ``True`` if the given stage contains any of the given stages, ``False`` otherwise. """ @@ -65,24 +67,28 @@ def or_check( class DataFields: + """This class provides utility methods for selecting data field names based + on their assigned processing stages. + """ + @staticmethod def get_joint_names( - datafields, - stages, - ): + datafields: dict, + stages: int | Sequence[int], + ) -> list[str]: """Returns the list of data field names that match at least one of the given stages, i.e. the joint set of data fields given the stages. Parameters ---------- - datafields : dict + datafields The dictionary of data field names as keys and stages as values. - stages : int | sequence of int + stages The stage(s) for which data field names should get returned. Returns ------- - datafield_names : list of str + datafield_names The list of data field names. """ datafield_names = [field for (field, stage) in datafields.items() if DataFieldStages.or_check(stage, stages)] diff --git a/skyllh/core/dataset.py b/skyllh/core/dataset.py index cdb221f103..0825be9dc1 100644 --- a/skyllh/core/dataset.py +++ b/skyllh/core/dataset.py @@ -7,7 +7,9 @@ import urllib.error import urllib.request import zipfile +from collections.abc import Callable, Sequence from copy import deepcopy +from typing import Any, cast, overload import numpy as np @@ -52,9 +54,7 @@ DataFieldRecordArray, create_FileLoader, ) -from skyllh.core.timing import ( - TaskTimer, -) +from skyllh.core.timing import TaskTimer, TimeLord class DatasetOrigin: @@ -64,27 +64,27 @@ class DatasetOrigin: def __init__( self, - base_path, - sub_path, - transfer_func, - filename=None, - host=None, - port=None, - username=None, - password=None, - post_transfer_func=None, - url=None, + base_path: str, + sub_path: str, + transfer_func: Callable, + filename: str | None = None, + host: str | None = None, + port: int | None = None, + username: str | None = None, + password: str | None = None, + post_transfer_func: Callable | None = None, + url: str | None = None, **kwargs, ): """Creates a new instance to define the origin of a dataset. Parameters ---------- - base_path : str + base_path The dataset's base directory at the origin. - sub_path : str + sub_path The dataset's sub directory at the origin. - transfer_func : callable + transfer_func The callable object that should be used to transfer the dataset. This function requires the following call signature:: @@ -96,20 +96,20 @@ def __init__( machine, ``user`` is the user name required to connect to the remote host, and ``password`` is the password for the user name required to connect to the remote host. - filename : str | None + filename If the origin is not a directory but a file, this specifies the filename. When ``url`` is set this becomes the local filename used when saving the downloaded file. - host : str | None + host The name or IP of the remote host. - port : int | None + port The port number to use when connecting to the remote host. - username : str | None + username The user name required to connect to the remote host. - password : str | None + password The password for the user name required to connect to the remote host. - post_transfer_func : callable | None + post_transfer_func The callable object that should be called after the dataset has been transferred by the ``transfer_func`` function. It can be used to extract an archive file. @@ -119,7 +119,7 @@ def __init__( where ``ds`` is an instance of ``Dataset``, and ``dst_path`` is the destination path. - url : str | None + url An optional complete download URL (e.g. an API endpoint with query parameters). When set, the ``transfer_func`` should use this URL directly instead of constructing one from ``host``, ``base_path``, @@ -307,13 +307,13 @@ def __str__(self): return s - def is_locally_available(self): + def is_locally_available(self) -> bool: """Checks if the dataset origin is available locally by checking if the given path exists on the local host. Returns ------- - check : bool + check ``True`` if the path specified in this dataset origin is an absolute path and exists on the local host, ``False`` otherwise. """ @@ -348,6 +348,17 @@ def __init__( text, mode=stat.S_IRUSR | stat.S_IWUSR | stat.S_IRGRP | stat.S_IROTH, ): + """Creates a new instance of TemporaryTextFile. + + Parameters + ---------- + pathfilename + The path and file name of the temporary text file. + text + The text that should be written into the file. + mode + The file access mode that should be set for the created file. + """ self.pathfilename = pathfilename self.text = text self.mode = mode @@ -381,23 +392,24 @@ class DatasetTransfer( """Base class for a dataset transfer mechanism.""" def __init__(self, **kwargs): + """Creates a new instance of DatasetTransfer.""" super().__init__(**kwargs) @staticmethod def execute_system_command( - cmd, + cmd: str, logger, - success_rcode=0, + success_rcode: int | None = 0, ): """Executes the given system command via a ``os.system`` call. Parameters ---------- - cmd : str + cmd The system command to execute. - logger : instance of logging.Logger + logger The logger to use for debug messages. - success_rcode : int | None + success_rcode The return code that indicates success of the system command. If set to ``None``, no return code checking is performed. @@ -413,13 +425,13 @@ def execute_system_command( @staticmethod def ensure_dst_path( - dst_path, + dst_path: str, ): """Ensures the existence of the given destination path. Parameters ---------- - dst_path : str + dst_path The destination path. """ if not os.path.isdir(dst_path): @@ -429,28 +441,28 @@ def ensure_dst_path( @abc.abstractmethod def transfer( self, - origin, - file_list, - dst_base_path, - username=None, - password=None, + origin: 'DatasetOrigin', + file_list: list[str], + dst_base_path: str, + username: str | None = None, + password: str | None = None, ): """This method is supposed to transfer the dataset origin path to the given destination path. Parameters ---------- - origin : instance of DatasetOrigin + origin The instance of DatasetOrigin defining the origin of the dataset. - file_list : list of str + file_list The list of files, relative to the origin base path, which should be transferred. - dst_base_path : str + dst_base_path The destination base path into which the dataset files will be transferred. - username : str | None + username The user name required to connect to the remote host. - password : str | None + password The password for the user name required to connect to the remote host. @@ -464,34 +476,38 @@ def transfer( class RSYNCDatasetTransfer( DatasetTransfer, ): + """This class provides a dataset transfer mechanism using the ``rsync`` + program. + """ + def __init__(self, **kwargs): + """Creates a new instance of RSYNCDatasetTransfer.""" super().__init__(**kwargs) def transfer( self, origin, - file_list, - dst_base_path, - username=None, - password=None, + file_list: list[str], + dst_base_path: str, + username: str | None = None, + password: str | None = None, ): """Transfers the given dataset to the given destination path using the ``rsync`` program. Parameters ---------- - ds : instance of Dataset - The instance of Dataset containing the origin property specifying - the origin of the dataset. - file_list : list of str + origin + The instance of DatasetOrigin specifying the origin of the dataset. + file_list The list of files, relative to the origin base path, which should be transferred. - dst_base_path : str + dst_base_path The destination base path into which the dataset files will be transferred. - username : str | None + username The user name required to connect to the remote host. - password : str | None + password The password for the user name required to connect to the remote host. """ @@ -585,7 +601,18 @@ def transfer( class WGETDatasetTransfer( DatasetTransfer, ): + """This class provides a dataset transfer mechanism using the ``wget`` + program. + """ + def __init__(self, protocol, **kwargs): + """Creates a new instance of WGETDatasetTransfer. + + Parameters + ---------- + protocol + The protocol to use for the transfer, e.g. ``"http"`` or ``"https"``. + """ super().__init__(**kwargs) self.protocol = protocol @@ -603,28 +630,28 @@ def protocol(self, obj): def transfer( self, - origin, - file_list, - dst_base_path, - username=None, - password=None, + origin: 'DatasetOrigin', + file_list: list[str], + dst_base_path: str, + username: str | None = None, + password: str | None = None, ): """Transfers the given dataset to the given destination path using the ``wget`` program. Parameters ---------- - origin : instance of DatasetOrigin + origin The instance of DatasetOrigin defining the origin of the dataset. - file_list : list of str + file_list The list of files relative to the origin's base path, which should be transferred. - dst_base_path : str + dst_base_path The destination base path into which the dataset will be transferred. - username : str | None + username The user name required to connect to the remote host. - password : str | None + password The password for the user name required to connect to the remote host. """ @@ -682,7 +709,18 @@ def transfer( class URLRetrieveDatasetTransfer( DatasetTransfer, ): + """This class provides a dataset transfer mechanism using Python's + :func:`urllib.request.urlretrieve` function. + """ + def __init__(self, protocol, **kwargs): + """Creates a new instance of URLRetrieveDatasetTransfer. + + Parameters + ---------- + protocol + The protocol to use for the transfer, e.g. ``"http"`` or ``"https"``. + """ super().__init__(**kwargs) self.protocol = protocol @@ -750,8 +788,8 @@ def transfer( # Use the full URL directly; save with an explicit output path # so the local filename matches origin.filename regardless of # what the server sends in the URL path or Content-Disposition. + url = origin.url output_file = os.path.join(dst_path, os.path.basename(file)) - transfer_args = [origin.url, output_file] else: url = f'{self.protocol}://{host}' if port is not None: @@ -759,7 +797,7 @@ def transfer( if path[0:1] != '/': url += '/' url += path - transfer_args = [url, dst_pathfilename] + output_file = dst_pathfilename if username is not None: password_mgr = urllib.request.HTTPPasswordMgrWithDefaultRealm() @@ -772,7 +810,7 @@ def transfer( opener.addheaders = [('User-Agent', 'Wget/1.25')] urllib.request.install_opener(opener) try: - urllib.request.urlretrieve(*transfer_args) + urllib.request.urlretrieve(url, output_file) except urllib.error.URLError as err: if os.path.exists(dst_pathfilename): os.remove(dst_pathfilename) @@ -835,18 +873,18 @@ class Dataset( """ @staticmethod - def get_combined_exp_pathfilenames(datasets): + def get_combined_exp_pathfilenames(datasets: 'Sequence[Dataset]') -> list: """Creates the combined list of exp pathfilenames of all the given datasets. Parameters ---------- - datasets : sequence of Dataset + datasets The sequence of Dataset instances. Returns ------- - exp_pathfilenames : list + exp_pathfilenames The combined list of exp pathfilenames. """ if not issequenceof(datasets, Dataset): @@ -859,18 +897,18 @@ def get_combined_exp_pathfilenames(datasets): return exp_pathfilenames @staticmethod - def get_combined_mc_pathfilenames(datasets): + def get_combined_mc_pathfilenames(datasets: 'Sequence[Dataset]') -> list: """Creates the combined list of mc pathfilenames of all the given datasets. Parameters ---------- - datasets : sequence of Dataset + datasets The sequence of Dataset instances. Returns ------- - mc_pathfilenames : list + mc_pathfilenames The combined list of mc pathfilenames. """ if not issequenceof(datasets, Dataset): @@ -883,38 +921,44 @@ def get_combined_mc_pathfilenames(datasets): return mc_pathfilenames @staticmethod - def get_combined_livetime(datasets): + def get_combined_livetime(datasets: 'Sequence[Dataset]') -> float: """Sums the live-time of all the given datasets. Parameters ---------- - datasets : sequence of Dataset + datasets The sequence of Dataset instances. Returns ------- - livetime : float + livetime The sum of all the individual live-times. """ if not issequenceof(datasets, Dataset): raise TypeError('The datasets argument must be a sequence of Dataset instances!') - livetime = np.sum([ds.livetime for ds in datasets]) + livetimes = [] + for ds in datasets: + if ds.livetime is None: + raise ValueError(f'The livetime of dataset "{ds.name}" is not set and thus cannot be combined!') + livetimes.append(ds.livetime) + + livetime = float(np.sum(livetimes)) return livetime def __init__( self, - name, - exp_pathfilenames, - mc_pathfilenames, - livetime, - default_sub_path_fmt, - version, - verqualifiers=None, - base_path=None, - sub_path_fmt=None, - origin=None, + name: str, + exp_pathfilenames: str | Sequence[str] | None, + mc_pathfilenames: str | Sequence[str] | None, + livetime: float | None, + default_sub_path_fmt: str, + version: int, + verqualifiers: dict | None = None, + base_path: str | None = None, + sub_path_fmt: str | None = None, + origin: 'DatasetOrigin | None' = None, **kwargs, ): """Creates a new dataset object that describes a self-consistent set of @@ -922,38 +966,38 @@ def __init__( Parameters ---------- - name : str + name The name of the dataset. - exp_pathfilenames : str | sequence of str | None + exp_pathfilenames The file name(s), including paths, of the experimental data file(s). This can be None, if a MC-only study is performed. - mc_pathfilenames : str | sequence of str | None + mc_pathfilenames The file name(s), including paths, of the monte-carlo data file(s). This can be None, if a MC-less analysis is performed. - livetime : float | None + livetime The integrated live-time in days of the dataset. It can be None for cases where the live-time is retrieved directly from the data files upon data loading. - default_sub_path_fmt : str + default_sub_path_fmt The default format of the sub path of the data set. This must be a string that can be formatted via the ``format`` method of the ``str`` class. - version : int + version The version number of the dataset. Higher version numbers indicate newer datasets. - verqualifiers : dict | None + verqualifiers If specified, this dictionary specifies version qualifiers. These can be interpreted as subversions of the dataset. The format of the dictionary must be 'qualifier (str): version (int)'. - base_path : str | None + base_path The user-defined base path of the data set. Usually, this is the path of the location of the data directory. If set to ``None``, the dataset configuration's repository ``base_path`` setting is used, which defaults to ``~/.cache/skyllh``. - sub_path_fmt : str | None + sub_path_fmt The user-defined format of the sub path of the data set. If set to ``None``, the ``default_sub_path_fmt`` will be used. - origin : instance of DatasetOrigin | None + origin The instance of DatasetOrigin defining the origin of the dataset, so the dataset can be transferred automatically to the user's device. @@ -1249,18 +1293,18 @@ def data_preparation_functions(self): """ return self._data_preparation_functions - def _gen_datafile_pathfilename_entry(self, pathfilename): + def _gen_datafile_pathfilename_entry(self, pathfilename: str) -> str: """Generates a string containing the given pathfilename and if it exists (FOUND) or not (NOT FOUND). Parameters ---------- - pathfilename : str + pathfilename The fully qualified path and filename of the file. Returns ------- - s : str + s The generated string. """ if os.path.exists(pathfilename): @@ -1297,20 +1341,20 @@ def __gt__(self, ds): # Look for version qualifiers that make this dataset older than the # reference dataset. - qs1 = self._verqualifiers.keys() - qs2 = ds._verqualifiers.keys() + vq1 = self._verqualifiers + vq2 = ds._verqualifiers # If a qualifier of self is also specified for ds, the version number # of the self qualifier must be larger than the version number of the ds # qualifier, in order to consider self as newer dataset. # If a qualifier is present in self but not in ds, self is considered # newer. - for q in qs1: - if q in qs2 and qs1[q] <= qs2[q]: + for q in vq1: + if q in vq2 and vq1[q] <= vq2[q]: return False # If there is a qualifier in ds but not in self, self is considered # older. - return all(q in qs1 for q in qs2) + return all(q in vq1 for q in vq2) def __str__(self): """Implementation of the pretty string representation of the Dataset @@ -1393,20 +1437,20 @@ def remove_data(self): def get_abs_pathfilename_list( self, - pathfilename_list, - ): + pathfilename_list: Sequence[str], + ) -> list[str]: """Returns a list where each entry of the given pathfilename_list is an absolute path. Relative paths will be prefixed with the root_dir property of this Dataset instance. Parameters ---------- - pathfilename_list : sequence of str + pathfilename_list The sequence of file names, either with relative, or absolute paths. Returns ------- - abs_pathfilename_list : list of str + abs_pathfilename_list The list of file names with absolute paths. """ root_dir = self.root_dir @@ -1422,13 +1466,13 @@ def get_abs_pathfilename_list( def get_missing_files( self, - ): + ) -> list[str]: """Determines which files of the dataset are missing and returns the list of files. Returns ------- - missing_files : list of str + missing_files The list of files that are missing. The files are relative to the dataset's root directory. """ @@ -1443,7 +1487,7 @@ def get_missing_files( def update_version_qualifiers( self, - verqualifiers, + verqualifiers: dict, ): """Updates the version qualifiers of the dataset. The update can only be done by increasing the version qualifier integer or by adding new @@ -1451,7 +1495,7 @@ def update_version_qualifiers( Parameters ---------- - verqualifiers : dict + verqualifiers The dictionary with the new version qualifiers. Raises @@ -1472,23 +1516,23 @@ def update_version_qualifiers( got_new_verqualifiers = True existing_verqualifiers_incremented = False - for q in verqualifiers: - if (q in self._verqualifiers) and (verqualifiers[q] > self._verqualifiers[q]): + for q, value in verqualifiers.items(): + if (q in self._verqualifiers) and (value > self._verqualifiers[q]): existing_verqualifiers_incremented = True - self._verqualifiers[q] = verqualifiers[q] + self._verqualifiers[q] = value if not (got_new_verqualifiers or existing_verqualifiers_incremented): raise ValueError('Version qualifier values did not increment and no new version qualifiers were added!') def create_file_list( self, - ): + ) -> list[str]: """Creates the list of files that are linked to this dataset. The file paths are relative to the dataset's root directory. Returns ------- - file_list : list of str + file_list The list of files of this dataset. """ file_list = self._exp_pathfilename_list + self._mc_pathfilename_list @@ -1500,9 +1544,9 @@ def create_file_list( def make_data_available( self, - username=None, - password=None, - ): + username: str | None = None, + password: str | None = None, + ) -> bool: """Makes the data of the dataset available. If the root directory of the dataset does not exist locally, the dataset is transferred from its origin to the local host. If the origin is @@ -1510,16 +1554,16 @@ def make_data_available( Parameters ---------- - username : str | None + username The user name required to connect to the remote host of the origin. If set to ``None``, the - password : str | None + password The password of the user name required to connect to the remote host of the origin. Returns ------- - success : bool + success ``True`` if the data was made available successfully, ``False`` otherwise. """ @@ -1573,6 +1617,7 @@ def make_data_available( if self.origin.is_directory: file_list = [os.path.join(self.origin.sub_path, pathfilename) for pathfilename in self.create_file_list()] else: + assert self.origin.filename is not None file_list = [os.path.join(self.origin.sub_path, self.origin.filename)] logger.info(f'Starting transfer of dataset "{self.name}" from origin into base path "{base_path}".') @@ -1594,13 +1639,13 @@ def make_data_available( def load_data( self, - livetime=None, - keep_fields=None, - dtc_dict=None, - dtc_except_fields=None, - efficiency_mode=None, - tl=None, - ): + livetime: Livetime | float | None = None, + keep_fields: list[str] | None = None, + dtc_dict: dict | None = None, + dtc_except_fields: str | Sequence[str] | None = None, + efficiency_mode: str | None = None, + tl: TimeLord | None = None, + ) -> 'DatasetData': """Loads the data, which is described by the dataset. .. note: @@ -1610,22 +1655,22 @@ def load_data( Parameters ---------- - livetime : instance of Livetime | float | None + livetime If not None, uses this livetime (if float, livetime in days) for the DatasetData instance, otherwise uses the Dataset livetime property value for the DatasetData instance. - keep_fields : list of str | None + keep_fields The list of user-defined data fields that should get loaded and kept in addition to the analysis required data fields. - dtc_dict : dict | None + dtc_dict This dictionary defines how data fields of specific data types (key) should get converted into other data types (value). This can be used to use less memory. If set to None, no data conversion is performed. - dtc_except_fields : str | sequence of str | None + dtc_except_fields The sequence of field names whose data type should not get converted. - efficiency_mode : str | None + efficiency_mode The efficiency mode the data should get loaded with. Possible values are: @@ -1640,12 +1685,12 @@ def load_data( The default value is ``'time'``. If set to ``None``, the default value will be used. - tl : instance of TimeLord | None + tl The TimeLord instance to use to time the data loading procedure. Returns ------- - data : instance of DatasetData + data A instance of DatasetData holding the experimental and monte-carlo data. """ @@ -1690,6 +1735,7 @@ def _conv_new2orig_field_names( + keep_fields, self._exp_field_name_renaming_dict, ) + or [] ) ) @@ -1714,25 +1760,31 @@ def _conv_new2orig_field_names( # But the renaming dictionary can differ for exp and MC fields. keep_fields_mc = list( set( - _conv_new2orig_field_names( - DataFields.get_joint_names( - datafields=datafields, stages=(DFS.DATAPREPARATION_EXP | DFS.ANALYSIS_EXP) + ( + _conv_new2orig_field_names( + DataFields.get_joint_names( + datafields=datafields, stages=(DFS.DATAPREPARATION_EXP | DFS.ANALYSIS_EXP) + ) + + keep_fields, + self._exp_field_name_renaming_dict, ) - + keep_fields, - self._exp_field_name_renaming_dict, + or [] ) - + _conv_new2orig_field_names( - DataFields.get_joint_names( - datafields=datafields, - stages=( - DFS.DATAPREPARATION_EXP - | DFS.ANALYSIS_EXP - | DFS.DATAPREPARATION_MC - | DFS.ANALYSIS_MC - ), + + ( + _conv_new2orig_field_names( + DataFields.get_joint_names( + datafields=datafields, + stages=( + DFS.DATAPREPARATION_EXP + | DFS.ANALYSIS_EXP + | DFS.DATAPREPARATION_MC + | DFS.ANALYSIS_MC + ), + ) + + keep_fields, + self._mc_field_name_renaming_dict, ) - + keep_fields, - self._mc_field_name_renaming_dict, + or [] ) ) ) @@ -1757,21 +1809,21 @@ def _conv_new2orig_field_names( def load_aux_data( self, - name, - tl=None, + name: str, + tl: TimeLord | None = None, ): """Loads the auxiliary data for the given auxiliary data definition. Parameters ---------- - name : str + name The name of the auxiliary data. - tl : instance of TimeLord | None + tl The TimeLord instance to use to time the data loading procedure. Returns ------- - data : unspecified + data The loaded auxiliary data. """ name = str_cast(name, 'The name argument must be cast-able to type str!') @@ -1794,13 +1846,13 @@ def load_aux_data( def add_data_preparation( self, - func, + func: Callable, ): """Adds the given data preparation function to the dataset. Parameters ---------- - func : callable + func The object with call signature __call__(data) that will prepare the data after it was loaded. The argument 'data' is a DatasetData instance holding the experimental and monte-carlo data. The function @@ -1813,13 +1865,13 @@ def add_data_preparation( def remove_data_preparation( self, - key=-1, + key: int | str = -1, ): """Removes a data preparation function from the dataset. Parameters ---------- - key : str, int, optional + key The name or the index of the data preparation function that should be removed. Default value is ``-1``, i.e. the last added function. @@ -1851,17 +1903,17 @@ def remove_data_preparation( def prepare_data( self, - data, - tl=None, + data: 'DatasetData', + tl: TimeLord | None = None, ): """Prepares the data by calling the data preparation callback functions of this dataset. Parameters ---------- - data : instance of DatasetData + data The instance of DatasetData holding the data. - tl : instance of TimeLord | None + tl The instance TimeLord that should be used to time the data preparation. """ @@ -1871,13 +1923,13 @@ def prepare_data( def load_and_prepare_data( self, - livetime=None, - keep_fields=None, - dtc_dict=None, - dtc_except_fields=None, - efficiency_mode=None, - tl=None, - ): + livetime: float | None = None, + keep_fields: Sequence[str] | None = None, + dtc_dict: dict | None = None, + dtc_except_fields: str | Sequence[str] | None = None, + efficiency_mode: str | None = None, + tl: TimeLord | None = None, + ) -> 'DatasetData': """Loads and prepares the experimental and monte-carlo data of this dataset by calling its ``load_data`` and ``prepare_data`` methods. After loading the data it drops all unnecessary data fields if they are @@ -1887,22 +1939,22 @@ def load_and_prepare_data( Parameters ---------- - livetime : float | None + livetime The user-defined livetime in days of the data set. If not set to None, livetime information from the data set will get ignored and this value of the livetime will be used. - keep_fields : sequence of str | None + keep_fields The list of additional data fields that should get kept. By default only the required data fields are kept. - dtc_dict : dict | None + dtc_dict This dictionary defines how data fields of specific data types (key) should get converted into other data types (value). This can be used to use less memory. If set to None, no data conversion is performed. - dtc_except_fields : str | sequence of str | None + dtc_except_fields The sequence of field names whose data type should not get converted. - efficiency_mode : str | None + efficiency_mode The efficiency mode the data should get loaded with. Possible values are: @@ -1917,13 +1969,13 @@ def load_and_prepare_data( The default value is ``'time'``. If set to ``None``, the default value will be used. - tl : instance of TimeLord | None + tl The instance of TimeLord that should be used to time the data loading and preparation. Returns ------- - data : instance of DatasetData + data The instance of DatasetData holding the experimental and monte-carlo data. """ @@ -1970,13 +2022,13 @@ def load_and_prepare_data( def add_binning_definition( self, - binning, + binning: BinningDefinition, ): """Adds a binning setting to this dataset. Parameters ---------- - binning : BinningDefinition + binning The BinningDefinition object holding the binning information. """ if not isinstance(binning, BinningDefinition): @@ -1988,18 +2040,18 @@ def add_binning_definition( def get_binning_definition( self, - name, - ): + name: str, + ) -> BinningDefinition: """Gets the BinningDefinition object for the given binning name. Parameters ---------- - name : str + name The name of the binning definition. Returns ------- - binning_definition : instance of BinningDefinition + binning_definition The requested instance of BinningDefinition. """ if name not in self._binning_definitions: @@ -2008,13 +2060,13 @@ def get_binning_definition( def remove_binning_definition( self, - name, + name: str, ): """Removes the BinningDefinition object from the dataset. Parameters ---------- - name : str + name The name of the binning definition. """ @@ -2027,41 +2079,41 @@ def remove_binning_definition( def has_binning_definition( self, - name, - ): + name: str, + ) -> bool: """Checks if the dataset has a defined binning definition with the given name. Parameters ---------- - name : str + name The name of the binning definition. Returns ------- - check : bool + check True if the binning definition exists, False otherwise. """ return name in self._binning_definitions def define_binning( self, - name, - binedges, - ): + name: str, + binedges: np.ndarray | Sequence, + ) -> BinningDefinition: """Defines a binning for ``name``, and adds it as binning definition. Parameters ---------- - name : str + name The name of the binning setting. - binedges : sequence - The sequence of the bin edges, which should be used for this binning + binedges + The numpy array of the bin edges, which should be used for this binning definition. Returns ------- - binning : instance of BinningDefinition + binning The instance of BinningDefinition which was created and added to this dataset. """ @@ -2069,13 +2121,13 @@ def define_binning( self.add_binning_definition(binning) return binning - def replace_binning_definition(self, binning): + def replace_binning_definition(self, binning: BinningDefinition): """Replaces an already defined binning definition of this dataset by the given binning definition. Parameters ---------- - binning : instance of BinningDefinition + binning The instance of BinningDefinition that will replace the dataset's BinningDefinition instance of the same name. """ @@ -2088,18 +2140,18 @@ def replace_binning_definition(self, binning): def add_aux_data_definition( self, - name, - pathfilenames, + name: str, + pathfilenames: str | Sequence[str], ): """Adds the given data files as auxiliary data definition to the dataset. Parameters ---------- - name : str + name The name of the auxiliary data definition. The name is used as identifier for the data within SkyLLH. - pathfilenames : str | sequence of str + pathfilenames The file name(s) (including paths) of the data file(s). """ name = str_cast( @@ -2120,13 +2172,13 @@ def add_aux_data_definition( def get_aux_data_definition( self, - name, - ): + name: str, + ) -> list[str]: """Returns the auxiliary data definition from the dataset. Parameters ---------- - name : str + name The name of the auxiliary data definition. Raises @@ -2136,7 +2188,7 @@ def get_aux_data_definition( Returns ------- - aux_data_definition : list of str + aux_data_definition The locations (pathfilenames) of the files defined in the auxiliary data as auxiliary data definition. """ @@ -2147,17 +2199,17 @@ def get_aux_data_definition( def set_aux_data_definition( self, - name, - pathfilenames, + name: str, + pathfilenames: str | Sequence[str], ): """Sets the files of the auxiliary data definition, which has the given name. Parameters ---------- - name : str + name The name of the auxiliary data definition. - pathfilenames : str | sequence of str + pathfilenames The file name(s) (including paths) of the data file(s). """ name = str_cast( @@ -2182,13 +2234,13 @@ def set_aux_data_definition( def remove_aux_data_definition( self, - name, + name: str, ): """Removes the auxiliary data definition from the dataset. Parameters ---------- - name : str + name The name of the data definition that should get removed. """ if name not in self._aux_data_definitions: @@ -2200,13 +2252,13 @@ def remove_aux_data_definition( def remove_aux_data_definitions( self, - names, + names: Sequence[str], ): """Removes the auxiliary data definition from the dataset. Parameters ---------- - names : sequence of str + names The names of the data definitions that should get removed. """ for name in names: @@ -2214,16 +2266,16 @@ def remove_aux_data_definitions( def add_aux_data( self, - name, + name: str, data, ): """Adds the given data as auxiliary data to this data set. Parameters ---------- - name : str + name The name under which the auxiliary data will be stored. - data : unspecified + data The data that should get stored. This can be of any data type. Raises @@ -2240,23 +2292,23 @@ def add_aux_data( def get_aux_data( self, - name, - default=None, + name: str, + default: Any = None, ): """Retrieves the auxiliary data that is stored in this data set under the given name. Parameters ---------- - name : str + name The name under which the auxiliary data is stored. - default : any | None + default If not ``None``, it specifies the returned default value when the auxiliary data does not exists. Returns ------- - data : unspecified + data The retrieved auxiliary data. Raises @@ -2276,14 +2328,14 @@ def get_aux_data( def remove_aux_data( self, - name, + name: str, ): """Removes the auxiliary data that is stored in this data set under the given name. Parameters ---------- - name : str + name The name of the dataset that should get removed. """ if name not in self._aux_data: @@ -2299,14 +2351,14 @@ class DatasetCollection: the ``add_datasets`` method. """ - def __init__(self, name, description=''): + def __init__(self, name: str, description: str = ''): """Creates a new DatasetCollection instance. Parameters ---------- - name : str + name The name of the collection. - description : str + description The (longer) description of the dataset collection. """ self.name = name @@ -2357,26 +2409,32 @@ def verqualifiers(self): ds_name = list(self._datasets.keys())[0] # noqa: RUF015 return self._datasets[ds_name].verqualifiers + @overload + def __getitem__(self, key: str) -> 'Dataset': ... + @overload + def __getitem__(self, key: list[str]) -> 'list[Dataset]': ... + @overload + def __getitem__(self, key: tuple[str, ...]) -> 'list[Dataset]': ... def __getitem__( self, - key, - ): + key: str | Sequence[str], + ) -> 'Dataset | list[Dataset]': """Implementation of the access operator ``[key]``. Parameters ---------- - key : str | sequence of str + key The name or names of the dataset(s) that should get retrieved from this dataset collection. Returns ------- - datasets : instance of Dataset | list of instance of Dataset + datasets The dataset instance or the list of dataset instances corresponding to the given key. """ if not issequence(key): - return self.get_dataset(key) + return self.get_dataset(cast(str, key)) if not issequenceof(key, str): raise TypeError( @@ -2413,7 +2471,7 @@ def __str__(self): def add_aux_data( self, - name, + name: str, data, ): """Adds the given data as auxiliary data to all datasets of this @@ -2421,9 +2479,9 @@ def add_aux_data( Parameters ---------- - name : str + name The name under which the auxiliary data will be stored. - data : unspecified + data The data that should get stored. This can be of any data type. Raises @@ -2441,26 +2499,29 @@ def add_aux_data( def add_datasets( self, - datasets, - ): + datasets: 'Dataset | Sequence[Dataset]', + ) -> 'DatasetCollection': """Adds the given Dataset object(s) to this dataset collection. Parameters ---------- - datasets : instance of Dataset | sequence of instance of Dataset + datasets The instance of Dataset or the sequence of instance of Dataset that should be added to the dataset collection. Returns ------- - self : instance of DatasetCollection + self This instance of DatasetCollection in order to be able to chain several ``add_datasets`` calls. """ - if not issequence(datasets): - datasets = [datasets] + _datasets: Sequence[Dataset | Sequence[Dataset]] + if not issequence(datasets): # noqa SIM108 + _datasets = [cast('Dataset', datasets)] + else: + _datasets = cast('Sequence[Dataset]', datasets) - for dataset in datasets: + for dataset in _datasets: if not isinstance(dataset, Dataset): raise TypeError('The dataset object must be a sub-class of Dataset!') @@ -2473,13 +2534,13 @@ def add_datasets( def remove_dataset( self, - name, + name: str, ): """Removes the given dataset from the collection. Parameters ---------- - name : str + name The name of the dataset that should get removed. """ if name not in self._datasets: @@ -2489,18 +2550,18 @@ def remove_dataset( def get_dataset( self, - name, - ): + name: str, + ) -> 'Dataset': """Retrieves a Dataset object from this dataset collection. Parameters ---------- - name : str + name The name of the dataset. Returns ------- - dataset : Dataset instance + dataset The Dataset object holding all the information about the dataset. Raises @@ -2521,18 +2582,18 @@ def get_dataset( def get_datasets( self, - names, - ): + names: str | Sequence[str], + ) -> 'list[Dataset]': """Retrieves a list of Dataset objects from this dataset collection. Parameters ---------- - names : str | sequence of str + names The name or sequence of names of the datasets to retrieve. Returns ------- - datasets : list of Dataset instances + datasets The list of Dataset instances for the given list of data set names. Raises @@ -2541,20 +2602,23 @@ def get_datasets( If one of the requested data sets is not present in this data set collection. """ - if not issequence(names): - names = [names] - if not issequenceof(names, str): + _names: Sequence[str] + if not issequence(names): # noqa SIM108 + _names = [cast(str, names)] + else: + _names = cast(Sequence[str], names) + if not issequenceof(_names, str): raise TypeError('The names argument must be an instance of str or a sequence of str instances!') datasets = [] - for name in names: + for name in _names: datasets.append(self.get_dataset(name)) return datasets def set_exp_field_name_renaming_dict( self, - d, + d: dict, ): """Sets the dictionary with the data field names of the experimental data that needs to be renamed just after loading the data. The @@ -2562,7 +2626,7 @@ def set_exp_field_name_renaming_dict( Parameters ---------- - d : dict + d The dictionary with the old field names as keys and the new field names as values. """ @@ -2571,7 +2635,7 @@ def set_exp_field_name_renaming_dict( def set_mc_field_name_renaming_dict( self, - d, + d: dict, ): """Sets the dictionary with the data field names of the monte-carlo data that needs to be renamed just after loading the data. The @@ -2579,7 +2643,7 @@ def set_mc_field_name_renaming_dict( Parameters ---------- - d : dict + d The dictionary with the old field names as keys and the new field names as values. """ @@ -2588,17 +2652,17 @@ def set_mc_field_name_renaming_dict( def set_dataset_prop( self, - name, - value, + name: str, + value: object, ): """Sets the given property to the given name for all data sets of this data set collection. Parameters ---------- - name : str + name The name of the property. - value : object + value The value to set for the given property. Raises @@ -2613,17 +2677,17 @@ def set_dataset_prop( def define_binning( self, - name, - binedges, + name: str, + binedges: np.ndarray | Sequence, ): """Defines a binning definition and adds it to all the datasets of this dataset collection. Parameters ---------- - name : str + name The name of the binning definition. - binedges : sequence + binedges The sequence of the bin edges, that should be used for the binning. """ for dataset in self._datasets.values(): @@ -2631,14 +2695,14 @@ def define_binning( def add_data_preparation( self, - func, + func: Callable, ): """Adds the data preparation function to all the datasets of this dataset collection. Parameters ---------- - func : callable + func The object with call signature ``__call__(data)`` that will prepare the data after it was loaded. The argument 'data' is the DatasetData instance holding the experimental and monte-carlo data. @@ -2649,14 +2713,14 @@ def add_data_preparation( def remove_data_preparation( self, - key=-1, + key: int | str = -1, ): """Removes data preparation function from all the datasets of this dataset collection. Parameters ---------- - key : str, int, optional + key The name or the index of the data preparation function that should be removed. Default value is ``-1``, i.e. the last added function. @@ -2684,24 +2748,24 @@ def update_version_qualifiers( def load_data( self, - livetime=None, - tl=None, - ppbar=None, + livetime: float | dict[str, float] | None = None, + tl: TimeLord | None = None, + ppbar: ProgressBar | None = None, **kwargs, - ): + ) -> dict: """Loads the data of all data sets of this data set collection. Parameters ---------- - livetime : float | dict of str => float | None + livetime If not None, uses this livetime (in days) as livetime for (all) the DatasetData instances, otherwise uses the live time from the Dataset instance. If a dictionary of data set names and floats is given, it defines the livetime for the individual data sets. - tl : instance of TimeLord | None + tl The instance of TimeLord that should be used to time the data load operation. - ppbar : instance of ProgressBar | None + ppbar The optional parent progress bar. **kwargs Additional keyword arguments are passed to the @@ -2710,7 +2774,7 @@ def load_data( Returns ------- - data_dict : dictionary str => instance of DatasetData + data_dict The dictionary with the DatasetData instance holding the data of an individual data set as value and the data set's name as key. """ @@ -2744,21 +2808,21 @@ class DatasetData: def __init__( self, - data_exp, - data_mc, - livetime, + data_exp: DataFieldRecordArray | None, + data_mc: DataFieldRecordArray | None, + livetime: 'Livetime | float | None', **kwargs, ): """Creates a new DatasetData instance. Parameters ---------- - data_exp : instance of DataFieldRecordArray | None + data_exp The instance of DataFieldRecordArray holding the experimental data. This can be None for a MC-only study. - data_mc : instance of DataFieldRecordArray + data_mc The instance of DataFieldRecordArray holding the monte-carlo data. - livetime : float + livetime The integrated livetime in days of the data. """ super().__init__(**kwargs) @@ -2794,7 +2858,7 @@ def mc(self, data): self._mc = data @property - def livetime(self): + def livetime(self) -> float | None: """The integrated livetime in days of the data. This is None, if there is no live-time provided. """ @@ -2817,22 +2881,26 @@ def exp_field_names(self): @property def mc_field_names(self): - """(read-only) The list of field names present in the monte-carlo data.""" + """(read-only) The list of field names present in the monte-carlo data. + This is an empty list if there is no monte-carlo data available. + """ + if self._mc is None: + return [] return self._mc.field_name_list def assert_data_format( - dataset, - data, + dataset: 'Dataset', + data: 'DatasetData', ): """Checks the format of the experimental and monte-carlo data. Parameters ---------- - dataset : instance of Dataset + dataset The instance of Dataset describing the dataset and holding the local configuration. - data : instance of DatasetData + data The instance of DatasetData holding the actual experimental and simulation data of the data set. @@ -2844,6 +2912,9 @@ def assert_data_format( cfg = dataset.cfg def _get_missing_keys(keys, required_keys): + """Returns the list of required keys that are not present in the given + list of keys. + """ missing_keys = [] for reqkey in required_keys: if reqkey not in keys: @@ -2875,23 +2946,23 @@ def _get_missing_keys(keys, required_keys): def remove_events( - data_exp, - mjds, -): + data_exp: DataFieldRecordArray, + mjds: float | np.ndarray, +) -> DataFieldRecordArray: """Utility function to remove events having the specified MJD time stamps. Parameters ---------- - data_exp : instance of DataFieldRecordArray + data_exp The instance of DataFieldRecordArray holding the experimental data events. - mjds : float | array of floats + mjds The MJD time stamps of the events, that should get removed from the experimental data array. Returns ------- - data_exp : instance of DataFieldRecordArray + data_exp The instance of DataFieldRecordArray holding the experimental data events with the specified events removed. """ @@ -2909,21 +2980,21 @@ def remove_events( def generate_base_path( - default_base_path, - base_path=None, -): + default_base_path: str, + base_path: str | None = None, +) -> str: """Generates the base path. If base_path is None, default_base_path is used. Parameters ---------- - default_base_path : str + default_base_path The default base path if base_path is None. - base_path : str | None + base_path The user-specified base path. Returns ------- - base_path : str + base_path The generated base path. """ if base_path is None: @@ -2937,24 +3008,24 @@ def generate_base_path( def generate_sub_path( - sub_path_fmt, - version, - verqualifiers, -): + sub_path_fmt: str, + version: int, + verqualifiers: dict, +) -> str: """Generates the sub path of the dataset based on the given sub path format. Parameters ---------- - sub_path_fmt : str + sub_path_fmt The format string of the sub path. - version : int + version The version of the dataset. - verqualifiers : dict + verqualifiers The dictionary holding the version qualifiers of the dataset. Returns ------- - sub_path : str + sub_path The generated sub path. """ fmt_dict = dict([('version', version), *verqualifiers.items()]) @@ -2964,13 +3035,13 @@ def generate_sub_path( def generate_data_file_root_dir( - default_base_path, - default_sub_path_fmt, - version, - verqualifiers, - base_path=None, - sub_path_fmt=None, -): + default_base_path: str, + default_sub_path_fmt: str, + version: int, + verqualifiers: dict, + base_path: str | None = None, + sub_path_fmt: str | None = None, +) -> str: """Generates the root directory of the data files based on the given base path and sub path format. If base_path is None, default_base_path is used. If sub_path_fmt is None, default_sub_path_fmt is used. @@ -2986,22 +3057,22 @@ def generate_data_file_root_dir( Parameters ---------- - default_base_path : str + default_base_path The default base path if base_path is None. - default_sub_path_fmt : str + default_sub_path_fmt The default sub path format if sub_path_fmt is None. - version : int + version The version of the data sample. - verqualifiers : dict + verqualifiers The dictionary holding the version qualifiers of the data sample. - base_path : str | None + base_path The user-specified base path. - sub_path_fmt : str | None + sub_path_fmt The user-specified sub path format. Returns ------- - root_dir : str + root_dir The generated root directory of the data files. This will have no trailing directory separator. """ @@ -3022,31 +3093,31 @@ def generate_data_file_root_dir( def get_data_subset( - data, - livetime, - t_start, + data: 'DatasetData', + livetime: Livetime, + t_start: float, t_stop, -): +) -> 'tuple[DatasetData, Livetime]': """Gets instance of DatasetData and instance of Livetime with data subsets between the given time range from ``t_start`` to ``t_stop``. Parameters ---------- - data : DatasetData + data The DatasetData object. - livetime : Livetime + livetime The Livetime object. - t_start : float + t_start The MJD start time of the time range to consider. - t_stop : float + t_stop The MJD stop time of the time range to consider. Returns ------- - data_subset : instance of DatasetData + data_subset The instance of DatasetData with subset of the data between the given time range from ``t_start`` to ``t_stop``. - livetime_subset : instance of Livetime + livetime_subset The instance of Livetime for a subset of the data between the given time range from ``t_start`` to ``t_stop``. """ @@ -3054,6 +3125,8 @@ def get_data_subset( raise TypeError('The "data" argument must be of type DatasetData!') if not isinstance(livetime, Livetime): raise TypeError('The "livetime" argument must be of type Livetime!') + assert data.exp is not None + assert data.mc is not None exp_slice = np.logical_and(data.exp['time'] >= t_start, data.exp['time'] < t_stop) mc_slice = np.logical_and(data.mc['time'] >= t_start, data.mc['time'] < t_stop) diff --git a/skyllh/core/detsigyield.py b/skyllh/core/detsigyield.py index a47dacc39d..2d193f2552 100644 --- a/skyllh/core/detsigyield.py +++ b/skyllh/core/detsigyield.py @@ -1,4 +1,5 @@ import abc +from collections.abc import Callable, Sequence from skyllh.core.config import ( Config, @@ -22,6 +23,7 @@ issequence, issequenceof, ) +from skyllh.core.source_model import SourceModel from skyllh.core.types import ( SourceHypoGroup_t, ) @@ -48,10 +50,10 @@ class DetSigYield( def __init__( self, - param_names, - dataset, - fluxmodel, - livetime, + param_names: Sequence[str], + dataset: Dataset, + fluxmodel: FluxModel, + livetime: float | Livetime, **kwargs, ): """Constructs a new detector signal yield object. It takes @@ -60,14 +62,14 @@ def __init__( Parameters ---------- - param_names : sequence of str + param_names The sequence of parameter names this detector signal yield depends on. These are either fixed or floating parameters. - dataset : Dataset instance + dataset The Dataset instance holding the monte-carlo event data. - fluxmodel : FluxModel + fluxmodel The flux model instance. Must be an instance of FluxModel. - livetime : float | Livetime + livetime The live-time in days to use for the detector signal yield. """ super().__init__(**kwargs) @@ -142,7 +144,7 @@ def livetime(self, lt): @abc.abstractmethod def sources_to_recarray( self, - sources, + sources: SourceModel | Sequence[SourceModel], ): """This method is supposed to convert a (list of) source model(s) into a numpy record array that is understood by the detector signal yield @@ -157,12 +159,12 @@ def sources_to_recarray( Parameters ---------- - sources : SourceModel | sequence of SourceModel + sources The source model(s) containing the information of the source(s). Returns ------- - recarr : numpy record ndarray + recarr The generated (N_sources,)-shaped 1D numpy record ndarray holding the information for each source. """ @@ -178,12 +180,12 @@ def __call__( Parameters ---------- - src_recarray : (N_sources,)-shaped numpy record ndarray + src_recarray The numpy record array containing the information of the sources. The required fields of this record array are implementation dependent. In the most generic case for a point-like source, it must contain the following three fields: ra, dec. - src_params_recarray : (N_sources,)-shaped numpy record ndarray + src_params_recarray The numpy record ndarray containing the parameter values of the sources. The parameter values can be different for the different sources. @@ -196,10 +198,10 @@ def __call__( Returns ------- - detsigyield : (N_sources,)-shaped 1D ndarray of float + detsigyield The array with the mean number of signal in the detector for each given source. - grads : dict + grads The dictionary holding the gradient values for each global fit parameter. The key is the global fit parameter index and the value is the (N_sources,)-shaped numpy ndarray holding the gradient value @@ -255,7 +257,7 @@ def assert_types_of_construct_detsigyield_arguments( f'The ppbar argument must be an instance of ProgressBar! Its current type is {classname(ppbar)}.' ) - def get_detsigyield_construction_factory(self): + def get_detsigyield_construction_factory(self) -> Callable | None: """This method is supposed to return a callable with the call-signature .. code:: @@ -275,7 +277,7 @@ def get_detsigyield_construction_factory(self): Returns ------- - factory : callable | None + factory This default implementation returns ``None``, indicating that a factory is not supported by this builder. """ @@ -284,11 +286,11 @@ def get_detsigyield_construction_factory(self): @abc.abstractmethod def construct_detsigyield( self, - dataset, - data, - shg, - ppbar=None, - ): + dataset: Dataset, + data: DatasetData, + shg: SourceHypoGroup_t, + ppbar: ProgressBar | None = None, + ) -> 'DetSigYield': """Abstract method to construct the DetSigYield instance. This method must be called by the derived class method implementation to ensure the compatibility check of the given flux model with the @@ -296,19 +298,19 @@ def construct_detsigyield( Parameters ---------- - dataset : instance of Dataset + dataset The instance of Dataset holding possible dataset specific settings. - data : instance of DatasetData + data The instance of DatasetData holding the monte-carlo event data. - shg : instance of SourceHypoGroup + shg The instance of SourceHypoGroup (i.e. sources and flux model) for which the detector signal yield should be constructed. - ppbar : instance of ProgressBar | None + ppbar The instance of ProgressBar of the optional parent progress bar. Returns ------- - detsigyield : instance of DetSigYield + detsigyield An instance derived from DetSigYield. """ @@ -320,14 +322,14 @@ class NullDetSigYieldBuilder(DetSigYieldBuilder): def __init__( self, - cfg=None, + cfg: Config | None = None, **kwargs, ): """Creates a new instance of NullDetSigYieldBuilder. Parameters ---------- - cfg : instance of Config | None + cfg The instance of Config holding the local configuration. Since this detector signal yield builder does nothing, this argument is optional. If not provided the default configuration is used. diff --git a/skyllh/core/display.py b/skyllh/core/display.py index 6744479ef4..0766e82ce8 100644 --- a/skyllh/core/display.py +++ b/skyllh/core/display.py @@ -24,19 +24,19 @@ class ANSIColors: UNDERLINE = '\033[4m' -def add_leading_text_line_padding(padwidth, text): +def add_leading_text_line_padding(padwidth: int, text: str) -> str: """Adds leading white spaces to all the lines of the given text. Parameters ---------- - padwidth : int + padwidth The width of the padding. - text : str + text The text with new line characters for each line. Returns ------- - padded_text : str + padded_text The text where each line is padded with the given number of whitespaces. """ return '\n'.join([' ' * padwidth + line for line in text.split('\n')]) diff --git a/skyllh/core/event_selection.py b/skyllh/core/event_selection.py index 6e3c157fa3..66de99e91b 100644 --- a/skyllh/core/event_selection.py +++ b/skyllh/core/event_selection.py @@ -1,5 +1,7 @@ import abc import inspect +from collections.abc import Callable, Sequence +from typing import Literal, cast, overload import numpy as np import scipy.sparse @@ -13,15 +15,20 @@ SourceHypoGroupManager, ) from skyllh.core.source_model import ( + IsPointlike, SourceModel, ) -from skyllh.core.timing import ( - TaskTimer, -) +from skyllh.core.storage import DataFieldRecordArray +from skyllh.core.timing import TaskTimer, TimeLord from skyllh.core.utils.coords import ( angular_separation, ) +# Return type for select_events: always a 2-tuple, optionally 3-tuple with original indices. +_SelectEventsReturn2 = tuple[DataFieldRecordArray, tuple[np.ndarray, np.ndarray]] +_SelectEventsReturn3 = tuple[DataFieldRecordArray, tuple[np.ndarray, np.ndarray], np.ndarray] +_SelectEventsReturn = _SelectEventsReturn2 | _SelectEventsReturn3 + class EventSelectionMethod(metaclass=abc.ABCMeta): """This is the abstract base class for all event selection method classes. @@ -30,12 +37,12 @@ class EventSelectionMethod(metaclass=abc.ABCMeta): are implemented through derived classes of this base class. """ - def __init__(self, shg_mgr, **kwargs): + def __init__(self, shg_mgr: SourceHypoGroupManager | None, **kwargs): """Creates a new event selection method instance. Parameters ---------- - shg_mgr : instance of SourceHypoGroupManager | None + shg_mgr The instance of SourceHypoGroupManager that defines the list of sources, i.e. the list of SourceModel instances. It can be ``None`` if the event selection method does not depend on @@ -66,32 +73,32 @@ def shg_mgr(self): """ return self._shg_mgr - def __and__(self, other): + def __and__(self, other: 'EventSelectionMethod') -> 'IntersectionEventSelectionMethod': """Implements the AND operator (&) for creating an event selection method, which is the intersection of this event selection method and another one using the expression ``intersection = self & other``. Parameters ---------- - other : instance of EventSelectionMethod + other The instance of EventSelectionMethod that is the other event selection method. Returns ------- - intersection : instance of IntersectionEventSelectionMethod + intersection The instance of IntersectionEventSelectionMethod that creates the intersection of this event selection method and the other. """ return IntersectionEventSelectionMethod(self, other) - def change_shg_mgr(self, shg_mgr): + def change_shg_mgr(self, shg_mgr: SourceHypoGroupManager | None): """Changes the SourceHypoGroupManager instance of the event selection method. This will also recreate the internal source numpy record array. Parameters ---------- - shg_mgr : instance of SourceHypoGroupManager | None + shg_mgr The new SourceHypoGroupManager instance, that should be used for this event selection method. It can be ``None`` if the event selection method does not depend on @@ -110,7 +117,7 @@ def change_shg_mgr(self, shg_mgr): self._src_arr = self.sources_to_array(sources=self._shg_mgr.source_list) - def sources_to_array(self, sources): + def sources_to_array(self, sources: Sequence[SourceModel]) -> np.ndarray | None: """This method is supposed to convert a sequence of SourceModel instances into a structured numpy ndarray with the source information in a format that is best understood by the actual event selection @@ -118,49 +125,72 @@ def sources_to_array(self, sources): Parameters ---------- - sources : sequence of SourceModel + sources The sequence of source models containing the necessary information of the source. Returns ------- - arr : numpy record ndarray | None + arr The generated numpy record ndarray holding the necessary information for each source. By default ``None`` is returned. """ return + @overload + def select_events( + self, + events: DataFieldRecordArray, + src_evt_idxs: tuple | None = None, + ret_original_evt_idxs: Literal[False] = False, + tl: TimeLord | None = None, + ) -> _SelectEventsReturn2: ... + @overload + def select_events( + self, + events: DataFieldRecordArray, + src_evt_idxs: tuple | None = None, + *, + ret_original_evt_idxs: Literal[True], + tl: TimeLord | None = None, + ) -> _SelectEventsReturn3: ... @abc.abstractmethod - def select_events(self, events, src_evt_idxs=None, ret_original_evt_idxs=False, tl=None): + def select_events( + self, + events: DataFieldRecordArray, + src_evt_idxs: tuple | None = None, + ret_original_evt_idxs: bool = False, + tl: TimeLord | None = None, + ) -> _SelectEventsReturn: """This method selects the events, which will contribute to the log-likelihood ratio function. Parameters ---------- - events : instance of DataFieldRecordArray + events The instance of DataFieldRecordArray of length N_events, holding the events. - src_evt_idxs : 2-tuple of 1d ndarrays of ints | None + src_evt_idxs The 2-element tuple holding the two 1d ndarrays of int of length N_values, specifying to which sources the given events belong to. - ret_original_evt_idxs : bool + ret_original_evt_idxs Flag if the original indices of the selected events should get returned as well. - tl : instance of TimeLord | None + tl The optional instance of TimeLord that should be used to collect timing information about this method. Returns ------- - selected_events : instance of DataFieldRecordArray + selected_events The instance of DataFieldRecordArray of length N_selected_events, holding the selected events, i.e. a subset of the ``events`` argument. (src_idxs, evt_idxs) : 1d ndarrays of ints The two 1d ndarrays of int of length N_values, holding the indices of the sources and the selected events. - original_evt_idxs : 1d ndarray of ints + original_evt_idxs The (N_selected_events,)-shaped numpy ndarray holding the original indices of the selected events, if ``ret_original_evt_idxs`` is set to ``True``. @@ -173,16 +203,16 @@ class IntersectionEventSelectionMethod(EventSelectionMethod): ``evt_sel_method1 & evt_sel_method2``. """ - def __init__(self, evt_sel_method1, evt_sel_method2, **kwargs): + def __init__(self, evt_sel_method1: 'EventSelectionMethod', evt_sel_method2: 'EventSelectionMethod', **kwargs): """Creates a compounded event selection method of two given event selection methods. Parameters ---------- - evt_sel_method1 : instance of EventSelectionMethod + evt_sel_method1 The instance of EventSelectionMethod for the first event selection method. - evt_sel_method2 : instance of EventSelectionMethod + evt_sel_method2 The instance of EventSelectionMethod for the second event selection method. """ @@ -225,14 +255,14 @@ def evt_sel_method2(self, method): ) self._evt_sel_method2 = method - def change_shg_mgr(self, shg_mgr): + def change_shg_mgr(self, shg_mgr: SourceHypoGroupManager | None): """Changes the SourceHypoGroupManager instance of the event selection method. This will call the ``change_shg_mgr`` of the individual event selection methods. Parameters ---------- - shg_mgr : instance of SourceHypoGroupManager | None + shg_mgr The new SourceHypoGroupManager instance, that should be used for this event selection method. It can be ``None`` if the event selection method does not depend on @@ -241,52 +271,87 @@ def change_shg_mgr(self, shg_mgr): self._evt_sel_method1.change_shg_mgr(shg_mgr=shg_mgr) self._evt_sel_method2.change_shg_mgr(shg_mgr=shg_mgr) - def select_events(self, events, src_evt_idxs=None, ret_original_evt_idxs=False, tl=None): + @overload + def select_events( + self, + events: DataFieldRecordArray, + src_evt_idxs: tuple | None = None, + ret_original_evt_idxs: Literal[False] = False, + tl: TimeLord | None = None, + ) -> _SelectEventsReturn2: ... + @overload + def select_events( + self, + events: DataFieldRecordArray, + src_evt_idxs: tuple | None = None, + *, + ret_original_evt_idxs: Literal[True], + tl: TimeLord | None = None, + ) -> _SelectEventsReturn3: ... + def select_events( + self, + events: DataFieldRecordArray, + src_evt_idxs: tuple | None = None, + ret_original_evt_idxs: bool = False, + tl: TimeLord | None = None, + ) -> _SelectEventsReturn: """Selects events by calling the ``select_events`` methods of the individual event selection methods. Parameters ---------- - events : instance of DataFieldRecordArray + events The instance of DataFieldRecordArray holding the events. - src_evt_idxs : 2-tuple of 1d ndarrays of ints | None + src_evt_idxs The 2-element tuple holding the two 1d ndarrays of int of length N_values, specifying to which sources the given events belong to. - ret_original_evt_idxs : bool + ret_original_evt_idxs Flag if the original indices of the selected events should get returned as well. - tl : instance of TimeLord | None + tl The optional instance of TimeLord that should be used to collect timing information about this method. Returns ------- - selected_events : DataFieldRecordArray + selected_events The instance of DataFieldRecordArray holding the selected events, i.e. a subset of the `events` argument. (src_idxs, evt_idxs) : 1d ndarrays of ints The indices of the sources and the selected events. - original_evt_idxs : 1d ndarray of ints + original_evt_idxs The (N_selected_events,)-shaped numpy ndarray holding the original indices of the selected events, if ``ret_original_evt_idxs`` is set to ``True``. """ if ret_original_evt_idxs: - (events, src_evt_idxs, org_evt_idxs1) = self._evt_sel_method1.select_events( - events=events, src_evt_idxs=src_evt_idxs, ret_original_evt_idxs=True + (events, src_evt_idxs, org_evt_idxs1) = cast( + 'tuple[DataFieldRecordArray, tuple[np.ndarray, np.ndarray], np.ndarray]', + self._evt_sel_method1.select_events( + events=events, src_evt_idxs=src_evt_idxs, ret_original_evt_idxs=True + ), ) - (events, src_evt_idxs, org_evt_idxs2) = self._evt_sel_method2.select_events( - events=events, src_evt_idxs=src_evt_idxs, ret_original_evt_idxs=True + (events, src_evt_idxs, org_evt_idxs2) = cast( + 'tuple[DataFieldRecordArray, tuple[np.ndarray, np.ndarray], np.ndarray]', + self._evt_sel_method2.select_events( + events=events, src_evt_idxs=src_evt_idxs, ret_original_evt_idxs=True + ), ) org_evt_idxs = np.take(org_evt_idxs1, org_evt_idxs2) return (events, src_evt_idxs, org_evt_idxs) - (events, src_evt_idxs) = self._evt_sel_method1.select_events(events=events, src_evt_idxs=src_evt_idxs) + (events, src_evt_idxs) = cast( + 'tuple[DataFieldRecordArray, tuple[np.ndarray, np.ndarray]]', + self._evt_sel_method1.select_events(events=events, src_evt_idxs=src_evt_idxs), + ) - (events, src_evt_idxs) = self._evt_sel_method2.select_events(events=events, src_evt_idxs=src_evt_idxs) + (events, src_evt_idxs) = cast( + 'tuple[DataFieldRecordArray, tuple[np.ndarray, np.ndarray]]', + self._evt_sel_method2.select_events(events=events, src_evt_idxs=src_evt_idxs), + ) return (events, src_evt_idxs) @@ -294,12 +359,12 @@ def select_events(self, events, src_evt_idxs=None, ret_original_evt_idxs=False, class AllEventSelectionMethod(EventSelectionMethod): """This event selection method selects all events.""" - def __init__(self, shg_mgr): + def __init__(self, shg_mgr: SourceHypoGroupManager): """Creates a new event selection method instance. Parameters ---------- - shg_mgr : instance of SourceHypoGroupManager + shg_mgr The instance of SourceHypoGroupManager that defines the list of sources, i.e. the list of SourceModel instances. For this particular event selection method it has no meaning, but it is an interface @@ -307,54 +372,78 @@ def __init__(self, shg_mgr): """ super().__init__(shg_mgr=shg_mgr) - def sources_to_array(self, sources): + def sources_to_array(self, sources) -> None: """Creates the source array from the given list of sources. This event selection method does not depend on the sources. Hence, ``None`` is returned. Returns ------- - arr : None + arr The generated numpy record ndarray holding the necessary information for each source. Since this event selection method does not depend on any source, ``None`` is returned. """ return - def select_events(self, events, src_evt_idxs=None, ret_original_evt_idxs=False, tl=None): + @overload + def select_events( + self, + events: DataFieldRecordArray, + src_evt_idxs: tuple | None = None, + ret_original_evt_idxs: Literal[False] = False, + tl: TimeLord | None = None, + ) -> _SelectEventsReturn2: ... + @overload + def select_events( + self, + events: DataFieldRecordArray, + src_evt_idxs: tuple | None = None, + *, + ret_original_evt_idxs: Literal[True], + tl: TimeLord | None = None, + ) -> _SelectEventsReturn3: ... + def select_events( + self, + events: DataFieldRecordArray, + src_evt_idxs: tuple | None = None, + ret_original_evt_idxs: bool = False, + tl: TimeLord | None = None, + ) -> _SelectEventsReturn: """Selects all of the given events. Hence, the returned event array is the same as the given array. Parameters ---------- - events : instance of DataFieldRecordArray + events The instance of DataFieldRecordArray holding the events, for which the selection method should get applied. - src_evt_idxs : 2-tuple of 1d ndarrays of ints | None + src_evt_idxs The 2-element tuple holding the two 1d ndarrays of int of length N_values, specifying to which sources the given events belong to. - ret_original_evt_idxs : bool + ret_original_evt_idxs Flag if the original indices of the selected events should get returned as well. - tl : instance of TimeLord | None + tl The optional instance of TimeLord that should be used to collect timing information about this method. Returns ------- - selected_events : DataFieldRecordArray + selected_events The instance of DataFieldRecordArray holding the selected events, i.e. a subset of the `events` argument. (src_idxs, evt_idxs) : 1d ndarrays of ints The indices of sources and the selected events. - original_evt_idxs : 1d ndarray of ints + original_evt_idxs The (N_selected_events,)-shaped numpy ndarray holding the original indices of the selected events, if ``ret_original_evt_idxs`` is set to ``True``. """ with TaskTimer(tl, 'ESM: Calculate indices of selected events.'): if src_evt_idxs is None: - n_sources = self.shg_mgr.n_sources + assert self._shg_mgr is not None + n_sources = self._shg_mgr.n_sources src_idxs = np.repeat(np.arange(n_sources), len(events)) evt_idxs = np.tile(events.indices, n_sources) else: @@ -371,31 +460,31 @@ class SpatialEventSelectionMethod(EventSelectionMethod, metaclass=abc.ABCMeta): selection methods. """ - def __init__(self, shg_mgr, **kwargs): + def __init__(self, shg_mgr: SourceHypoGroupManager, **kwargs): """Creates a new event selection method instance. Parameters ---------- - shg_mgr : instance of SourceHypoGroupManager + shg_mgr The instance of SourceHypoGroupManager that defines the list of sources, i.e. the list of SourceModel instances. """ super().__init__(shg_mgr=shg_mgr, **kwargs) - def sources_to_array(self, sources): + def sources_to_array(self, sources: Sequence[SourceModel]) -> np.ndarray: """Converts the given sequence of SourceModel instances into a structured numpy ndarray holding the necessary source information needed for this event selection method. Parameters ---------- - sources : sequence of SourceModel + sources The sequence of source models containing the necessary information of the source. Returns ------- - arr : numpy record ndarray + arr The generated numpy record ndarray holding the necessary information for each source. It contains the following data fields: 'ra', 'dec'. """ @@ -409,8 +498,9 @@ def sources_to_array(self, sources): arr = np.empty((len(sources),), dtype=[('ra', np.float64), ('dec', np.float64)], order='F') for i, src in enumerate(sources): - arr['ra'][i] = src.ra - arr['dec'][i] = src.dec + pointlike_src = cast(IsPointlike, src) + arr['ra'][i] = pointlike_src.ra + arr['dec'][i] = pointlike_src.dec return arr @@ -420,16 +510,16 @@ class DecBandEventSectionMethod(SpatialEventSelectionMethod): around a list of point-like source positions. """ - def __init__(self, shg_mgr, delta_angle): + def __init__(self, shg_mgr: SourceHypoGroupManager, delta_angle: float): """Creates and configures a spatial declination band event selection method object. Parameters ---------- - shg_mgr : instance of SourceHypoGroupManager + shg_mgr The instance of SourceHypoGroupManager that defines the list of sources, i.e. the list of SourceModel instances. - delta_angle : float + delta_angle The half-opening angle around the source in declination for which events should get selected. """ @@ -449,41 +539,65 @@ def delta_angle(self, angle): angle = float_cast(angle, 'The delta_angle property must be castable to type float!') self._delta_angle = angle - def select_events(self, events, src_evt_idxs=None, ret_original_evt_idxs=False, tl=None): + @overload + def select_events( + self, + events: DataFieldRecordArray, + src_evt_idxs: tuple | None = None, + ret_original_evt_idxs: Literal[False] = False, + tl: TimeLord | None = None, + ) -> _SelectEventsReturn2: ... + @overload + def select_events( + self, + events: DataFieldRecordArray, + src_evt_idxs: tuple | None = None, + *, + ret_original_evt_idxs: Literal[True], + tl: TimeLord | None = None, + ) -> _SelectEventsReturn3: ... + def select_events( + self, + events: DataFieldRecordArray, + src_evt_idxs: tuple | None = None, + ret_original_evt_idxs: bool = False, + tl: TimeLord | None = None, + ) -> _SelectEventsReturn: """Selects the events within the declination band. Parameters ---------- - events : instance of DataFieldRecordArray + events The instance of DataFieldRecordArray that holds the event data. The following data fields must exist: ``'dec'`` : float The declination of the event. - src_evt_idxs : 2-tuple of 1d ndarrays of ints | None + src_evt_idxs The 2-element tuple holding the two 1d ndarrays of int of length N_values, specifying to which sources the given events belong to. - ret_original_evt_idxs : bool + ret_original_evt_idxs Flag if the original indices of the selected events should get returned as well. - tl : instance of TimeLord | None + tl The optional instance of TimeLord that should be used to collect timing information about this method. Returns ------- - selected_events : instance of DataFieldRecordArray + selected_events The instance of DataFieldRecordArray holding only the selected events. (src_idxs, evt_idxs) : 1d ndarrays of ints The indices of sources and the selected events. - original_evt_idxs : 1d ndarray of ints + original_evt_idxs The (N_selected_events,)-shaped numpy ndarray holding the original indices of the selected events, if ``ret_original_evt_idxs`` is set to ``True``. """ delta_angle = self._delta_angle + assert self._src_arr is not None src_arr = self._src_arr # Calculates the minus and plus declination around each source and @@ -526,16 +640,16 @@ class RABandEventSectionMethod(SpatialEventSelectionMethod): around a list of point-like source positions. """ - def __init__(self, shg_mgr, delta_angle): + def __init__(self, shg_mgr: SourceHypoGroupManager, delta_angle: float): """Creates and configures a right-ascension band event selection method object. Parameters ---------- - shg_mgr : instance of SourceHypoGroupManager + shg_mgr The instance of SourceHypoGroupManager that defines the list of sources, i.e. the list of SourceModel instances. - delta_angle : float + delta_angle The half-opening angle around the source in right-ascension for which events should get selected. """ @@ -555,7 +669,30 @@ def delta_angle(self, angle): angle = float_cast(angle, 'The delta_angle property must be castable to type float!') self._delta_angle = angle - def select_events(self, events, src_evt_idxs=None, ret_original_evt_idxs=False, tl=None): + @overload + def select_events( + self, + events: DataFieldRecordArray, + src_evt_idxs: tuple | None = None, + ret_original_evt_idxs: Literal[False] = False, + tl: TimeLord | None = None, + ) -> _SelectEventsReturn2: ... + @overload + def select_events( + self, + events: DataFieldRecordArray, + src_evt_idxs: tuple | None = None, + *, + ret_original_evt_idxs: Literal[True], + tl: TimeLord | None = None, + ) -> _SelectEventsReturn3: ... + def select_events( + self, + events: DataFieldRecordArray, + src_evt_idxs: tuple | None = None, + ret_original_evt_idxs: bool = False, + tl: TimeLord | None = None, + ) -> _SelectEventsReturn: """Selects the events within the right-ascention band. The solid angle dOmega = dRA * dSinDec = dRA * dDec * cos(dec) is a @@ -564,7 +701,7 @@ def select_events(self, events, src_evt_idxs=None, ret_original_evt_idxs=False, Parameters ---------- - events : instance of DataFieldRecordArray + events The instance of DataFieldRecordArray that holds the event data. The following data fields must exist: @@ -573,29 +710,30 @@ def select_events(self, events, src_evt_idxs=None, ret_original_evt_idxs=False, ``'dec'`` : float The declination of the event. - src_evt_idxs : 2-tuple of 1d ndarrays of ints | None + src_evt_idxs The 2-element tuple holding the two 1d ndarrays of int of length N_values, specifying to which sources the given events belong to. - ret_original_evt_idxs : bool + ret_original_evt_idxs Flag if the original indices of the selected events should get returned as well. - tl : instance of TimeLord | None + tl The optional instance of TimeLord that should be used to collect timing information about this method. Returns ------- - selected_events : instance of DataFieldRecordArray + selected_events The instance of DataFieldRecordArray holding only the selected events. (src_idxs, evt_idxs) : 1d ndarrays of ints The indices of the sources and the selected events. - original_evt_idxs : 1d ndarray of ints + original_evt_idxs The (N_selected_events,)-shaped numpy ndarray holding the original indices of the selected events, if ``ret_original_evt_idxs`` is set to ``True``. """ delta_angle = self._delta_angle + assert self._src_arr is not None src_arr = self._src_arr # Get the minus and plus declination around the sources. @@ -655,15 +793,15 @@ class SpatialBoxEventSelectionMethod(SpatialEventSelectionMethod): positions. """ - def __init__(self, shg_mgr, delta_angle): + def __init__(self, shg_mgr: SourceHypoGroupManager, delta_angle: float): """Creates and configures a spatial box event selection method object. Parameters ---------- - shg_mgr : instance of SourceHypoGroupManager + shg_mgr The instance of SourceHypoGroupManager that defines the list of sources, i.e. the list of SourceModel instances. - delta_angle : float + delta_angle The half-opening angle around the source for which events should get selected. """ @@ -683,7 +821,30 @@ def delta_angle(self, angle): angle = float_cast(angle, 'The delta_angle property must be castable to type float!') self._delta_angle = angle - def select_events(self, events, src_evt_idxs=None, ret_original_evt_idxs=False, tl=None): + @overload + def select_events( + self, + events: DataFieldRecordArray, + src_evt_idxs: tuple | None = None, + ret_original_evt_idxs: Literal[False] = False, + tl: TimeLord | None = None, + ) -> _SelectEventsReturn2: ... + @overload + def select_events( + self, + events: DataFieldRecordArray, + src_evt_idxs: tuple | None = None, + *, + ret_original_evt_idxs: Literal[True], + tl: TimeLord | None = None, + ) -> _SelectEventsReturn3: ... + def select_events( + self, + events: DataFieldRecordArray, + src_evt_idxs: tuple | None = None, + ret_original_evt_idxs: bool = False, + tl: TimeLord | None = None, + ) -> _SelectEventsReturn: """Selects the events within the spatial box in right-ascention and declination. @@ -693,7 +854,7 @@ def select_events(self, events, src_evt_idxs=None, ret_original_evt_idxs=False, Parameters ---------- - events : instance of DataFieldRecordArray + events The instance of DataFieldRecordArray that holds the event data. The following data fields must exist: @@ -702,29 +863,30 @@ def select_events(self, events, src_evt_idxs=None, ret_original_evt_idxs=False, ``'dec'`` : float The declination of the event. - src_evt_idxs : 2-tuple of 1d ndarrays of ints | None + src_evt_idxs The 2-element tuple holding the two 1d ndarrays of int of length N_values, specifying to which sources the given events belong to. - ret_original_evt_idxs : bool + ret_original_evt_idxs Flag if the original indices of the selected events should get returned as well. - tl : instance of TimeLord | None + tl The optional instance of TimeLord that should be used to collect timing information about this method. Returns ------- - selected_events : instance of DataFieldRecordArray + selected_events The instance of DataFieldRecordArray holding only the selected events. (src_idxs, evt_idxs) : 1d ndarrays of ints | None The indices of sources and the selected events. - original_evt_idxs : 1d ndarray of ints + original_evt_idxs The (N_selected_events,)-shaped numpy ndarray holding the original indices of the selected events, if ``ret_original_evt_idxs`` is set to ``True``. """ delta_angle = self._delta_angle + assert self._src_arr is not None src_arr = self._src_arr n_sources = len(src_arr) @@ -814,17 +976,17 @@ class PsiFuncEventSelectionMethod(EventSelectionMethod): of the provided function. """ - def __init__(self, shg_mgr, psi_name, func, axis_name_list): + def __init__(self, shg_mgr: SourceHypoGroupManager, psi_name: str, func: Callable, axis_name_list: list[str]): """Creates a new PsiFuncEventSelectionMethod instance. Parameters ---------- - shg_mgr : instance of SourceHypoGroupManager + shg_mgr The instance of SourceHypoGroupManager that defines the list of sources, i.e. the list of SourceModel instances. - psi_name : str + psi_name The name of the data field that provides the psi value of the event. - func : callable + func The function that should get evaluated for each event. The call signature must be @@ -833,7 +995,7 @@ def __init__(self, shg_mgr, psi_name, func, axis_name_list): where ``*axis_data`` is the event data of each required axis. The number of axes must match the provided axis names through the ``axis_name_list``. - axis_name_list : list of str + axis_name_list The list of data field names for each axis of the function ``func``. All field names must be valid field names of the trial data's DataFieldRecordArray instance. @@ -852,7 +1014,8 @@ def __init__(self, shg_mgr, psi_name, func, axis_name_list): f'of arguments is {n_func_args}.' ) - n_sources = self.shg_mgr.n_sources + assert self._shg_mgr is not None + n_sources = self._shg_mgr.n_sources if n_sources != 1: raise ValueError( 'The `PsiFuncEventSelectionMethod.select_events` currently ' @@ -905,13 +1068,36 @@ def axis_name_list(self, names): ) self._axis_name_list = list(names) - def select_events(self, events, src_evt_idxs=None, ret_original_evt_idxs=False, tl=None): + @overload + def select_events( + self, + events: DataFieldRecordArray, + src_evt_idxs: tuple | None = None, + ret_original_evt_idxs: Literal[False] = False, + tl: TimeLord | None = None, + ) -> _SelectEventsReturn2: ... + @overload + def select_events( + self, + events: DataFieldRecordArray, + src_evt_idxs: tuple | None = None, + *, + ret_original_evt_idxs: Literal[True], + tl: TimeLord | None = None, + ) -> _SelectEventsReturn3: ... + def select_events( + self, + events: DataFieldRecordArray, + src_evt_idxs: tuple | None = None, + ret_original_evt_idxs: bool = False, + tl: TimeLord | None = None, + ) -> _SelectEventsReturn: """Selects the events whose psi value is smaller than the value of the predefined function. Parameters ---------- - events : instance of DataFieldRecordArray + events The instance of DataFieldRecordArray that holds the event data. The following data fields must exist: @@ -921,24 +1107,24 @@ def select_events(self, events, src_evt_idxs=None, ret_original_evt_idxs=False, The name of the axis required for the function ``func`` to be evaluated. - src_evt_idxs : 2-tuple of 1d ndarrays of ints | None + src_evt_idxs The 2-element tuple holding the two 1d ndarrays of int of length N_values, specifying to which sources the given events belong to. - ret_original_evt_idxs : bool + ret_original_evt_idxs Flag if the original indices of the selected events should get returned as well. - tl : instance of TimeLord | None + tl The optional instance of TimeLord that should be used to collect timing information about this method. Returns ------- - selected_events : instance of DataFieldRecordArray + selected_events The instance of DataFieldRecordArray holding only the selected events. (src_idxs, evt_idxs) : 1d ndarrays of ints The indices of the sources and the selected events. - original_evt_idxs : 1d ndarray of ints + original_evt_idxs The (N_selected_events,)-shaped numpy ndarray holding the original indices of the selected events, if ``ret_original_evt_idxs`` is set to ``True``. @@ -978,29 +1164,23 @@ class AngErrOfPsiEventSelectionMethod(SpatialEventSelectionMethod): is larger than the value of the provided function at a given psi value. """ - def __init__(self, shg_mgr, func, psi_floor=None, **kwargs): + def __init__(self, shg_mgr: SourceHypoGroupManager, func: Callable, psi_floor: float | None = None, **kwargs): """Creates and configures a spatial box and psi func event selection method object. Parameters ---------- - shg_mgr : instance of SourceHypoGroupManager + shg_mgr The instance of SourceHypoGroupManager that defines the list of sources, i.e. the list of SourceModel instances. - delta_angle : float - The half-opening angle around the source for which events should - get selected. - psi_name : str | None - The name of the data field that provides the psi value of the event. - If set to ``None``, the psi value will be calculated automatically. - func : callable + func The function that should get evaluated for each event. The call signature must be ``func(psi)``, where ``psi`` is the opening angle between the source and the event. - psi_floor : float | None + psi_floor The psi func event selection is excluded for events having psi value below the ``psi_floor``. If None, set it to default 5 degrees. """ @@ -1039,7 +1219,30 @@ def psi_floor(self, psi): psi = float_cast(psi, 'The psi_floor property must be castable to type float!') self._psi_floor = psi - def select_events(self, events, src_evt_idxs=None, ret_original_evt_idxs=False, tl=None): + @overload + def select_events( + self, + events: DataFieldRecordArray, + src_evt_idxs: tuple | None = None, + ret_original_evt_idxs: Literal[False] = False, + tl: TimeLord | None = None, + ) -> _SelectEventsReturn2: ... + @overload + def select_events( + self, + events: DataFieldRecordArray, + src_evt_idxs: tuple | None = None, + *, + ret_original_evt_idxs: Literal[True], + tl: TimeLord | None = None, + ) -> _SelectEventsReturn3: ... + def select_events( + self, + events: DataFieldRecordArray, + src_evt_idxs: tuple | None = None, + ret_original_evt_idxs: bool = False, + tl: TimeLord | None = None, + ) -> _SelectEventsReturn: """Selects the events within the spatial box in right-ascention and declination and performs an additional selection of events whose ang_err value is larger than the value of the provided function at a given psi @@ -1051,7 +1254,7 @@ def select_events(self, events, src_evt_idxs=None, ret_original_evt_idxs=False, Parameters ---------- - events : instance of DataFieldRecordArray + events The instance of DataFieldRecordArray that holds the event data. The following data fields must exist: @@ -1060,32 +1263,34 @@ def select_events(self, events, src_evt_idxs=None, ret_original_evt_idxs=False, ``'dec'`` : float The declination of the event. - src_evt_idxs : 2-tuple of 1d ndarrays of ints | None + src_evt_idxs The 2-element tuple holding the two 1d ndarrays of int of length N_values, specifying to which sources the given events belong to. If set to ``None`` all given events will be considered to for all sources. - ret_original_evt_idxs : bool + ret_original_evt_idxs Flag if the original indices of the selected events should get returned as well. - tl : instance of TimeLord | None + tl The optional instance of TimeLord that should be used to collect timing information about this method. Returns ------- - selected_events : instance of DataFieldRecordArray + selected_events The instance of DataFieldRecordArray holding only the selected events. (src_idxs, evt_idxs) : 1d ndarrays of ints The indices of the sources and the selected events. - original_evt_idxs : 1d ndarray of ints + original_evt_idxs The (N_selected_events,)-shaped numpy ndarray holding the original indices of the selected events, if ``ret_original_evt_idxs`` is set to ``True``. """ + assert self._src_arr is not None + src_arr = self._src_arr if src_evt_idxs is None: - n_sources = len(self._src_arr) + n_sources = len(src_arr) n_events = len(events) src_idxs = np.repeat(np.arange(n_sources), n_events) evt_idxs = np.tile(np.arange(n_events), n_sources) @@ -1095,8 +1300,8 @@ def select_events(self, events, src_evt_idxs=None, ret_original_evt_idxs=False, # Perform selection based on psi values. with TaskTimer(tl, 'ESM: Calculate psi values.'): psi = angular_separation( - ra1=np.take(self._src_arr['ra'], src_idxs), - dec1=np.take(self._src_arr['dec'], src_idxs), + ra1=np.take(src_arr['ra'], src_idxs), + dec1=np.take(src_arr['dec'], src_idxs), ra2=np.take(events['ra'], evt_idxs), dec2=np.take(events['dec'], evt_idxs), ) @@ -1108,7 +1313,7 @@ def select_events(self, events, src_evt_idxs=None, ret_original_evt_idxs=False, # Have to define the shape argument in order to not truncate # the mask in case last events are not selected. mask_sky = scipy.sparse.csr_matrix( - (mask_psi, (src_idxs, evt_idxs)), shape=(len(self._src_arr), len(events)) + (mask_psi, (src_idxs, evt_idxs)), shape=(len(src_arr), len(events)) ).toarray() mask = np.any(mask_sky, axis=0) diff --git a/skyllh/core/expectation_maximization.py b/skyllh/core/expectation_maximization.py index ce43170007..09953de34a 100644 --- a/skyllh/core/expectation_maximization.py +++ b/skyllh/core/expectation_maximization.py @@ -3,37 +3,37 @@ def em_expectation_step( - ns, - mu, - sigma, - t, - sob, -): + ns: np.ndarray | list, + mu: np.ndarray | list, + sigma: np.ndarray | list, + t: np.ndarray, + sob: np.ndarray, +) -> tuple[np.ndarray, float]: """Expectation step of expectation maximization algorithm. Parameters ---------- - ns : instance of ndarray + ns The (n_flares,)-shaped numpy ndarray holding the number of signal neutrinos, as weight for each gaussian flare. - mu : instance of ndarray + mu The (n_flares,)-shaped numpy ndarray holding the mean for each gaussian flare. - sigma: instance of ndarray + sigma The (n_flares,)-shaped numpy ndarray holding the sigma for each gaussian flare. - t : instance of ndarray + t The (n_events,)-shaped numpy ndarray holding the time of each event. - sob : instance of ndarray + sob The (n_events,)-shaped numpy ndarray holding the signal-over-background values of each event. Returns ------- - expectations : instane of ndarray + expectations The (n_flares, n_events)-shaped numpy ndarray holding the expectation of each flare and event. - llh : float + llh The log-likelihood value, which is the sum of log of the signal and background expectations. """ @@ -43,7 +43,7 @@ def em_expectation_step( N = len(t) e_sig = np.empty((n_flares, N), dtype=np.float64) for i in range(n_flares): - e_sig[i] = norm(loc=mu[i], scale=sigma[i]).pdf(t) + e_sig[i] = norm(loc=mu[i], scale=sigma[i]).pdf(t) # pyright: ignore[reportAttributeAccessIssue] e_sig[i] *= sob e_sig[i] *= ns[i] e_bkg = (N - np.sum(ns)) / (np.max(t) - np.min(t)) / b_term @@ -56,81 +56,81 @@ def em_expectation_step( def em_maximization_step( - e, - t, -): + e: np.ndarray, + t: np.ndarray, +) -> tuple[list[float], list[float], list[float]]: """The maximization step of the expectation maximization algorithm. Parameters ---------- - e : instance of ndarray + e The (n_flares, n_events)-shaped numpy ndarray holding the expectation for each event and flare. - t : 1d ndarray of float + t The times of each event. Returns ------- - mu : list of float + mu Best fit mean time of the gaussian flare. - sigma : list of float + sigma Best fit sigma of the gaussian flare. - ns : list of float + ns Best fit number of signal neutrinos, as weight for the gaussian flare. """ - mu = [] - sigma = [] - ns = [] + mu: list[float] = [] + sigma: list[float] = [] + ns: list[float] = [] for i in range(e.shape[0]): - mu.append(np.average(t, weights=e[i])) - sigma.append(np.sqrt(np.average(np.square(t - mu[i]), weights=e[i]))) - ns.append(np.sum(e[i])) + mu.append(float(np.average(t, weights=e[i]))) + sigma.append(float(np.sqrt(np.average(np.square(t - mu[i]), weights=e[i])))) + ns.append(float(np.sum(e[i]))) sigma = [max(1, s) for s in sigma] return (mu, sigma, ns) def em_fit( - x, - weights, - n=1, - tol=1.0e-200, - iter_max=500, - weight_thresh=0, - initial_width=5000, - remove_x=None, -): + x: np.ndarray, + weights: np.ndarray, + n: int = 1, + tol: float = 1.0e-200, + iter_max: int = 500, + weight_thresh: float = 0, + initial_width: float = 5000, + remove_x: float | None = None, +) -> tuple[list[float], list[float], list[float]]: """Perform the expectation maximization fit. Parameters ---------- - x : array of float + x The quantity to run EM on (e.g. the time if EM should find time flares). - weights : array of float + weights The weights for each x value (e.g. the signal over background ratio). - n : int + n How many Gaussians flares we are looking for. - tol : float + tol The stopping criteria for the expectation maximization. This is the difference in the normalized likelihood over the last 20 iterations. - iter_max : int + iter_max The maximum number of iterations, even if stopping criteria tolerance (``tol``) is not yet reached. - weight_thresh : float + weight_thresh Set a minimum threshold for event weights. Events with smaller weights will be removed. - initial_width : float + initial_width The starting width for the gaussian flare in days. - remove_x : float | None + remove_x Specific x of event that should be removed. Returns ------- - mu : list of float + mu The list of size ``n`` with the determined mean values. - sigma : list of float + sigma The list of size ``n`` with the standard deviation values. - ns : list of float + ns The list of size ``n`` with the normalization factor values. """ if weight_thresh > 0: @@ -171,4 +171,4 @@ def em_fit( (mu, sigma, ns) = em_maximization_step(e=e, t=x) - return (mu, sigma, ns) + return (list(map(float, mu)), list(map(float, sigma)), list(map(float, ns))) diff --git a/skyllh/core/flux_model.py b/skyllh/core/flux_model.py index 4012c49d01..8f1137bcbb 100644 --- a/skyllh/core/flux_model.py +++ b/skyllh/core/flux_model.py @@ -10,6 +10,7 @@ class for the most generic flux model is `FluxModel`, which is an abstract base """ import abc +from collections.abc import Callable import numpy as np import scipy.special @@ -42,6 +43,10 @@ class for the most generic flux model is `FluxModel`, which is an abstract base IsPointlike, ) +# numpy deprecated the np.trapz function in favor of np.trapezoid, but to maintain compatibility with older numpy +# versions, we define _trapezoid as np.trapz if np.trapezoid is not available. +_trapezoid = np.trapezoid if hasattr(np, 'trapezoid') else np.trapz # pyright: ignore[reportAttributeAccessIssue] + class FluxProfile(MathFunction, HasConfig, metaclass=abc.ABCMeta): """The abstract base class for a flux profile math function.""" @@ -57,12 +62,12 @@ def __init__( class SpatialFluxProfile(FluxProfile, metaclass=abc.ABCMeta): """The abstract base class for a spatial flux profile function.""" - def __init__(self, angle_unit=None, **kwargs): + def __init__(self, angle_unit: units.UnitBase | None = None, **kwargs): """Creates a new SpatialFluxProfile instance. Parameters ---------- - angle_unit : instance of astropy.units.UnitBase | None + angle_unit The used unit for angles. If set to ``Ǹone``, the configured default angle unit for fluxes is used. @@ -87,24 +92,26 @@ def angle_unit(self, unit): self._angle_unit = unit @abc.abstractmethod - def __call__(self, ra, dec, unit=None): + def __call__( + self, ra: float | np.ndarray, dec: float | np.ndarray, unit: units.UnitBase | None = None + ) -> np.ndarray: """This method is supposed to return the spatial profile value for the given celestrial coordinates. Parameters ---------- - ra : float | 1d numpy ndarray of float + ra The right-ascention coordinate. - dec : float | 1d numpy ndarray of float + dec The declination coordinate. - unit : instance of astropy.units.UnitBase | None + unit The unit of the given celestrial angles. If ``None``, the set angle unit of this SpatialFluxProfile is assumed. Returns ------- - values : 1D numpy ndarray + values The spatial profile values. """ @@ -116,14 +123,14 @@ class UnitySpatialFluxProfile(SpatialFluxProfile): def __init__( self, - angle_unit=None, + angle_unit: units.UnitBase | None = None, **kwargs, ): """Creates a new UnitySpatialFluxProfile instance. Parameters ---------- - angle_unit : instance of astropy.units.UnitBase | None + angle_unit The used unit for angles. If set to ``Ǹone``, the configured default angle unit for fluxes is used. @@ -137,22 +144,24 @@ def math_function_str(self): """ return '1' - def __call__(self, ra, dec, unit=None): + def __call__( + self, ra: float | np.ndarray, dec: float | np.ndarray, unit: units.UnitBase | None = None + ) -> np.ndarray: """Returns 1 as numpy ndarray in same shape as ra and dec. Parameters ---------- - ra : float | 1d numpy ndarray of float + ra The right-ascention coordinate. - dec : float | 1d numpy ndarray of float + dec The declination coordinate. - unit : instance of astropy.units.UnitBase | None + unit The unit of the given celestrial angles. By the definition of this class this argument is ignored. Returns ------- - values : 1D numpy ndarray + values 1 in same shape as ra and dec. """ (ra, dec) = np.atleast_1d(ra, dec) @@ -169,9 +178,9 @@ class PointSpatialFluxProfile(SpatialFluxProfile): def __init__( self, - ra, - dec, - angle_unit=None, + ra: float | None, + dec: float | None, + angle_unit: units.UnitBase | None = None, **kwargs, ): """Creates a new spatial flux profile for a point at equatorial @@ -179,15 +188,15 @@ def __init__( Parameters ---------- - ra : float | None + ra The right-ascention of the point. In case it is None, the evaluation of this spatial flux profile will return zero, unless evaluated for ra=None. - dec : float | None + dec The declination of the point. In case it is None, the evaluation of this spatial flux profile will return zero, unless evaluated for dec=None. - angle_unit : instance of astropy.units.UnitBase | None + angle_unit The used unit for angles. If set to ``Ǹone``, the configured default angle unit for fluxes is used. @@ -237,26 +246,28 @@ def math_function_str(self): return s - def __call__(self, ra, dec, unit=None): + def __call__( + self, ra: float | np.ndarray, dec: float | np.ndarray, unit: units.UnitBase | None = None + ) -> np.ndarray: """Returns a numpy ndarray in same shape as ra and dec with 1 if `ra` equals `self.ra` and `dec` equals `self.dec`, and 0 otherwise. Parameters ---------- - ra : float | 1d numpy ndarray of float + ra The right-ascention coordinate at which to evaluate the spatial flux profile. The unit must be the internally used angle unit. - dec : float | 1d numpy ndarray of float + dec The declination coordinate at which to evaluate the spatial flux profile. The unit must be the internally used angle unit. - unit : instance of astropy.units.UnitBase | None + unit The unit of the given celestrial angles. If set to ``None``, the set angle unit of this SpatialFluxProfile instance is assumed. Returns ------- - value : 1D numpy ndarray of int8 + value A numpy ndarray in same shape as ra and dec with 1 if `ra` equals `self.ra` and `dec` equals `self.dec`, and 0 otherwise. """ @@ -277,13 +288,13 @@ def __call__(self, ra, dec, unit=None): class EnergyFluxProfile(FluxProfile, metaclass=abc.ABCMeta): """The abstract base class for an energy flux profile function.""" - def __init__(self, energy_unit=None, **kwargs): + def __init__(self, energy_unit: units.UnitBase | None = None, **kwargs): """Creates a new energy flux profile with a given energy unit to be used for flux calculation. Parameters ---------- - energy_unit : instance of astropy.units.UnitBase | None + energy_unit The used unit for energy. If set to ``None``, the configured default energy unit for fluxes is used. @@ -307,31 +318,31 @@ def energy_unit(self, unit): self._energy_unit = unit @abc.abstractmethod - def __call__(self, E, unit=None): + def __call__(self, E: float | np.ndarray, unit: units.UnitBase | None = None) -> np.ndarray: """This method is supposed to return the energy profile value for the given energy value. Parameters ---------- - E : float | 1d numpy ndarray of float + E The energy value for which to retrieve the energy profile value. - unit : instance of astropy.units.UnitBase | None + unit The unit of the given energy. If set to ``None``, the set energy unit of this EnergyFluxProfile is assumed. Returns ------- - values : 1D numpy ndarray of float + values The energy profile values for the given energies. """ def get_integral( self, - E1, - E2, - unit=None, - ): + E1: float | np.ndarray, + E2: float | np.ndarray, + unit: units.UnitBase | None = None, + ) -> np.ndarray: """This is the default implementation for calculating the integral value of this energy flux profile in the range ``[E1, E2]``. @@ -344,18 +355,18 @@ def get_integral( Parameters ---------- - E1 : float | 1d numpy ndarray of float + E1 The lower energy bound of the integration. - E2 : float | 1d numpy ndarray of float + E2 The upper energy bound of the integration. - unit : instance of astropy.units.UnitBase | None + unit The unit of the given energies. If set to ``None``, the set energy unit of this EnergyFluxProfile instance is assumed. Returns ------- - integral : instance of ndarray + integral The (n,)-shaped numpy ndarray holding the integral values of the given integral ranges. """ @@ -378,12 +389,12 @@ def get_integral( class UnityEnergyFluxProfile(EnergyFluxProfile): """Energy flux profile for the constant function 1.""" - def __init__(self, energy_unit=None, **kwargs): + def __init__(self, energy_unit: units.UnitBase | None = None, **kwargs): """Creates a new UnityEnergyFluxProfile instance. Parameters ---------- - energy_unit : instance of astropy.units.UnitBase | None + energy_unit The used unit for energy. If set to ``None``, the configured default energy unit for fluxes is used. @@ -397,20 +408,20 @@ def math_function_str(self): """ return '1' - def __call__(self, E, unit=None): + def __call__(self, E: float | np.ndarray, unit: units.UnitBase | None = None) -> np.ndarray: """Returns 1 as numpy ndarray in some shape as E. Parameters ---------- - E : float | 1D numpy ndarray of float + E The energy value for which to retrieve the energy profile value. - unit : instance of astropy.units.UnitBase | None + unit The unit of the given energies. By definition of this specific class, this argument is ignored. Returns ------- - values : 1D numpy ndarray of int8 + values 1 in same shape as E. """ E = np.atleast_1d(E) @@ -419,24 +430,26 @@ def __call__(self, E, unit=None): return values - def get_integral(self, E1, E2, unit=None): + def get_integral( + self, E1: float | np.ndarray, E2: float | np.ndarray, unit: units.UnitBase | None = None + ) -> np.ndarray: """Computes the integral of this energy flux profile in the range [``E1``, ``E2``], which by definition is ``E2 - E1``. Parameters ---------- - E1 : float | 1d numpy ndarray of float + E1 The lower energy bound of the integration. - E2 : float | 1d numpy ndarray of float + E2 The upper energy bound of the integration. - unit : instance of astropy.units.UnitBase | None + unit The unit of the given energies. If set to ``None``, the set energy unit of this EnergyFluxProfile instance is assumed. Returns ------- - integral : 1d ndarray of float + integral The integral values of the given integral ranges. """ E1 = np.atleast_1d(E1) @@ -464,17 +477,17 @@ class PowerLawEnergyFluxProfile( """ - def __init__(self, E0, gamma, energy_unit=None, **kwargs): + def __init__(self, E0: float, gamma: float, energy_unit: units.UnitBase | None = None, **kwargs): """Creates a new power law flux profile with the reference energy ``E0`` and spectral index ``gamma``. Parameters ---------- - E0 : castable to float + E0 The reference energy. - gamma : castable to float + gamma The spectral index. - energy_unit : instance of astropy.units.UnitBase | None + energy_unit The used unit for energy. If set to ``None``, the configured default energy unit for fluxes is used. @@ -522,22 +535,22 @@ def math_function_str(self): return s - def __call__(self, E, unit=None): + def __call__(self, E: float | np.ndarray, unit: units.UnitBase | None = None) -> np.ndarray: """Returns the power law values for the given energies as numpy ndarray in same shape as E. Parameters ---------- - E : float | 1D numpy ndarray of float + E The energy value for which to retrieve the energy profile value. - unit : instance of astropy.units.UnitBase | None + unit The unit of the given energies. If set to ``None``, the set energy unit of this EnergyFluxProfile instance is assumed. Returns ------- - values : 1D numpy ndarray of float + values The energy profile values for the given energies. """ E = np.atleast_1d(E) @@ -549,24 +562,26 @@ def __call__(self, E, unit=None): return value - def get_integral(self, E1, E2, unit=None): + def get_integral( + self, E1: float | np.ndarray, E2: float | np.ndarray, unit: units.UnitBase | None = None + ) -> np.ndarray: """Computes the integral value of this power-law energy flux profile in the range ``[E1, E2]``. Parameters ---------- - E1 : float | 1d numpy ndarray of float + E1 The lower energy bound of the integration. - E2 : float | 1d numpy ndarray of float + E2 The upper energy bound of the integration. - unit : instance of astropy.units.UnitBase | None + unit The unit of the given energies. If set to ``None``, the set energy unit of this EnergyFluxProfile instance is assumed. Returns ------- - integral : 1d ndarray of float + integral The integral values of the given integral ranges. """ E1 = np.atleast_1d(E1) @@ -602,10 +617,10 @@ class CutoffPowerLawEnergyFluxProfile( def __init__( self, - E0, - gamma, - Ecut, - energy_unit=None, + E0: float, + gamma: float, + Ecut: float, + energy_unit: units.UnitBase | None = None, **kwargs, ): """Creates a new cut-off power law flux profile with the reference @@ -613,13 +628,13 @@ def __init__( Parameters ---------- - E0 : castable to float + E0 The reference energy. - gamma : castable to float + gamma The spectral index. - Ecut : castable to float + Ecut The cut-off energy. - energy_unit : instance of astropy.units.UnitBase | None + energy_unit The used unit for energy. If set to ``None``, the configured default energy unit for fluxes is used. @@ -649,24 +664,24 @@ def math_function_str(self): def __call__( self, - E, - unit=None, - ): + E: float | np.ndarray, + unit: units.UnitBase | None = None, + ) -> np.ndarray: """Returns the cut-off power law values for the given energies as numpy ndarray in the same shape as E. Parameters ---------- - E : float | instance of numpy ndarray + E The energy value(s) for which to retrieve the energy profile value. - unit : instance of astropy.units.UnitBase | None + unit The unit of the given energies. If set to ``None``, the set energy unit of this EnergyFluxProfile instance is assumed. Returns ------- - values : instance of numpy ndarray + values The energy profile values for the given energies. """ E = np.atleast_1d(E) @@ -681,10 +696,10 @@ def __call__( def get_integral( self, - E1, - E2, - unit=None, - ): + E1: float | np.ndarray, + E2: float | np.ndarray, + unit: units.UnitBase | None = None, + ) -> np.ndarray: """This falls back to the default implementation for the CutoffPowerLawEnergyFluxProfile to avoid using the inhereted function .. note:: @@ -696,18 +711,18 @@ def get_integral( Parameters ---------- - E1 : float | 1d numpy ndarray of float + E1 The lower energy bound of the integration. - E2 : float | 1d numpy ndarray of float + E2 The upper energy bound of the integration. - unit : instance of astropy.units.UnitBase | None + unit The unit of the given energies. If set to ``None``, the set energy unit of this EnergyFluxProfile instance is assumed. Returns ------- - integral : instance of ndarray + integral The (n,)-shaped numpy ndarray holding the integral values of the given integral ranges. """ @@ -724,7 +739,7 @@ def get_integral( for i, (E1_i, E2_i) in enumerate(zip(E1, E2, strict=True)): # integration for log(energy) for hopefully better numerical stability tmp_e = np.linspace(np.log10(E1_i), np.log10(E2_i), 5000) - tmp_int = np.trapz(np.log(10) * self(10**tmp_e) * 10**tmp_e, tmp_e) + tmp_int = _trapezoid(np.log(10) * self(10**tmp_e) * 10**tmp_e, tmp_e) # make sure it is always positive (probably not an issue any more with np.trapz. # used to be an issue using the spline integrate self.function.integrate) @@ -747,22 +762,22 @@ class LogParabolaPowerLawEnergyFluxProfile( def __init__( self, - E0, - alpha, - beta, - energy_unit=None, + E0: float, + alpha: float, + beta: float, + energy_unit: units.UnitBase | None = None, **kwargs, ): """ Parameters ---------- - E0 : castable to float + E0 The reference energy. - alpha : float + alpha The alpha parameter of the log-parabola spectral index. - beta : float + beta The beta parameter of the log-parabola spectral index. - energy_unit : instance of astropy.units.UnitBase | None + energy_unit The used unit for energy. If set to ``None``, the configured default energy unit for fluxes is used. @@ -804,24 +819,24 @@ def math_function_str(self): def __call__( self, - E, - unit=None, - ): + E: float | np.ndarray, + unit: units.UnitBase | None = None, + ) -> np.ndarray: """Returns the log-parabola power-law values for the given energies as numpy ndarray in the same shape as E. Parameters ---------- - E : float | instance of numpy ndarray + E The energy value(s) for which to retrieve the energy profile value. - unit : instance of astropy.units.UnitBase | None + unit The unit of the given energies. If set to ``None``, the set energy unit of this EnergyFluxProfile instance is assumed. Returns ------- - values : instance of numpy ndarray + values The energy profile values for the given energies. """ E = np.atleast_1d(E) @@ -845,23 +860,23 @@ class PhotosplineEnergyFluxProfile( def __init__( self, splinetable, - crit_log10_energy_lower, - crit_log10_energy_upper, - energy_unit=None, + crit_log10_energy_lower: float, + crit_log10_energy_upper: float, + energy_unit: units.UnitBase | None = None, **kwargs, ): """Creates a new instance of PhotosplineEnergyFluxProfile. Parameters ---------- - splinetable : instance of photospline.SplineTable + splinetable The instance of photospline.SplineTable representing the energy flux profile as a spline. - crit_log10_energy_lower : float + crit_log10_energy_lower The lower edge of the spline's supported energy range in log10(E). - crit_log10_energy_upper : float + crit_log10_energy_upper The upper edge of the spline's supported energy range in log10(E). - energy_unit : instance of astropy.units.UnitBase | None + energy_unit The used unit for energy. If set to ``None``, the configured default energy unit for fluxes is used. @@ -916,13 +931,13 @@ def crit_log10_energy_upper(self, v): class FunctionEnergyFluxProfile(EnergyFluxProfile): r"""Energy flux profile for a callable function with energy as argument.""" - def __init__(self, function, energy_unit=None, **kwargs): + def __init__(self, function: Callable, energy_unit: units.UnitBase | None = None, **kwargs): """Creates a new flux profile following a given function. Parameters ---------- - function : callable function, takes energy as argument - energy_unit : instance of astropy.units.UnitBase | None + function + energy_unit The used unit for energy. If set to ``None``, the configured default energy unit for fluxes is used. @@ -931,22 +946,22 @@ def __init__(self, function, energy_unit=None, **kwargs): self.function = function - def __call__(self, E, unit=None): + def __call__(self, E: float | np.ndarray, unit: units.UnitBase | None = None) -> np.ndarray: """Returns the function values for the given energies as numpy ndarray in same shape as E. Parameters ---------- - E : float | 1D numpy ndarray of float + E The energy value for which to retrieve the energy profile value. - unit : instance of astropy.units.UnitBase | None + unit The unit of the given energies. If set to ``None``, the set energy unit of this EnergyFluxProfile instance is assumed. Returns ------- - values : 1D numpy ndarray of float + values The energy profile values for the given energies. """ E = np.atleast_1d(E) @@ -959,7 +974,7 @@ def __call__(self, E, unit=None): return value @property - def math_function_str(self): + def math_function_str(self) -> str: """(read-only) The string representation of the mathematical function of this energy flux profile. """ @@ -972,15 +987,17 @@ class EpeakFunctionEnergyProfile(FunctionEnergyFluxProfile): will be optimized. """ - def __init__(self, function, e_peak_orig, e_peak_offset, energy_unit=None, **kwargs): + def __init__( + self, function: Callable, e_peak_orig, e_peak_offset, energy_unit: units.UnitBase | None = None, **kwargs + ): """Creates a new flux profile with the peak energy ``E0``. Parameters ---------- - function : callable function, takes energy as argument - e_peak_orig : log10(energy) for which the original array reaches its peak value - e_peak_offset : log10(energy) to which the peak flux should be shifted - energy_unit : instance of astropy.units.UnitBase | None + function + e_peak_orig + e_peak_offset + energy_unit The used unit for energy. If set to ``None``, the configured default energy unit for fluxes is used. @@ -1016,22 +1033,22 @@ def e_peak_orig(self, e): e = float_cast(e, 'Property e must be castable to type float!') self._e_peak_orig = e - def __call__(self, E, unit=None): + def __call__(self, E: float | np.ndarray, unit: units.UnitBase | None = None) -> np.ndarray: """Returns the function values for the given energies as numpy ndarray in same shape as E. Parameters ---------- - E : float | 1D numpy ndarray of float + E The energy value for which to retrieve the energy profile value. - unit : instance of astropy.units.UnitBase | None + unit The unit of the given energies. If set to ``None``, the set energy unit of this EnergyFluxProfile instance is assumed. Returns ------- - values : 1D numpy ndarray of float + values The energy profile values for the given energies. """ E = np.atleast_1d(E) @@ -1052,10 +1069,10 @@ def __call__(self, E, unit=None): def get_integral( self, - E1, - E2, - unit=None, - ): + E1: float | np.ndarray, + E2: float | np.ndarray, + unit: units.UnitBase | None = None, + ) -> np.ndarray: """This is the default implementation for calculating the integral value of this energy flux profile in the range ``[E1, E2]``. @@ -1068,18 +1085,18 @@ def get_integral( Parameters ---------- - E1 : float | 1d numpy ndarray of float + E1 The lower energy bound of the integration. - E2 : float | 1d numpy ndarray of float + E2 The upper energy bound of the integration. - unit : instance of astropy.units.UnitBase | None + unit The unit of the given energies. If set to ``None``, the set energy unit of this EnergyFluxProfile instance is assumed. Returns ------- - integral : instance of ndarray + integral The (n,)-shaped numpy ndarray holding the integral values of the given integral ranges. """ @@ -1097,7 +1114,7 @@ def get_integral( # integration for log(energy) for hopefully better numerical stability tmp_e = np.linspace(np.log10(E1_i), np.log10(E2_i)) - tmp_int = np.trapz(np.log(10) * self(10**tmp_e) * 10**tmp_e, tmp_e) + tmp_int = _trapezoid(np.log(10) * self(10**tmp_e) * 10**tmp_e, tmp_e) # make sure it is always positive (probably not an issue any more with np.trapz. # used to be an issue using the spline integrate self.function.integrate) @@ -1119,20 +1136,22 @@ class TimeFluxProfile( ): """The abstract base class for a time flux profile function.""" - def __init__(self, t_start=-np.inf, t_stop=np.inf, time_unit=None, **kwargs): + def __init__( + self, t_start: float = -np.inf, t_stop: float = np.inf, time_unit: units.UnitBase | None = None, **kwargs + ): """Creates a new time flux profile instance. Parameters ---------- - t_start : float + t_start The start time of the time profile. If set to -inf, it means, that the profile starts at the beginning of the entire time-span of the dataset. - t_stop : float + t_stop The stop time of the time profile. If set to +inf, it means, that the profile ends at the end of the entire time-span of the dataset. - time_unit : instance of astropy.units.UnitBase | None + time_unit The used unit for time. If set to ``None``, the configured default time unit for fluxes is used. @@ -1192,58 +1211,58 @@ def time_unit(self, unit): ) self._time_unit = unit - def get_total_integral(self): + def get_total_integral(self) -> 'float | np.ndarray': """Calculates the total integral of the time profile from t_start to t_stop. Returns ------- - integral : float + integral The integral value of the entire time profile. The value is in the set time unit of this TimeFluxProfile instance. """ - integral = self.get_integral(self._t_start, self._t_stop).squeeze() + integral = np.asarray(self.get_integral(self._t_start, self._t_stop)).squeeze() return integral @abc.abstractmethod def __call__( self, - t, - unit=None, - ): + t: float | np.ndarray, + unit: units.UnitBase | None = None, + ) -> np.ndarray: """This method is supposed to return the time profile value for the given times. Parameters ---------- - t : float | 1D numpy ndarray of float + t The time(s) for which to get the time flux profile values. - unit : instance of astropy.units.UnitBase | None + unit The unit of the given times. If set to ``None``, the set time unit of this TimeFluxProfile instance is assumed. Returns ------- - values : 1D numpy ndarray of float + values The time profile values. """ @abc.abstractmethod def move( self, - dt, - unit=None, + dt: float, + unit: units.UnitBase | None = None, ): """Abstract method to move the time profile by the given amount of time. Parameters ---------- - dt : float + dt The time difference of how far to move the time profile in time. This can be a positive or negative time shift value. - unit : instance of astropy.units.UnitBase | None + unit The unit of the given time difference. If set to ``Ǹone``, the set time unit of this TimeFluxProfile instance is assumed. @@ -1252,27 +1271,27 @@ def move( @abc.abstractmethod def get_integral( self, - t1, - t2, - unit=None, - ): + t1: float | np.ndarray, + t2: float | np.ndarray, + unit: units.UnitBase | None = None, + ) -> 'float | np.ndarray': """This method is supposed to calculate the integral of the time profile from time ``t1`` to time ``t2``. Parameters ---------- - t1 : float | array of float + t1 The start time of the integration. - t2 : float | array of float + t2 The end time of the integration. - unit : instance of astropy.units.UnitBase | None + unit The unit of the given times. If set to ``Ǹone``, the set time unit of this TimeFluxProfile instance is assumed. Returns ------- - integral : array of float + integral The integral value(s) of the time profile. The values are in the set time unit of this TimeFluxProfile instance. """ @@ -1285,14 +1304,14 @@ class UnityTimeFluxProfile( def __init__( self, - time_unit=None, + time_unit: units.UnitBase | None = None, **kwargs, ): """Creates a new instance of UnityTimeFluxProfile. Parameters ---------- - time_unit : instance of astropy.units.UnitBase | None + time_unit The used unit for time. If set to ``None``, the configured default time unit for fluxes is used. @@ -1301,26 +1320,29 @@ def __init__( @property def math_function_str(self): + """The string representation of the mathematical function of this time + flux profile. + """ return '1' def __call__( self, - t, - unit=None, - ): + t: float | np.ndarray, + unit: units.UnitBase | None = None, + ) -> np.ndarray: """Returns 1 as numpy ndarray in same shape as t. Parameters ---------- - t : float | 1D numpy ndarray of float + t The time(s) for which to get the time flux profile values. - unit : instance of astropy.units.UnitBase | None + unit The unit of the given times. By definition of this specific class, this argument is ignored. Returns ------- - values : 1D numpy ndarray of int8 + values 1 in same shape as ``t``. """ t = np.atleast_1d(t) @@ -1331,8 +1353,8 @@ def __call__( def move( self, - dt, - unit=None, + dt: float, + unit: units.UnitBase | None = None, ): """Moves the time profile by the given amount of time. By definition this method does nothing, because the profile is 1 over the entire @@ -1340,10 +1362,10 @@ def move( Parameters ---------- - dt : float + dt The time difference of how far to move the time profile in time. This can be a positive or negative time shift value. - unit : instance of astropy.units.UnitBase | None + unit The unit of the given time difference. If set to ``None``, the set time unit of this TimeFluxProfile instance is assumed. @@ -1351,31 +1373,31 @@ def move( def get_integral( self, - t1, - t2, - unit=None, - ): + t1: float | np.ndarray, + t2: float | np.ndarray, + unit: units.UnitBase | None = None, + ) -> 'float | np.ndarray': """Calculates the integral of the time profile from time t1 to time t2. Parameters ---------- - t1 : float | array of float + t1 The start time of the integration. - t2 : float | array of float + t2 The end time of the integration. - unit : instance of astropy.units.UnitBase | None + unit The unit of the given times. If set to ``None``, the set time unit of this TimeFluxProfile instance is assumed. Returns ------- - integral : array of float + integral The integral value(s) of the time profile. The values are in the set time unit of this TimeFluxProfile instance. """ if (unit is not None) and (unit != self._time_unit): - time_unit_conv_factor = unit.to(self._time_unit) + time_unit_conv_factor = float(unit.to(self._time_unit)) # pyright: ignore[reportArgumentType] t1 = t1 * time_unit_conv_factor t2 = t2 * time_unit_conv_factor @@ -1401,21 +1423,21 @@ class BoxTimeFluxProfile( @classmethod def from_start_and_stop_time( cls, - start, - stop, - time_unit=None, + start: float, + stop: float, + time_unit: units.UnitBase | None = None, **kwargs, - ): + ) -> 'BoxTimeFluxProfile': """Constructs a BoxTimeFluxProfile instance from the given start and stop time. Parameters ---------- - start : float + start The start time of the box profile. - stop : float + stop The stop time of the box profile. - time_unit : instance of astropy.units.UnitBase | None + time_unit The used unit for time. If set to ``None``, the configured default time unit for fluxes is used. @@ -1425,7 +1447,7 @@ def from_start_and_stop_time( Returns ------- - profile : instance of BoxTimeFluxProfile + profile The newly created instance of BoxTimeFluxProfile. """ t0 = 0.5 * (start + stop) @@ -1437,20 +1459,20 @@ def from_start_and_stop_time( def __init__( self, - t0, - tw, - time_unit=None, + t0: float, + tw: float, + time_unit: units.UnitBase | None = None, **kwargs, ): """Creates a new box-shaped time profile instance. Parameters ---------- - t0 : float + t0 The mid time of the box profile. - tw : float + tw The width of the box profile. - time_unit : instance of astropy.units.UnitBase | None + time_unit The used unit for time. If set to ``None``, the configured default time unit for fluxes is used. @@ -1504,24 +1526,24 @@ def math_function_str(self): def __call__( self, - t, - unit=None, - ): + t: float | np.ndarray, + unit: units.UnitBase | None = None, + ) -> np.ndarray: """Returns 1 for all t within the interval [t0-tw/2; t0+tw/2], and 0 otherwise. Parameters ---------- - t : float | 1D numpy ndarray of float + t The time(s) for which to get the time flux profile values. - unit : instance of astropy.units.UnitBase | None + unit The unit of the given times. If set to ``None``, the set time unit of this TimeFluxProfile instance is assumed. Returns ------- - values : 1D numpy ndarray of int8 + values The value(s) of the time flux profile for the given time(s). """ t = np.atleast_1d(t) @@ -1537,25 +1559,25 @@ def __call__( def cdf( self, - t, - unit=None, - ): + t: float | np.ndarray, + unit: units.UnitBase | None = None, + ) -> np.ndarray: """Calculates the cumulative distribution function value for the given time values ``t``. Parameters ---------- - t : float | instance of numpy ndarray + t The (N_times,)-shaped numpy ndarray holding the time values for which to calculate the CDF values. - unit : instance of astropy.units.UnitBase | None + unit The unit of the given times. If set to ``None``, the set time unit of this TimeFluxProfile is assumed. Returns ------- - values : instance of numpy ndarray + values The (N_times,)-shaped numpy ndarray holding the cumulative distribution function values for each time ``t``. """ @@ -1579,50 +1601,50 @@ def cdf( def move( self, - dt, - unit=None, + dt: float, + unit: units.UnitBase | None = None, ): """Moves the box-shaped time profile by the time difference dt. Parameters ---------- - dt : float + dt The time difference of how far to move the time profile in time. This can be a positive or negative time shift value. - unit : instance of astropy.units.UnitBase | None + unit The unit of ``dt``. If set to ``None``, the set time unit of this TimeFluxProfile instance is assumed. """ if (unit is not None) and (unit != self._time_unit): - dt = dt * unit.to(self._time_unit) + dt = dt * unit.to(self._time_unit) # pyright: ignore[reportAssignmentType] self._t_start += dt self._t_stop += dt def get_integral( self, - t1, - t2, - unit=None, - ): + t1: float | np.ndarray, + t2: float | np.ndarray, + unit: units.UnitBase | None = None, + ) -> np.ndarray: """Calculates the integral of the box-shaped time flux profile from time t1 to time t2. Parameters ---------- - t1 : float | array of float + t1 The start time(s) of the integration. - t2 : float | array of float + t2 The end time(s) of the integration. - unit : instance of astropy.units.UnitBase | None + unit The unit of the given times. If set to ``None``, the set time unit of this TimeFluxProfile instance is assumed. Returns ------- - integral : array of float + integral The integral value(s). The values are in the set time unit of this TimeFluxProfile instance. """ @@ -1659,19 +1681,21 @@ class GaussianTimeFluxProfile( The one-sigma width of the gaussian profile. """ - def __init__(self, t0, sigma_t, tol=1e-12, time_unit=None, **kwargs): + def __init__( + self, t0: float, sigma_t: float, tol: float = 1e-12, time_unit: units.UnitBase | None = None, **kwargs + ): """Creates a new gaussian-shaped time flux profile instance. Parameters ---------- - t0 : float + t0 The mid time of the gaussian profile. - sigma_t : float + sigma_t The one-sigma width of the gaussian profile. - tol : float + tol The tolerance of the gaussian value. This defines the start and end time of the gaussian profile. - time_unit : instance of astropy.units.UnitBase | None + time_unit The used unit for time. If set to ``None``, the configured default time unit for fluxes is used. @@ -1693,6 +1717,9 @@ def __init__(self, t0, sigma_t, tol=1e-12, time_unit=None, **kwargs): @property def math_function_str(self): + """The string representation of the mathematical function of this time + flux profile. + """ return 'exp(-(t-t0)^2/(2 sigma_t^2))' @property @@ -1701,7 +1728,7 @@ def t0(self): The unit of the value is the set time unit of this TimeFluxProfile instance. """ - return 0.5 * (self._t_start + self._t_stop) + return 0.5 * (self.t_start + self.t_stop) @t0.setter def t0(self, t): @@ -1725,23 +1752,23 @@ def sigma_t(self, sigma): def __call__( self, - t, - unit=None, - ): + t: float | np.ndarray, + unit: units.UnitBase | None = None, + ) -> np.ndarray: """Returns the gaussian profile value for the given time ``t``. Parameters ---------- - t : float | 1D numpy ndarray of float + t The time(s) for which to get the time flux profile values. - unit : instance of astropy.units.UnitBase | None + unit The unit of the given times. If set to ``None``, the set time unit of this TimeFluxProfile is assumed. Returns ------- - values : 1D numpy ndarray of float + values The value(s) of the time flux profile for the given time(s). """ t = np.atleast_1d(t) @@ -1754,7 +1781,7 @@ def __call__( s = self._sigma_t twossq = 2 * s * s - t0 = 0.5 * (self._t_stop + self._t_start) + t0 = 0.5 * (self.t_stop + self.t_start) dt = t[m] - t0 values = np.zeros_like(t) @@ -1764,25 +1791,25 @@ def __call__( def cdf( self, - t, - unit=None, - ): + t: float | np.ndarray, + unit: units.UnitBase | None = None, + ) -> np.ndarray: """Calculates the cumulative distribution function values for the given time values ``t``. Parameters ---------- - t : float | instance of numpy ndarray + t The (N_times,)-shaped numpy ndarray holding the time values for which to calculate the CDF values. - unit : instance of astropy.units.UnitBase | None + unit The unit of the given times. If set to ``None``, the set time unit of this TimeFluxProfile is assumed. Returns ------- - values : instance of numpy ndarray + values The (N_times,)-shaped numpy ndarray holding the cumulative distribution function values for each time ``t``. """ @@ -1806,59 +1833,59 @@ def cdf( def move( self, - dt, - unit=None, + dt: float, + unit: units.UnitBase | None = None, ): """Moves the gaussian time profile by the given amount of time. Parameters ---------- - dt : float + dt The time difference of how far to move the time profile in time. This can be a positive or negative time shift value. - unit : instance of astropy.units.UnitBase | None + unit The unit of the given time difference. If set to ``None``, the set time unit of this TimeFluxProfile is assumed. """ if (unit is not None) and (unit != self._time_unit): - dt = dt * unit.to(self._time_unit) + dt = dt * unit.to(self._time_unit) # pyright: ignore[reportAssignmentType] self._t_start += dt self._t_stop += dt def get_integral( self, - t1, - t2, - unit=None, - ): + t1: float | np.ndarray, + t2: float | np.ndarray, + unit: units.UnitBase | None = None, + ) -> np.ndarray: """Calculates the integral of the gaussian time profile from time ``t1`` to time ``t2``. Parameters ---------- - t1 : float | array of float + t1 The start time(s) of the integration. - t2 : float | array of float + t2 The end time(s) of the integration. - unit : instance of astropy.units.UnitBase | None + unit The unit of the given times. If set to ``None``, the set time unit of this TimeFluxProfile instance is assumed. Returns ------- - integral : array of float + integral The integral value(s). The values are in the set time unit of this TimeFluxProfile instance. """ if (unit is not None) and (unit != self._time_unit): - time_unit_conv_factor = unit.to(self._time_unit) + time_unit_conv_factor = float(unit.to(self._time_unit)) # pyright: ignore[reportArgumentType] t1 = t1 * time_unit_conv_factor t2 = t2 * time_unit_conv_factor - t0 = 0.5 * (self._t_stop + self._t_start) + t0 = 0.5 * (self.t_stop + self.t_start) sigma_t = self._sigma_t c1 = np.sqrt(np.pi / 2) * sigma_t @@ -1890,40 +1917,47 @@ class FluxModel( """ @staticmethod - def get_default_units(cfg): + def get_default_units(cfg: Config) -> dict: """Returns the configured default units for flux models. Parameters ---------- - cfg : instance of Config + cfg The instance of Config holding the local configuration. Returns ------- - units_dict : dict + units_dict The dictionary holding the configured default units used for flux models. """ return cfg['units']['defaults']['fluxes'] - def __init__(self, angle_unit=None, energy_unit=None, length_unit=None, time_unit=None, **kwargs): + def __init__( + self, + angle_unit: units.UnitBase | None = None, + energy_unit: units.UnitBase | None = None, + length_unit: units.UnitBase | None = None, + time_unit: units.UnitBase | None = None, + **kwargs, + ): """Creates a new FluxModel instance and defines the user-defined units. Parameters ---------- - angle_unit : instance of astropy.units.UnitBase | None + angle_unit The used unit for angles. If set to ``None``, the configured default angle unit for fluxes is used. - energy_unit : instance of astropy.units.UnitBase | None + energy_unit The used unit for energy. If set to ``None``, the configured default energy unit for fluxes is used. - length_unit : instance of astropy.units.UnitBase | None + length_unit The used unit for length. If set to ``None``, the configured default length unit for fluxes is used. - time_unit : instance of astropy.units.UnitBase | None + time_unit The used unit for time. If set to ``None``, the configured default time unit for fluxes is used. @@ -2031,44 +2065,53 @@ def __str__(self): return f'{self.math_function_str} {self.unit_str}' @abc.abstractmethod - def __call__(self, ra=None, dec=None, E=None, t=None, angle_unit=None, energy_unit=None, time_unit=None): + def __call__( + self, + ra: float | np.ndarray | None = None, + dec: float | np.ndarray | None = None, + E: float | np.ndarray | None = None, + t: float | np.ndarray | None = None, + angle_unit: units.UnitBase | None = None, + energy_unit: units.UnitBase | None = None, + time_unit: units.UnitBase | None = None, + ) -> np.ndarray: """The call operator to retrieve a flux value for a given celestrial position, energy, and observation time. Parameters ---------- - ra : float | (Ncoord,)-shaped 1D numpy ndarray of float + ra The right-ascention coordinate for which to retrieve the flux value. - dec : float | (Ncoord,)-shaped 1D numpy ndarray of float + dec The declination coordinate for which to retrieve the flux value. - E : float | (Nenergy,)-shaped 1D numpy ndarray of float + E The energy for which to retrieve the flux value. - t : float | (Ntime,)-shaped 1D numpy ndarray of float + t The observation time for which to retrieve the flux value. - angle_unit : instance of astropy.units.UnitBase | None + angle_unit The unit of the given angles. If ``None``, the set angle unit of the flux model is assumed. - energy_unit : instance of astropy.units.UnitBase | None + energy_unit The unit of the given energies. If ``None``, the set energy unit of the flux model is assumed. - time_unit : instance of astropy.units.UnitBase | None + time_unit The unit of the given times. If ``None``, the set time unit of the flux model is assumed. Returns ------- - flux : (Ncoord,Nenergy,Ntime)-shaped ndarray of float + flux The flux values are in unit of the set flux model units [energy]^{-1} [angle]^{-2} [length]^{-2} [time]^{-1}. """ - def to_internal_flux_unit(self): + def to_internal_flux_unit(self) -> float: """Calculates the conversion factor to convert the flux unit of this flux model instance to the SkyLLH internally used flux unit. Returns ------- - factor : float + factor The conversion factor. """ self_flux_unit = 1 / (self.angle_unit**2 * self.energy_unit * self.length_unit**2 * self.time_unit) @@ -2081,7 +2124,7 @@ def to_internal_flux_unit(self): * internal_units['time'] ) - factor = (self_flux_unit).to(internal_flux_unit).value + factor = (self_flux_unit).to(internal_flux_unit).value # pyright: ignore[reportAttributeAccessIssue] return factor @@ -2097,14 +2140,14 @@ class NullFluxModel( def __init__( self, *args, - cfg=None, + cfg: Config | None = None, **kwargs, ): """Creates a new instance of NullFluxModel. Parameters ---------- - cfg : instance of Config | None + cfg The instance of Config holding the local configuration. Since this flux model does nothing, this argument is optional. If not provided the default configuration is used. @@ -2114,6 +2157,7 @@ def __init__( super().__init__(*args, cfg=cfg, **kwargs) + @property def math_function_str(self): """Since this is a dummy flux model, calling this method will raise a NotImplementedError. @@ -2151,35 +2195,35 @@ class FactorizedFluxModel( def __init__( self, - Phi0, - spatial_profile, - energy_profile, - time_profile, - length_unit=None, + Phi0: float, + spatial_profile: 'SpatialFluxProfile | None', + energy_profile: 'EnergyFluxProfile | None', + time_profile: 'TimeFluxProfile | None', + length_unit: units.UnitBase | None = None, **kwargs, ): """Creates a new factorized flux model. Parameters ---------- - Phi0 : float + Phi0 The flux normalization constant. - spatial_profile : instance of SpatialFluxProfile | None + spatial_profile The SpatialFluxProfile instance providing the spatial profile function of the flux. If set to None, an instance of UnitySpatialFluxProfile will be used, which represents the constant function 1. - energy_profile : instance of EnergyFluxProfile | None + energy_profile The EnergyFluxProfile instance providing the energy profile function of the flux. If set to None, an instance of UnityEnergyFluxProfile will be used, which represents the constant function 1. - time_profile : instance of TimeFluxProfile | None + time_profile The TimeFluxProfile instance providing the time profile function of the flux. If set to None, an instance of UnityTimeFluxProfile will be used, which represents the constant function 1. - length_unit : instance of astropy.units.UnitBase | None + length_unit The used unit for length. If set to ``None``, the configured default length unit for fluxes is used. @@ -2330,47 +2374,47 @@ def param_names(self): @param_names.setter def param_names(self, names): - super(FactorizedFluxModel, type(self)).param_names.fset(self, names) + super(FactorizedFluxModel, type(self)).param_names.fset(self, names) # pyright: ignore[reportAttributeAccessIssue] def __call__( self, - ra=None, - dec=None, - E=None, - t=None, - angle_unit=None, - energy_unit=None, - time_unit=None, - ): + ra: float | np.ndarray | None = None, + dec: float | np.ndarray | None = None, + E: float | np.ndarray | None = None, + t: float | np.ndarray | None = None, + angle_unit: units.UnitBase | None = None, + energy_unit: units.UnitBase | None = None, + time_unit: units.UnitBase | None = None, + ) -> np.ndarray: """Calculates the flux values for the given celestrial positions, energies, and observation times. Parameters ---------- - ra: float | (Ncoord,)-shaped 1d numpy ndarray of float | None + ra The right-ascention coordinate for which to retrieve the flux value. - dec : float | (Ncoord,)-shaped 1d numpy ndarray of float | None + dec The declination coordinate for which to retrieve the flux value. - E : float | (Nenergy,)-shaped 1d numpy ndarray of float | None + E The energy for which to retrieve the flux value. - t : float | (Ntime,)-shaped 1d numpy ndarray of float | None + t The observation time for which to retrieve the flux value. - angle_unit : instance of astropy.units.UnitBase | None + angle_unit The unit of the given angles. If ``None``, the set angle unit of the spatial flux profile is assumed. - energy_unit : instance of astropy.units.UnitBase | None + energy_unit The unit of the given energies. If ``None``, the set energy unit of the energy flux profile is assumed. - time_unit : instance of astropy.units.UnitBase | None + time_unit The unit of the given times. If ``None``, the set time unit of the time flux profile is assumed. Returns ------- - flux : (Ncoord,Nenergy,Ntime)-shaped ndarray of float + flux The flux values are in unit [energy]^{-1} [angle]^{-2} [length]^{-2} [time]^{-1}. """ @@ -2398,18 +2442,18 @@ def __call__( return flux - def get_param(self, name): + def get_param(self, name: str) -> float: """Retrieves the value of the given parameter. It returns ``np.nan`` if the parameter does not exist. Parameters ---------- - name : str + name The name of the parameter. Returns ------- - value : float | np.nan + value The value of the parameter. """ for obj in (super(), self._spatial_profile, self._energy_profile, self._time_profile): @@ -2419,19 +2463,19 @@ def get_param(self, name): return np.nan - def set_params(self, pdict): + def set_params(self, pdict: dict) -> bool: """Sets the parameters of the flux model. For this factorized flux model it means that it sets the parameters of the spatial, energy, and time profiles. Parameters ---------- - pdict : dict + pdict The flux parameter dictionary. Returns ------- - updated : bool + updated Flag if parameter values were actually updated. """ updated = False @@ -2456,39 +2500,39 @@ class PointlikeFFM( def __init__( self, - Phi0, - energy_profile, - time_profile, - ra=None, - dec=None, - angle_unit=None, - length_unit=None, + Phi0: float, + energy_profile: 'EnergyFluxProfile | None', + time_profile: 'TimeFluxProfile | None', + ra: float | None = None, + dec: float | None = None, + angle_unit: units.UnitBase | None = None, + length_unit: units.UnitBase | None = None, **kwargs, ): """Creates a new factorized flux model for a point-like source. Parameters ---------- - Phi0 : float + Phi0 The flux normalization constant in unit of flux. - energy_profile : instance of EnergyFluxProfile | None + energy_profile The EnergyFluxProfile instance providing the energy profile function of the flux. If set to None, an instance of UnityEnergyFluxProfile will be used, which represents the constant function 1. - time_profile : instance of TimeFluxProfile | None + time_profile The TimeFluxProfile instance providing the time profile function of the flux. If set to None, an instance of UnityTimeFluxProfile will be used, which represents the constant function 1. - ra : float | None + ra The right-ascention of the point. - dec : float | None + dec The declination of the point. - angle_unit : instance of astropy.units.UnitBase | None + angle_unit The unit for angles used for the flux unit. If set to ``None``, the configured internal angle unit is used. - length_unit : instance of astropy.units.UnitBase | None + length_unit The unit for length used for the flux unit. If set to ``None``, the configured internal length unit is used. """ @@ -2547,13 +2591,13 @@ class SteadyPointlikeFFM( def __init__( self, - Phi0, - energy_profile, - ra=None, - dec=None, - angle_unit=None, - length_unit=None, - time_unit=None, + Phi0: float, + energy_profile: 'EnergyFluxProfile | None', + ra: float | None = None, + dec: float | None = None, + angle_unit: units.UnitBase | None = None, + length_unit: units.UnitBase | None = None, + time_unit: units.UnitBase | None = None, **kwargs, ): """Creates a new factorized flux model for a point-like source with no @@ -2561,26 +2605,26 @@ def __init__( Parameters ---------- - Phi0 : float + Phi0 The flux normalization constant. - energy_profile : instance of EnergyFluxProfile | None + energy_profile The EnergyFluxProfile instance providing the energy profile function of the flux. If set to None, an instance of UnityEnergyFluxProfile will be used, which represents the constant function 1. - ra : float | None + ra The right-ascention of the point. - dec : float | None + dec The declination of the point. - angle_unit : instance of astropy.units.UnitBase | None + angle_unit The unit for angles used for the flux unit. If set to ``None``, the configured default angle unit for fluxes is used. - length_unit : instance of astropy.units.UnitBase | None + length_unit The unit for length used for the flux unit. If set to ``None``, the configured default length unit for fluxes is used. - time_unit : instance of astropy.units.UnitBase | None + time_unit The used unit for time. If set to ``None``, the configured default time unit for fluxes is used. diff --git a/skyllh/core/interpolate.py b/skyllh/core/interpolate.py index a937adf597..681b011933 100644 --- a/skyllh/core/interpolate.py +++ b/skyllh/core/interpolate.py @@ -1,6 +1,8 @@ """This module provides functionality for interpolation.""" import abc +from collections.abc import Callable +from typing import cast import numpy as np @@ -8,6 +10,7 @@ from skyllh.core.py import ( classname, ) +from skyllh.core.trialdata import TrialDataManager class GridManifoldInterpolationMethod( @@ -24,8 +27,8 @@ class GridManifoldInterpolationMethod( def __init__( self, - func, - param_grid_set, + func: Callable, + param_grid_set: ParameterGrid | ParameterGridSet, **kwargs, ): """Constructor for a GridManifoldInterpolationMethod object. @@ -33,7 +36,7 @@ def __init__( Parameters ---------- - func : callable R^D -> R + func The function that takes D parameter grid values as input and returns the value of the D-dimensional manifold at this point for each given trial event and source. @@ -43,17 +46,17 @@ def __init__( The arguments are as follows: - tdm : instance of TrialDataManager + tdm The TrialDataManager instance holding the trial event data. - eventdata : instance of numpy ndarray + eventdata A two-dimensional (V,N_events)-shaped numpy ndarray holding the event data, where N_events is the number of trial events, and V the dimensionality of the event data. - gridparams_recarray : instance of numpy.ndarray + gridparams_recarray The structured numpy ndarray of length ``len(src_idxs)`` with the D parameter names and values on the grid for all sources. - n_values : int + n_values The length of the output numpy ndarray of shape (n_values,). ``**kwargs`` Additional keyword arguments required by ``func``. @@ -64,7 +67,7 @@ def __init__( The length of the array, i.e. n_values, depends on the ``src_evt_idx`` property of the TrialDataManager. In the worst case n_values is N_sources * N_events. - param_grid_set : instance of ParameterGrid | instance of ParameterGridSet + param_grid_set The set of D parameter grids. This defines the grid of the manifold. """ @@ -85,7 +88,7 @@ def func(self, f): self._func = f @property - def param_grid_set(self): + def param_grid_set(self) -> 'ParameterGridSet': """The ParameterGridSet instance defining the set of D parameter grids. This defines the grid of the manifold. """ @@ -107,24 +110,24 @@ def ndim(self): @abc.abstractmethod def __call__( self, - tdm, - eventdata, - params_recarray, + tdm: TrialDataManager, + eventdata: np.ndarray, + params_recarray: np.ndarray, **kwargs, - ): + ) -> tuple[np.ndarray, np.ndarray]: """Retrieves the interpolated value of the manifold at the D-dimensional point ``params_recarray`` for all given events and sources, along with the D gradients, i.e. partial derivatives. Parameters ---------- - tdm : instance of TrialDataManager + tdm The TrialDataManager instance holding the trial data. - eventdata : numpy ndarray + eventdata The 2D (V,N_events)-shaped numpy ndarray holding the event data, where N_events is the number of trial events, and V the dimensionality of the event data. - params_recarray : instance of numpy record ndarray + params_recarray The structured numpy ndarray holding the N_sources set of parameter names and values, that define the point (for each source) on the manifold for which the value should get calculated for each event. @@ -133,10 +136,10 @@ def __call__( Returns ------- - values : instance of numpy.ndarray + values The (N,)-shaped numpy ndarray holding the interpolated manifold values for the given events and sources. - grads : instance of numpy.ndarray + grads The (D,N)-shaped numpy ndarray holding the D manifold gradients for the N given values, where D is the number of parameters. The order of the D parameters is defined by the ParameterGridSet @@ -157,22 +160,22 @@ class NullGridManifoldInterpolationMethod( def __init__( self, - func, - param_grid_set, + func: Callable, + param_grid_set: ParameterGrid | ParameterGridSet, **kwargs, ): """Creates a new NullGridManifoldInterpolationMethod instance. Parameters ---------- - func : callable R^d -> R + func The function that takes d parameter grid values as input and returns the value of the d-dimensional manifold at this point for each given trial event and source. See the documentation of the :class:`~skyllh.core.interpolate.GridManifoldInterpolationMethod` class for more details. - param_grid_set : instance of ParameterGrid | instance of ParameterGridSet + param_grid_set The set of d parameter grids. This defines the grid of the manifold. """ @@ -180,24 +183,24 @@ class for more details. def __call__( self, - tdm, - eventdata, - params_recarray, + tdm: TrialDataManager, + eventdata: np.ndarray, + params_recarray: np.ndarray, **kwargs, - ): + ) -> tuple[np.ndarray, np.ndarray]: """Calculates the non-interpolated manifold value and its gradient (zero) for each given event and source at the points given by ``params_recarray``. Parameters ---------- - tdm : instance of TrialDataManager + tdm The TrialDataManager instance holding the trial data. - eventdata : instance of numpy.ndarray + eventdata The (V,N_events)-shaped numpy ndarray holding the event data, where N_events is the number of events, and V the dimensionality of the event data. - params_recarray : instance of numpy.ndarray + params_recarray The structured numpy ndarray of length N_sources holding the parameter names and values of the sources, defining the point on the manifold for which the values should get calculated. @@ -206,10 +209,10 @@ def __call__( Returns ------- - values : instance of numpy.ndarray + values The (N,)-shaped numpy ndarray holding the interpolated manifold values for the given events and sources. - grads : instance of numpy.ndarray + grads The (D,N)-shaped ndarray of float holding the D manifold gradients for the N values, where D is the number of parameters of the manifold. @@ -225,8 +228,15 @@ def __call__( pvalues = params_recarray[pname] gridparams_recarray[pname] = p_grid.round_to_nearest_grid_point(pvalues) - values = self._func( - tdm=tdm, eventdata=eventdata, gridparams_recarray=gridparams_recarray, n_values=tdm.get_n_values(), **kwargs + values = cast( + np.ndarray, + self._func( + tdm=tdm, + eventdata=eventdata, + gridparams_recarray=gridparams_recarray, + n_values=tdm.get_n_values(), + **kwargs, + ), ) grads = np.zeros((len(self.param_grid_set), len(values)), dtype=np.float64) @@ -243,22 +253,22 @@ class Linear1DGridManifoldInterpolationMethod( def __init__( self, - func, - param_grid_set, + func: Callable, + param_grid_set: ParameterGrid | ParameterGridSet, **kwargs, ): """Creates a new Linear1DGridManifoldInterpolationMethod instance. Parameters ---------- - func : callable R -> R + func The function that takes the parameter grid value as input and returns the value of the 1-dimensional manifold at this point for each given source and trial event. See the documentation of the :class:`~skyllh.core.interpolate.GridManifoldInterpolationMethod` class for more details. - param_grid_set : instance of ParameterGrid | instance of ParameterGridSet + param_grid_set The one parameter grid. This defines the grid of the manifold. """ super().__init__(func=func, param_grid_set=param_grid_set, **kwargs) @@ -277,26 +287,26 @@ class for more details. def _create_cache( self, - trial_data_state_id, - x0, - m, - b, + trial_data_state_id: int | None, + x0: np.ndarray | None, + m: np.ndarray | None, + b: np.ndarray | None, ): """Creates a cache for the line parameterization for the last manifold grid point for the N_events different events. Parameters ---------- - trial_data_state_id : int | None + trial_data_state_id The trial data state id of the TrialDataManager. - x0 : instance of ndarray | None + x0 The (N_sources,)-shaped numpy ndarray holding the parameter grid value of the lower point of the grid manifold for each source used to estimate the line. - m : instance of ndarray | None + m The (N_values,)-shaped numpy ndarray holding the slope of the line for each trial event and source. - b : instance of ndarray | None + b The (N_values,)-shaped numpy ndarray holding the offset coefficient of the line for each trial event and source. """ @@ -308,18 +318,18 @@ def _is_cached( self, trial_data_state_id, x0, - ): + ) -> bool: """Checks if the given line parametrization are already cached for the given x0 values. Returns ------- - check : bool + check ``True`` if the line parametrization for x0 is already cached, ``False`` otherwise. """ self__cache = self._cache - return ( + return bool( (self__cache['trial_data_state_id'] is not None) and (self__cache['trial_data_state_id'] == trial_data_state_id) and (np.all(np.isclose(self__cache['x0'], x0))) @@ -327,23 +337,23 @@ def _is_cached( def __call__( self, - tdm, - eventdata, - params_recarray, + tdm: TrialDataManager, + eventdata: np.ndarray, + params_recarray: np.ndarray, **kwargs, - ): + ) -> tuple[np.ndarray, np.ndarray]: """Calculates the interpolated manifold value and its gradient for each given source and trial event at the point ``params_recarray``. Parameters ---------- - tdm : instance of TrialDataManager + tdm The TrialDataManager instance holding the trial data. - eventdata : instance of numpy ndarray + eventdata The (V,N_events)-shaped numpy ndarray holding the event data, where N_events is the number of events, and V the dimensionality of the event data. - params_recarray : numpy record ndarray + params_recarray The numpy record ndarray of length N_sources holding the parameter names and values for each source, defining the point on the manifold for which the value should get calculated. @@ -354,10 +364,10 @@ def __call__( Returns ------- - values : instance of numpy.ndarray + values The (N_values,)-shaped numpy ndarray of float holding the interpolated manifold values for all sources and trial events. - grads : instance of numpy.ndarray + grads The (D,N_values)-shaped numpy ndarray of float holding the D manifold gradients for the N_values values for all sources and trial events, where D is the number of interpolation parameters. @@ -393,13 +403,19 @@ def __call__( m = np.empty((n_values,), dtype=np.float64) gridparams_recarray = np.array(x0, dtype=[(xname, np.float64)]) - M0 = self__func( - tdm=tdm, eventdata=eventdata, gridparams_recarray=gridparams_recarray, n_values=n_values, **kwargs + M0: np.ndarray = cast( + np.ndarray, + self__func( + tdm=tdm, eventdata=eventdata, gridparams_recarray=gridparams_recarray, n_values=n_values, **kwargs + ), ) gridparams_recarray = np.array(x1, dtype=[(xname, np.float64)]) - M1 = self__func( - tdm=tdm, eventdata=eventdata, gridparams_recarray=gridparams_recarray, n_values=n_values, **kwargs + M1: np.ndarray = cast( + np.ndarray, + self__func( + tdm=tdm, eventdata=eventdata, gridparams_recarray=gridparams_recarray, n_values=n_values, **kwargs + ), ) # Broadcast x0 and x1 to the values array. @@ -426,22 +442,22 @@ class Parabola1DGridManifoldInterpolationMethod( def __init__( self, - func, - param_grid_set, + func: Callable, + param_grid_set: ParameterGrid | ParameterGridSet, **kwargs, ): """Creates a new Parabola1DGridManifoldInterpolationMethod instance. Parameters ---------- - func : callable R -> R + func The function that takes the parameter grid value as input and returns the value of the 1-dimensional manifold at this point for each given source and trial event. See the documentation of the :class:`~skyllh.core.interpolate.GridManifoldInterpolationMethod` class for more details. - param_grid_set : instance of ParameterGrid | instance of ParameterGridSet + param_grid_set The one parameter grid. This defines the grid of the manifold. """ super().__init__(func=func, param_grid_set=param_grid_set, **kwargs) @@ -460,37 +476,37 @@ class for more details. def _create_cache( self, - trial_data_state_id, - x1, - M1, - a, - b, - ): + trial_data_state_id: int | None, + x1: np.ndarray | None, + M1: np.ndarray | None, + a: np.ndarray | None, + b: np.ndarray | None, + ) -> dict: """Creates a cache for the parabola parameterization for the last manifold grid point for the N_events different events. Parameters ---------- - trial_data_state_id : int | None + trial_data_state_id The trial data state ID of the TrialDataManager. - x1 : instance of numpy.ndarray | None + x1 The (N_sources,)-shaped numpy ndarray of float holding the parameter grid value for the middle point of the grid manifold for all sources used to estimate the parabola. - M1 : instance of numpy.ndarray + M1 The (N_values,)-shaped numpy ndarray of float holding the grid manifold value for each source and trial event of the middle point (x1,). - a : instance of numpy.ndarray + a The (N_values,)-shaped numpy ndarray of float holding the parabola coefficient ``a`` for each source and trial event. - b : instance of numpy.ndarray + b The (N_values,)-shaped numpy ndarray of float holding the parabola coefficient ``b`` for each source and trial event. Returns ------- - cache : dict + cache The dictionary holding the cache data. """ cache = { @@ -521,23 +537,23 @@ def _is_cached( def __call__( self, - tdm, - eventdata, - params_recarray, + tdm: TrialDataManager, + eventdata: np.ndarray, + params_recarray: np.ndarray, **kwargs, - ): + ) -> tuple[np.ndarray, np.ndarray]: """Calculates the interpolated manifold value and its gradient for each given source and trial event at the point ``params_recarray``. Parameters ---------- - tdm : instance of TrialDataManager + tdm The TrialDataManager instance holding the trial data. - eventdata : instance of numpy ndarray + eventdata The (V,N_events)-shaped numpy ndarray holding the event data, where N_events is the number of events, and V the dimensionality of the event data. - params_recarray : numpy record ndarray + params_recarray The numpy record ndarray of length N_sources holding the parameter names and values for each source, defining the point on the manifold for which the value should get calculated. @@ -548,9 +564,9 @@ def __call__( Returns ------- - values : (N_values,) ndarray of float + values The interpolated manifold value for the N given events. - grads : (D,N_values) ndarray of float + grads The D manifold gradients for the N given events, where D is the number of parameters. """ @@ -579,18 +595,27 @@ def __call__( n_values = tdm.get_n_values() gridparams_recarray = np.array(x0, dtype=[(xname, np.float64)]) - M0 = self__func( - tdm=tdm, eventdata=eventdata, gridparams_recarray=gridparams_recarray, n_values=n_values, **kwargs + M0: np.ndarray = cast( + np.ndarray, + self__func( + tdm=tdm, eventdata=eventdata, gridparams_recarray=gridparams_recarray, n_values=n_values, **kwargs + ), ) gridparams_recarray = np.array(x1, dtype=[(xname, np.float64)]) - M1 = self__func( - tdm=tdm, eventdata=eventdata, gridparams_recarray=gridparams_recarray, n_values=n_values, **kwargs + M1: np.ndarray = cast( + np.ndarray, + self__func( + tdm=tdm, eventdata=eventdata, gridparams_recarray=gridparams_recarray, n_values=n_values, **kwargs + ), ) gridparams_recarray = np.array(x2, dtype=[(xname, np.float64)]) - M2 = self__func( - tdm=tdm, eventdata=eventdata, gridparams_recarray=gridparams_recarray, n_values=n_values, **kwargs + M2: np.ndarray = cast( + np.ndarray, + self__func( + tdm=tdm, eventdata=eventdata, gridparams_recarray=gridparams_recarray, n_values=n_values, **kwargs + ), ) a = 0.5 * (M0 - 2.0 * M1 + M2) / dx**2 diff --git a/skyllh/core/livetime.py b/skyllh/core/livetime.py index 39f3525b90..8ce99b139f 100644 --- a/skyllh/core/livetime.py +++ b/skyllh/core/livetime.py @@ -1,11 +1,14 @@ """The livetime module provides general functionality for detector up-time.""" +from collections.abc import Sequence + import numpy as np from skyllh.core.py import ( classname, issequence, ) +from skyllh.core.random import RandomStateService class Livetime: @@ -14,34 +17,32 @@ class Livetime: """ @staticmethod - def get_integrated_livetime(livetime): + def get_integrated_livetime(livetime: 'float | Livetime') -> float: """Gets the integrated live-time from the given livetime argument, which is either a scalar value or an instance of Livetime. Parameters ---------- - livetime : float | Livetime instance + livetime The live-time in days as float, or an instance of Livetime. Returns ------- - intgrated_livetime : float + integrated_livetime The integrated live-time. """ - intgrated_livetime = livetime - if isinstance(livetime, Livetime): - intgrated_livetime = livetime.livetime + return float(livetime.livetime) - return intgrated_livetime + return float(livetime) - def __init__(self, uptime_mjd_intervals_arr, **kwargs): + def __init__(self, uptime_mjd_intervals_arr: np.ndarray, **kwargs): """Creates a new Livetime object from a (N,2)-shaped ndarray holding the uptime intervals. Parameters ---------- - uptime_mjd_intervals_arr : (N,2)-shaped ndarray + uptime_mjd_intervals_arr The (N,2)-shaped ndarray holding the start and end times of each up-time interval. @@ -60,13 +61,13 @@ def __init__(self, uptime_mjd_intervals_arr, **kwargs): self.uptime_mjd_intervals_arr = uptime_mjd_intervals_arr - def assert_mjd_intervals_integrity(self, arr): + def assert_mjd_intervals_integrity(self, arr: np.ndarray): """Checks if the given MJD interval array conforms with all its data requirements. Parameters ---------- - arr : instance of numpy ndarray + arr The (N,2)-shaped numpy ndarray holding the up-time intervals. Raises @@ -158,14 +159,14 @@ def __str__(self): s = f'{classname(self)}(time_window=({self.time_window[0]:.6f}, {self.time_window[1]:.6f}))' return s - def _get_onoff_intervals(self): + def _get_onoff_intervals(self) -> np.ndarray: """A view on the uptime intervals where each time is a lower bin edge. Hence, odd array elements (bins) are on-time intervals, and even array elements are off-time intervals. Returns ------- - onoff_intervals : instance of numpy ndarray + onoff_intervals The (n_uptime_intervals*2,)-shaped numpy ndarray holding the time edges of the uptime intervals. """ @@ -173,7 +174,7 @@ def _get_onoff_intervals(self): return onoff_intervals - def _get_onoff_interval_indices(self, mjds): + def _get_onoff_interval_indices(self, mjds: np.ndarray) -> np.ndarray: """Retrieves the indices of the on-time and off-time intervals, which correspond to the given MJD values. @@ -187,12 +188,12 @@ def _get_onoff_interval_indices(self, mjds): Parameters ---------- - mjds : numpy array of floats + mjds The array of MJD values. Returns ------- - idxs : numpy array of ints + idxs The array of the on-off-time interval indices that correspond to the given MJD values. """ @@ -205,37 +206,39 @@ def _get_onoff_interval_indices(self, mjds): return idxs - def get_uptime_intervals_between(self, t_start, t_end): + def get_uptime_intervals_between(self, t_start: float, t_end: float) -> np.ndarray: """Creates a (N,2)-shaped ndarray holding the on-time detector intervals between the given time range from t_start to t_end. Parameters ---------- - t_start : float + t_start The MJD start time of the time range to consider. This might be the lower bound of the first on-time interval. - t_end : float + t_end The MJD end time of the time range to consider. This might be the upper bound of the last on-time interval. Returns ------- - ontime_intervals : (N,2)-shaped ndarray + ontime_intervals The (N,2)-shaped ndarray holding the on-time detector intervals. """ onoff_intervals = self._get_onoff_intervals() - (t_start_idx, t_end_idx) = self._get_onoff_interval_indices((t_start, t_end)) + (t_start_idx, t_end_idx) = self._get_onoff_interval_indices(np.array([t_start, t_end])) + t_start_f: float = t_start + t_end_f: float = t_end if t_start_idx % 2 == 0: # t_start is during off-time. Use the next on-time lower edge as # first on-time edge. - t_start = onoff_intervals[t_start_idx] + t_start_f = float(onoff_intervals[t_start_idx]) else: t_start_idx -= 1 if t_end_idx % 2 == 0: # t_end is during off-time. Use the previous on-time upper edge as # the last on-time edge. - t_end = onoff_intervals[t_end_idx - 1] + t_end_f = float(onoff_intervals[t_end_idx - 1]) else: t_end_idx += 1 @@ -244,8 +247,8 @@ def get_uptime_intervals_between(self, t_start, t_end): ontime_intervals_flat = np.empty((N_ontime_intervals * 2,), dtype=np.float64) # Set the first and last on-time interval edges. - ontime_intervals_flat[0] = t_start - ontime_intervals_flat[-1] = t_end + ontime_intervals_flat[0] = t_start_f + ontime_intervals_flat[-1] = t_end_f if N_ontime_intervals > 1: # Fill also the interval edges of the intermediate on-time bins. ontime_intervals_flat[1:-1] = onoff_intervals[t_start_idx + 1 : t_end_idx - 1] @@ -254,18 +257,18 @@ def get_uptime_intervals_between(self, t_start, t_end): return ontime_intervals - def get_livetime_upto(self, mjd): + def get_livetime_upto(self, mjd: float | np.ndarray) -> float | np.ndarray: """Calculates the cumulative detector livetime up to the given time. Parameters ---------- - mjd : float | array of floats + mjd The time in MJD up to which the detector livetime should be calculated. Returns ------- - livetimes : float | ndarray of floats + livetimes The ndarray holding the cumulative detector livetime corresponding to the the given MJD times. """ @@ -305,18 +308,18 @@ def get_livetime_upto(self, mjd): return livetimes - def is_on(self, mjd): + def is_on(self, mjd: float | Sequence[float] | np.ndarray) -> np.ndarray: """Checks if the detector is on at the given MJD time. MJD times outside any live-time interval will be masked as False. Parameters ---------- - mjd : float | sequence of float + mjd The time in MJD. Returns ------- - is_on : array of bool + is_on True if the detector was on at the given time. """ mjd = np.atleast_1d(mjd) @@ -331,37 +334,37 @@ def is_on(self, mjd): return is_on - def draw_ontimes(self, rss, size, t_min=None, t_max=None): + def draw_ontimes( + self, rss: RandomStateService, size: int, t_min: float | None = None, t_max: float | None = None + ) -> np.ndarray: """Draws random MJD times based on the detector on-time intervals. Parameters ---------- - rss : RandomStateService + rss The skyllh RandomStateService instance to use for drawing random numbers from. - size : int + size The number of random MJD times to generate. - t_min : float + t_min The optional minimal time to consider. If set to ``None``, the start time of this Livetime instance will be used. - t_max : float + t_max The optional maximal time to consider. If set to ``None``, the end time of this Livetime instance will be used. Returns ------- - ontimes : ndarray + ontimes The 1d array holding the generated MJD times. """ uptime_intervals_arr = self._uptime_mjd_intervals_arr if t_min is not None or t_max is not None: - if t_min is None: - t_min = self.time_start - if t_max is None: - t_max = self.time_stop + t_min_val = float(t_min) if t_min is not None else float(self.time_start) + t_max_val = float(t_max) if t_max is not None else float(self.time_stop) - uptime_intervals_arr = self.get_uptime_intervals_between(t_min, t_max) + uptime_intervals_arr = self.get_uptime_intervals_between(t_min_val, t_max_val) onoff_intervals = np.reshape(uptime_intervals_arr, (uptime_intervals_arr.size,)) diff --git a/skyllh/core/llhratio.py b/skyllh/core/llhratio.py index 12ff22cc4f..79194d0576 100644 --- a/skyllh/core/llhratio.py +++ b/skyllh/core/llhratio.py @@ -29,6 +29,7 @@ float_cast, issequenceof, ) +from skyllh.core.random import RandomStateService from skyllh.core.services import ( DatasetSignalWeightFactorsService, SrcDetSigYieldWeightsService, @@ -36,9 +37,7 @@ from skyllh.core.source_hypo_grouping import ( SourceHypoGroupManager, ) -from skyllh.core.timing import ( - TaskTimer, -) +from skyllh.core.timing import TaskTimer, TimeLord from skyllh.core.trialdata import ( TrialDataManager, ) @@ -52,15 +51,15 @@ class LLHRatio( ): """Abstract base class for a log-likelihood (LLH) ratio function.""" - def __init__(self, pmm, minimizer, **kwargs): + def __init__(self, pmm: ParameterModelMapper, minimizer: Minimizer, **kwargs): """Creates a new LLH ratio function instance. Parameters ---------- - pmm : instance of ParameterModelMapper + pmm The instance of ParameterModelMapper providing the mapping of global parameters to local parameters of individual models. - minimizer : instance of Minimizer + minimizer The Minimizer instance that should be used to minimize the negative of this log-likelihood ratio function. """ @@ -102,68 +101,70 @@ def minimizer(self, minimizer): self._minimizer = minimizer @abc.abstractmethod - def initialize_for_new_trial(self, tl=None, **kwargs): + def initialize_for_new_trial(self, tl: TimeLord | None = None, **kwargs): """This method will be called by the Analysis class after new trial data has been initialized to the trial data manager. Derived classes can make use of this call hook to perform LLHRatio specific trial initialization. Parameters ---------- - tl : instance of TimeLord | None + tl The optional instance of TimeLord to use for timing measurements. """ @abc.abstractmethod - def evaluate(self, fitparam_values, src_params_recarray=None, tl=None): + def evaluate( + self, fitparam_values: np.ndarray, src_params_recarray: np.ndarray | None = None, tl: TimeLord | None = None + ) -> tuple[float, np.ndarray]: """This method evaluates the LLH ratio function for the given set of fit parameter values. Parameters ---------- - fitparam_values : instance of numpy ndarray + fitparam_values The (N_fitparams,)-shaped numpy 1D ndarray holding the current values of the global fit parameters. - src_params_recarray : instance of numpy record ndarray | None + src_params_recarray The numpy record ndarray of length N_sources holding the parameter names and values of all sources. If set to ``None`` it will be created automatically from the ``fitparam_values`` array. See the documentation of the :meth:`skyllh.core.parameters.ParameterModelMapper.create_src_params_recarray` method for more information about this array. - tl : instance of TimeLord | None + tl The optional instance of TimeLord to use for measuring timing. Returns ------- - log_lambda : float + log_lambda The calculated log-lambda value. - grads : instance of numpy ndarray + grads The (N_fitparams,)-shaped 1D numpy ndarray holding the gradient value for each global fit parameter. """ - def maximize(self, rss, tl=None): + def maximize(self, rss: RandomStateService, tl: TimeLord | None = None) -> tuple[float, np.ndarray, dict]: """Maximize the log-likelihood ratio function, by using the ``evaluate`` method. Parameters ---------- - rss : instance of RandomStateService + rss The instance of RandomStateService to draw random numbers from. This is needed to generate random parameter initial values. - tl : instance of TimeLord | None + tl The optional instance of TimeLord that should be used to time the maximization process. Returns ------- - log_lambda_max : float + log_lambda_max The (maximum) value of the log-likelihood ratio (log_lambda) function for the best fit parameter values. - fitparam_values : instance of numpy ndarray + fitparam_values The (N_fitparam,)-shaped 1D numpy ndarray holding the values of the global fit parameters. - status : dict + status The dictionary with status information about the maximization process, i.e. from the minimizer. """ @@ -173,6 +174,10 @@ def maximize(self, rss, tl=None): self_evaluate = self.evaluate def negative_llhratio_func(fitparam_values, func_stats, tl=None): + """Evaluates the log-likelihood ratio function for the given fit + parameter values and returns the negative value and its negative + gradients, suitable for minimization. + """ src_params_recarray = self._pmm.create_src_params_recarray(fitparam_values) func_stats['n_calls'] += 1 @@ -208,18 +213,18 @@ class TCLLHRatio(LLHRatio, metaclass=abc.ABCMeta): components, i.e. signal and background. """ - def __init__(self, pmm, minimizer, mean_n_sig_0, **kwargs): + def __init__(self, pmm: ParameterModelMapper, minimizer: Minimizer, mean_n_sig_0: float, **kwargs): """Creates a new two-component LLH ratio function instance. Parameters ---------- - pmm : instance of ParameterModelMapper + pmm The instance of ParameterModelMapper providing the mapping of global floating parameters to individual models. - minimizer : instance of Minimizer + minimizer The Minimizer instance that should be used to minimize the negative of this log-likelihood ratio function. - mean_n_sig_0 : float + mean_n_sig_0 The fixed mean number of signal events for the null-hypothesis. """ super().__init__(pmm=pmm, minimizer=minimizer, **kwargs) @@ -239,58 +244,61 @@ def mean_n_sig_0(self, v): self._mean_n_sig_0 = v @abc.abstractmethod - def calculate_ns_grad2(self, ns, ns_pidx, src_params_recarray, tl=None, **kwargs): + def calculate_ns_grad2( + self, ns: float, ns_pidx: int, src_params_recarray: np.ndarray, tl: TimeLord | None = None, **kwargs + ) -> float: """This method is supposed to calculate the second derivative of the log-likelihood ratio function w.r.t. the fit parameter ns, the number of signal events in the data set. Parameters ---------- - fitparam_values : instance of numpy ndarray - The (N_fitparams,)-shaped 1D numpy ndarray holding the current - values of the global fit parameters. - ns_pidx : int + ns + The value of the global fit parameter ns. + ns_pidx The index of the global ns fit parameter. - src_params_recarray : instance of numpy record ndarray + src_params_recarray The numpy record ndarray of length N_sources holding the parameter names and values of all sources. See the documentation of the :meth:`skyllh.core.parameters.ParameterModelMapper.create_src_params_recarray` method for more information about this array. - tl : instance of TimeLord | None + tl The optional instance of TimeLord that should be used for timing measurements. Returns ------- - nsgrad2 : float + nsgrad2 The second derivative w.r.t. ns of the log-likelihood ratio function for the given fit parameter values. """ - def maximize_with_1d_newton_rapson_minimizer(self, rss, tl=None): + def maximize_with_1d_newton_rapson_minimizer( + self, rss: RandomStateService, tl: TimeLord | None = None + ) -> tuple[float, np.ndarray, dict]: """Maximizes this log-likelihood ratio function, by minimizing its negative using a 1D Newton-Rapson minimizer. Parameters ---------- - rss : instance of RandomStateService + rss The instance of RandomStateService that should be used to draw random numbers from. It is used by the minimizer to generate random fit parameter initial values. - tl : instance of TimeLord | None + tl The optional instance of TimeLord that should be used to time the maximization of the LLH-ratio function. Returns ------- - log_lambda_max : float + log_lambda_max The (maximum) value of the log-likelihood ratio (log_lambda) function for the best fit parameter values. - fitparam_values : instance of numpy ndarray + fitparam_values The (N_fitparam,)-shaped 1D numpy ndarray holding the global fit parameter values. - status : dict + status The dictionary with status information about the maximization process, i.e. from the minimizer. """ @@ -303,6 +311,10 @@ def maximize_with_1d_newton_rapson_minimizer(self, rss, tl=None): ns_pidx = self._pmm.get_gflp_idx(name='ns') def negative_llhratio_func_nr1d_ns(fitparam_values, tl): + """Evaluates the log-likelihood ratio function for the given fit + parameter values and returns the negative first and second + derivatives w.r.t. ns, suitable for the 1D Newton-Rapson minimizer. + """ ns = fitparam_values[ns_pidx] src_params_recarray = self._pmm.create_src_params_recarray(fitparam_values) with TaskTimer(tl, 'Evaluate llh-ratio function.'): @@ -323,7 +335,7 @@ def negative_llhratio_func_nr1d_ns(fitparam_values, tl): return (log_lambda_max, fitparam_values, status) - def maximize(self, rss, tl=None): + def maximize(self, rss: RandomStateService, tl: TimeLord | None = None) -> tuple[float, np.ndarray, dict]: """Maximizes this log-likelihood ratio function, by minimizing its negative. This method has a special implementation when a 1D Newton-Rapson @@ -332,23 +344,23 @@ def maximize(self, rss, tl=None): Parameters ---------- - rss : instance of RandomStateService + rss The instance of RandomStateService that should be used to draw random numbers from. It is used by the minimizer to generate random fit parameter initial values. - tl : instance of TimeLord | None + tl The optional instance of TimeLord that should be used to time the maximization of the LLH-ratio function. Returns ------- - log_lambda_max : float + log_lambda_max The (maximum) value of the log-likelihood ratio (log_lambda) function for the best fit parameter values. - fitparam_values : instance of numpy ndarray + fitparam_values The (N_fitparam,)-shaped 1D numpy ndarray holding the global fit parameter values. - status : dict + status The dictionary with status information about the maximization process, i.e. from the minimizer. """ @@ -363,25 +375,33 @@ class SingleDatasetTCLLHRatio(TCLLHRatio, metaclass=abc.ABCMeta): components, i.e. signal and background, for a single data set. """ - def __init__(self, pmm, minimizer, shg_mgr, tdm, mean_n_sig_0, **kwargs): + def __init__( + self, + pmm: ParameterModelMapper, + minimizer: Minimizer, + shg_mgr: SourceHypoGroupManager, + tdm: TrialDataManager, + mean_n_sig_0: float, + **kwargs, + ): """Creates a new two-component LLH ratio function instance for a single data set. Parameters ---------- - pmm : instance of ParameterModelMapper + pmm The instance of ParameterModelMapper providing the mapping of global floating parameters to individual models. - minimizer : instance of Minimizer + minimizer The Minimizer instance that should be used to minimize the negative of this log-likelihood ratio function. - shg_mgr : SourceHypoGroupManager instance + shg_mgr The SourceHypoGroupManager instance that defines the source hypothesis groups. - tdm : instance of TrialDataManager + tdm The instance of TrialDataManager that holds the trial event data and additional data fields for this LLH ratio function. - mean_n_sig_0 : float + mean_n_sig_0 The fixed mean number of signal events for the null-hypothesis. """ super().__init__(pmm=pmm, minimizer=minimizer, mean_n_sig_0=mean_n_sig_0, **kwargs) @@ -424,13 +444,13 @@ def tdm(self, mgr): ) self._tdm = mgr - def change_shg_mgr(self, shg_mgr): + def change_shg_mgr(self, shg_mgr: SourceHypoGroupManager): """Changes the source hypothesis group manager of this two-component LLH ratio function. Parameters ---------- - shg_mgr : instance of SourceHypoGroupManager + shg_mgr The new instance of SourceHypoGroupManager. """ self.shg_mgr = shg_mgr @@ -451,24 +471,32 @@ class ZeroSigH0SingleDatasetTCLLHRatio(SingleDatasetTCLLHRatio): # instance member, because it is supposed to be the same for all instances. _one_plus_alpha = 1e-3 - def __init__(self, pmm, minimizer, shg_mgr, tdm, pdfratio, **kwargs): + def __init__( + self, + pmm: ParameterModelMapper, + minimizer: Minimizer, + shg_mgr: SourceHypoGroupManager, + tdm: TrialDataManager, + pdfratio: PDFRatio, + **kwargs, + ): """Constructor of the two-component log-likelihood ratio function. Parameters ---------- - pmm : instance of ParameterModelMapper + pmm The instance of ParameterModelMapper providing the mapping of global floating parameters to individual models. - minimizer : instance of Minimizer + minimizer The Minimizer instance that should be used to minimize the negative of this log-likelihood ratio function. - shg_mgr : instance of SourceHypoGroupManager + shg_mgr The SourceHypoGroupManager instance that defines the source hypothesis groups. - tdm : instance of TrialDataManager + tdm The instance of TrialDataManager that holds the trial event data and additional data fields for this LLH ratio function. - pdfratio : instance of PDFRatio + pdfratio The instance of PDFRatio. A PDFRatio instance might depend on none, one, or several fit parameters. """ @@ -483,7 +511,7 @@ def __init__(self, pmm, minimizer, shg_mgr, tdm, pdfratio, **kwargs): @SingleDatasetTCLLHRatio.mean_n_sig_0.setter def mean_n_sig_0(self, v): - SingleDatasetTCLLHRatio.mean_n_sig_0.fset(self, v) + SingleDatasetTCLLHRatio.mean_n_sig_0.fset(self, v) # pyright: ignore[reportOptionalCall] if self._mean_n_sig_0 != 0: raise ValueError(f'The {classname(self)} class is only valid for mean_n_sig_0 = 0!') @@ -500,7 +528,7 @@ def pdfratio(self, r): ) self._pdfratio = r - def initialize_for_new_trial(self, tl=None, **kwargs): + def initialize_for_new_trial(self, tl: TimeLord | None = None, **kwargs): """Initializes the log-likelihood ratio function for a new trial. It calls the :meth:`~skyllh.core.pdfratio.PDFRatio.initialize_for_new_trial` method @@ -508,46 +536,41 @@ def initialize_for_new_trial(self, tl=None, **kwargs): Parameters ---------- - tl : instance of TimeLord + tl The optional instance of TimeLord to measure timing information. """ self._pdfratio.initialize_for_new_trial(tdm=self._tdm, tl=tl, **kwargs) - def calculate_log_lambda_and_grads(self, N, ns, ns_pidx, p_mask, Xi, dXi_dp): + def calculate_log_lambda_and_grads( + self, N: int, ns: float, ns_pidx: int, p_mask: np.ndarray, Xi: np.ndarray, dXi_dp: np.ndarray + ) -> tuple[float, np.ndarray]: """Calculates the log(Lambda) value and its gradient for each global fit parameter. This calculation is source and detector independent. Parameters ---------- - fitparam_values : instance of numpy ndarray - The (N_fitparams,)-shaped ndarray holding the current values of the - global fit parameters. - These numbers are used as cache key to validate the ``nsgrad_i`` - values for the given fit parameter values for a possible later - calculation of the second derivative w.r.t. ns of the log-likelihood - ratio function. - N : int + N The total number of events. - ns : float + ns The value of the global fit parameter ns. - ns_pidx : int + ns_pidx The index of the global fit parameter ns. - p_mask : instance of numpy ndarray + p_mask The (N_fitparam,)-shaped numpy ndarray of bool selecting all global fit parameters, except ns. - Xi : instance of numpy ndarray + Xi The (n_selected_events,)-shaped 1D numpy ndarray holding the X value of each selected event. - dXi_dp : instance of numpy ndarray + dXi_dp The (n_selected_events, N_fitparams-1,)-shaped 2D ndarray holding the derivative value for each fit parameter p (i.e. except ns) of each event's X value. Returns ------- - log_lambda : float + log_lambda The value of the log-likelihood ratio function. - grads : instance of numpy ndarray + grads The (N_fitparams,)-shaped numpy ndarray holding the gradient value of log_lambda for each fit parameter. """ @@ -580,6 +603,7 @@ def calculate_log_lambda_and_grads(self, N, ns, ns_pidx, p_mask, Xi, dXi_dp): np.log1p(alpha_i, where=m_stable, out=log_lambda_i) # Calculate the log_lambda_i value for the numerical unstable events. + tildealpha_i: np.ndarray = np.empty(0, dtype=np.float64) if any_unstable_events: tildealpha_i = (alpha_i[m_unstable] - alpha) / one_plus_alpha log_lambda_i[m_unstable] = np.log1p(alpha) + tildealpha_i - 0.5 * tildealpha_i**2 @@ -620,7 +644,13 @@ def calculate_log_lambda_and_grads(self, N, ns, ns_pidx, p_mask, Xi, dXi_dp): return (log_lambda, grads) - def calculate_ns_grad2(self, ns, ns_pidx=None, src_params_recarray=None, tl=None): + def calculate_ns_grad2( # pyright: ignore[reportIncompatibleMethodOverride] + self, + ns: float, + ns_pidx: int | None = None, + src_params_recarray: np.ndarray | None = None, + tl: TimeLord | None = None, + ) -> float: """Calculates the second derivative w.r.t. ns of the log-likelihood ratio function. This method tries to use cached values for the first derivative @@ -631,28 +661,25 @@ def calculate_ns_grad2(self, ns, ns_pidx=None, src_params_recarray=None, tl=None Parameters ---------- - fitparam_values : numpy (N_fitparams+1)-shaped 1D ndarray - The ndarray holding the current values of the global fit - parameters. - ns : float + ns The value of the global fit parameter ns. - ns_pidx : int + ns_pidx The parameter index of the global fit parameter ns. For this particular class this is an ignored interface parameter. - src_params_recarray : instance of numpy record ndarray + src_params_recarray The numpy record ndarray of length N_sources holding the parameter names and values of all sources. See the documentation of the :meth:`skyllh.core.parameters.ParameterModelMapper.create_src_params_recarray` method for more information about this array. For this particular class this is an ignored interface parameter. - tl : instance of TimeLord | None + tl The optional instance of TimeLord that should be used for timing measurements. Returns ------- - nsgrad2 : float + nsgrad2 The second derivative w.r.t. ns of the log-likelihood ratio function for the given fit parameter values. """ @@ -668,30 +695,32 @@ def calculate_ns_grad2(self, ns, ns_pidx=None, src_params_recarray=None, tl=None return nsgrad2 - def evaluate(self, fitparam_values, src_params_recarray=None, tl=None): + def evaluate( + self, fitparam_values: np.ndarray, src_params_recarray: np.ndarray | None = None, tl: TimeLord | None = None + ) -> tuple[float, np.ndarray]: """Evaluates the log-likelihood ratio function for the given set of data events. Parameters ---------- - fitparam_values : instance of numpy ndarray + fitparam_values The (N_fitparams,)-shaped 1D ndarray holding the current values of the global fit parameters. - src_params_recarray : instance of numpy structured ndarray | None + src_params_recarray The numpy record ndarray of length N_sources holding the local parameter names and values of all sources. If it is ``None``, it will be generated automatically from the ``fitparam_values`` argument using the :class:`~skyllh.core.parameters.ParameterModelMapper` instance. - tl : instance of TimeLord | None + tl The optional instance of TimeLord to measure the timing of evaluating the LLH ratio function. Returns ------- - log_lambda : float + log_lambda The calculated log-lambda value. - grads : instance of numpy ndarray + grads The (N_fitparams,)-shaped 1D numpy ndarray holding the gradient value for each global fit parameter. """ @@ -707,15 +736,16 @@ def evaluate(self, fitparam_values, src_params_recarray=None, tl=None): ns = fitparam_values[ns_pidx] N = tdm.n_events + assert N is not None # Calculate the data fields that depend on global fit parameters. if tdm.has_global_fitparam_data_fields: with TaskTimer(tl, 'Calculate global fit parameter dependent data fields.'): # Create the global_fitparams dictionary with the global fit # parameter names and values. - global_fitparams = self._pmm.get_global_floating_params_dict(gflp_values=fitparam_values) + global_fitparams = self._pmm.create_global_floating_params_dict(gflp_values=fitparam_values) tdm.calculate_global_fitparam_data_fields( - shg_mgr=self._shg_mgr, pmm=self._pmm, global_fitparams=global_fitparams + shg_mgr=self._shg_mgr, pmm=self._pmm, global_fitparams_dict=global_fitparams ) # Calculate the PDF ratio values for each selected event. @@ -772,25 +802,31 @@ class MultiDatasetTCLLHRatio(TCLLHRatio): """ def __init__( - self, pmm, minimizer, src_detsigyield_weights_service, ds_sig_weight_factors_service, llhratio_list, **kwargs + self, + pmm: ParameterModelMapper, + minimizer: Minimizer, + src_detsigyield_weights_service: SrcDetSigYieldWeightsService, + ds_sig_weight_factors_service: DatasetSignalWeightFactorsService, + llhratio_list: 'list[SingleDatasetTCLLHRatio]', + **kwargs, ): """Creates a new composite two-component log-likelihood ratio function. Parameters ---------- - pmm : instance of ParameterModelMapper + pmm The instance of ParameterModelMapper providing the mapping of global floating parameters to individual models. - minimizer : instance of Minimizer + minimizer The Minimizer instance that should be used to minimize the negative of this log-likelihood ratio function. - src_detsigyield_weights_service : instance of SrcDetSigYieldWeightsService + src_detsigyield_weights_service An instance of SrcDetSigYieldWeightsService, which provides the product of the source weights with the detector signal yield. - ds_sig_weight_factors_service : instance of DatasetSignalWeightFactorsService + ds_sig_weight_factors_service An instance of DatasetSignalWeightFactorsService, which provides the relative dataset signal weight factors. - llhratio_list : list of instance of SingleDatasetTCLLHRatio + llhratio_list The list of the two-component log-likelihood ratio functions, one for each dataset. """ @@ -877,18 +913,18 @@ def n_selected_events(self): @TCLLHRatio.mean_n_sig_0.setter def mean_n_sig_0(self, v): - TCLLHRatio.mean_n_sig_0.fset(self, v) + TCLLHRatio.mean_n_sig_0.fset(self, v) # pyright: ignore[reportOptionalCall] for llhratio in self._llhratio_list: llhratio.mean_n_sig_0 = self._mean_n_sig_0 - def change_shg_mgr(self, shg_mgr): + def change_shg_mgr(self, shg_mgr: SourceHypoGroupManager): """Changes the source hypo group manager of all objects of this LLH ratio function, hence, calling the ``change_shg_mgr`` method of all TCLLHRatio instances of this LLHRatio instance. Parameters ---------- - shg_mgr : instance of SourceHypoGroupManager + shg_mgr The instance of SourceHypoGroupManager that defines the groups of source hypotheses. """ @@ -897,7 +933,7 @@ def change_shg_mgr(self, shg_mgr): for llhratio in self._llhratio_list: llhratio.change_shg_mgr(shg_mgr=shg_mgr) - def initialize_for_new_trial(self, tl=None, **kwargs): + def initialize_for_new_trial(self, tl: TimeLord | None = None, **kwargs): """Initializes the log-likelihood-ratio function for a new trial. It calls the :meth:`~skyllh.core.llhratio.LLHRatio.initialize_for_new_trial` method @@ -906,22 +942,24 @@ def initialize_for_new_trial(self, tl=None, **kwargs): Parameters ---------- - tl : instance of TimeLord + tl The optional instance of TimeLord to measure timing information. """ for llhratio in self._llhratio_list: llhratio.initialize_for_new_trial(tl=tl, **kwargs) - def evaluate(self, fitparam_values, src_params_recarray=None, tl=None): + def evaluate( + self, fitparam_values: np.ndarray, src_params_recarray: np.ndarray | None = None, tl: TimeLord | None = None + ) -> tuple[float, np.ndarray]: """Evaluates the composite log-likelihood-ratio function and returns its value and global fit parameter gradients. Parameters ---------- - fitparam_values : instance of numpy ndarray + fitparam_values The (N_fitparams,)-shaped numpy 1D ndarray holding the current values of the global fit parameters. - src_params_recarray : instance of numpy record ndarray | None + src_params_recarray The numpy record ndarray of length N_sources holding the parameter names and values of all sources. See the documentation of the @@ -930,16 +968,16 @@ def evaluate(self, fitparam_values, src_params_recarray=None, tl=None): It case it is ``None``, it will be created automatically from the ``fitparam_values`` argument using the :class:`~skyllh.core.parameters.ParameterModelMapper` instance. - tl : instance of TimeLord | None + tl The optional instance of TimeLord that should be used for timing measurements. Returns ------- - log_lambda : float + log_lambda The calculated log-lambda value of the composite log-likelihood-ratio function. - grads : instance of numpy ndarray + grads The (N_fitparams,)-shaped 1D ndarray holding the gradient value of the composite log-likelihood-ratio function for each global fit parameter. @@ -967,6 +1005,8 @@ def evaluate(self, fitparam_values, src_params_recarray=None, tl=None): # f_grads is a dictionary holding (N_datasets,)-shaped 1D ndarrays for # each global fit parameter. (f, f_grads_dict) = self._ds_sig_weight_factors_service.get_weights() + assert f is not None + assert f_grads_dict is not None # Convert the f_grads dictionary into a (N_datasets,N_fitparams) f_grads = np.zeros((len(f), n_fitparams), dtype=np.float64) @@ -1016,7 +1056,9 @@ def evaluate(self, fitparam_values, src_params_recarray=None, tl=None): return (log_lambda, grads) - def calculate_ns_grad2(self, ns, ns_pidx, src_params_recarray, tl=None): + def calculate_ns_grad2( # pyright: ignore[reportIncompatibleMethodOverride] + self, ns: float, ns_pidx: int, src_params_recarray: np.ndarray, tl: TimeLord | None = None + ) -> float: """Calculates the second derivative w.r.t. ns of the log-likelihood ratio function. @@ -1028,30 +1070,28 @@ def calculate_ns_grad2(self, ns, ns_pidx, src_params_recarray, tl=None): Parameters ---------- - fitparam_values : instance of numpy ndarray - The (N_fitparams,)-shaped 1D ndarray holding the current values of - the global fit parameters. - ns : float + ns The value of the global fit parameter ns. - ns_pidx : int + ns_pidx The index of the global parameter ns. - src_params_recarray : instance of numpy record ndarray + src_params_recarray The numpy record ndarray of length N_sources holding the parameter names and values of all sources. See the documentation of the :meth:`skyllh.core.parameters.ParameterModelMapper.create_src_params_recarray` method for more information about this array. - tl : instance of TimeLord | None + tl The optional instance of TimeLord that should be used for timing measurements. Returns ------- - nsgrad2 : float + nsgrad2 The second derivative w.r.t. ns of the log-likelihood ratio function for the given fit parameter values. """ (f, _) = self._ds_sig_weight_factors_service.get_weights() + assert f is not None nsf = ns * f @@ -1082,7 +1122,14 @@ class NsProfileMultiDatasetTCLLHRatio(TCLLHRatio): the null-hypothesis. """ - def __init__(self, pmm, minimizer, mean_n_sig_0, llhratio, **kwargs): + def __init__( + self, + pmm: ParameterModelMapper, + minimizer: Minimizer, + mean_n_sig_0: float, + llhratio: 'MultiDatasetTCLLHRatio', + **kwargs, + ): r"""Creates a new ns-profile log-likelihood-ratio function with a null-hypothesis where :math:`n_{\mathrm{s}}` is fixed to ``mean_n_sig_0``. @@ -1136,7 +1183,7 @@ def llhratio(self, obj): ) self._llhratio = obj - def change_shg_mgr(self, shg_mgr): + def change_shg_mgr(self, shg_mgr: SourceHypoGroupManager): """Changes the source hypo group manager of all objects of this LLH ratio function, hence, calling the ``change_shg_mgr`` method of the underlying MultiDatasetTCLLHRatio instance of this @@ -1144,17 +1191,17 @@ def change_shg_mgr(self, shg_mgr): Parameters ---------- - shg_mgr : instance of SourceHypoGroupManager + shg_mgr The new instance of SourceHypoGroupManager. """ self._llhratio.change_shg_mgr(shg_mgr=shg_mgr) - def initialize_for_new_trial(self, tl=None, **kwargs): + def initialize_for_new_trial(self, tl: TimeLord | None = None, **kwargs): """Initializes the log-likelihood-ratio function for a new trial. Parameters ---------- - tl : instance of TimeLord | None + tl The optional instance of TimeLord that should be used for timing measurements. """ @@ -1165,33 +1212,35 @@ def initialize_for_new_trial(self, tl=None, **kwargs): fitparam_values_0 = np.array([self._mean_n_sig_0], dtype=np.float64) (self._logL_0, _) = self._llhratio.evaluate(fitparam_values=fitparam_values_0, tl=tl) - def evaluate(self, fitparam_values, src_params_recarray=None, tl=None): + def evaluate( + self, fitparam_values: np.ndarray, src_params_recarray: np.ndarray | None = None, tl: TimeLord | None = None + ) -> tuple[float, np.ndarray]: """Evaluates the log-likelihood-ratio function and returns its value and global fit parameter gradients. Parameters ---------- - fitparam_values : instance of numpy ndarray + fitparam_values The (1,)-shaped numpy 1D ndarray holding the current values of the global fit parameters. By definition of this LLH ratio function, it must contain the single fit parameter value for ns. - src_params_recarray : instance of numpy record ndarray + src_params_recarray The numpy record ndarray of length N_sources holding the parameter names and values of all sources. See the documentation of the :meth:`skyllh.core.parameters.ParameterModelMapper.create_src_params_recarray` method for more information about this array. - tl : instance of TimeLord | None + tl The optional instance of TimeLord that should be used for timing measurements. Returns ------- - log_lambda : float + log_lambda The calculated log-lambda value of this log-likelihood-ratio function. - grads : (1,)-shaped 1D ndarray + grads The ndarray holding the gradient value of this log-likelihood-ratio for ns. """ @@ -1199,34 +1248,37 @@ def evaluate(self, fitparam_values, src_params_recarray=None, tl=None): fitparam_values=fitparam_values, src_params_recarray=src_params_recarray, tl=tl ) + assert self._logL_0 is not None log_lambda = logL - self._logL_0 return (log_lambda, grads) - def calculate_ns_grad2(self, ns, ns_pidx, src_params_recarray, tl=None): + def calculate_ns_grad2( # pyright: ignore[reportIncompatibleMethodOverride] + self, ns: float, ns_pidx: int, src_params_recarray: np.ndarray, tl: TimeLord | None = None + ) -> float: """Calculates the second derivative w.r.t. ns of the log-likelihood ratio function. Parameters ---------- - ns : float + ns The value of the global fit parameter ns. - ns_pidx : int + ns_pidx The index of the global fit parameter ns. By definition this must be ``0``. - src_params_recarray : instance of numpy record ndarray + src_params_recarray The numpy record ndarray of length N_sources holding the parameter names and values of all sources. See the documentation of the :meth:`skyllh.core.parameters.ParameterModelMapper.create_src_params_recarray` method for more information about this array. - tl : instance of TimeLord | None + tl The optional instance of TimeLord that should be used for timing measurements. Returns ------- - nsgrad2 : float + nsgrad2 The second derivative w.r.t. ns of the log-likelihood ratio function for the given fit parameter values. """ diff --git a/skyllh/core/logging.py b/skyllh/core/logging.py index 78fdfbb8f8..5d5582082f 100644 --- a/skyllh/core/logging.py +++ b/skyllh/core/logging.py @@ -2,15 +2,17 @@ import os.path import sys +from skyllh.core.config import Config -def _resolve_log_level(level, default=logging.INFO): + +def _resolve_log_level(level: int | str | None, default: int = logging.INFO): """Converts a logging level representation into a numeric level. Parameters ---------- - level : int | str | None + level The level representation. - default : int + default Default level if ``level`` is None. Returns @@ -36,18 +38,18 @@ def _resolve_log_level(level, default=logging.INFO): raise TypeError(f'The logging level must be int, str, or None! Its current type is {type(level)}.') -def get_logger(name): +def get_logger(name: str) -> logging.Logger: """Retrieves the logger with the given name from the Python logging system. Parameters ---------- - name : str + name The name of the logger. Logger hierarchy is defined using dots as separators. Returns ------- - logger : logging.Logger + logger The Logger instance. """ logger = logging.getLogger(name) @@ -55,54 +57,54 @@ def get_logger(name): def setup_logger( - cfg, - name, - log_level=None, - log_format=None, - console=False, - console_level=None, + cfg: Config, + name: str, + log_level: int | str | None = None, + log_format: str | None = None, + console: bool = False, + console_level: int | None = None, stream=None, - log_file=None, - file_level=None, - file_mode='a', - propagate=False, - clear_existing_handlers=False, + log_file: str | None = None, + file_level: int | None = None, + file_mode: str = 'a', + propagate: bool = False, + clear_existing_handlers: bool = False, ): """Sets up a logger with the given local configuration and a name. Parameters ---------- - cfg : instance of Config + cfg Local configuration. - name : str + name The name of the logger to set up. Logger hierarchy is defined using dots as separators. - log_level : int | str | None + log_level The log level of the logger. If None, the log level is taken from the configuration. - log_format : str | None + log_format The format of log records in the final output. If None, the log format is taken from the configuration. - console : bool + console Whether to set up a console handler for the logger. Default: False. - console_level : int | None + console_level The log level of the console handler. If None, it uses `log_level`. - stream : data stream | None + stream The stream to which the console handler will write. If None, it defaults to `sys.stdout`. - log_file : str | None + log_file If not ``None``, file handlers for DEBUG messages will be installed and those messages will be stored in the given file. - file_level : int | None + file_level The log level of the file handler. If None, it uses `log_level`. - file_mode : str + file_mode File opening mode. Default is 'a' for appending. - propagate : bool + propagate Whether the logger should propagate messages to ancestor loggers. - Default: False. - clear_existing_handlers : bool + Default + clear_existing_handlers Optionally clear handlers before setting up new ones. - Default: False. + Default """ logger = logging.getLogger(name) @@ -166,31 +168,39 @@ def setup_logger( return logger -def setup_logging(cfg, name, log_format=None, log_level=None, console=True, log_file=None, reconfigure=False): +def setup_logging( + cfg: Config, + name: str, + log_format: str | None = None, + log_level: int | str | None = None, + console: bool = True, + log_file: str | None = None, + reconfigure: bool = False, +): """Initializes package and script loggers and returns the script logger. Parameters ---------- - cfg : instance of Config + cfg The instance of Config holding the local configuration. - name : str + name The name of the user-defined logger to set up. - log_format : str | None + log_format The format template of the log message. If ``None``, the format is taken from ``cfg['logging']['log_format']``. - log_level : int | str | None + log_level The log level of the loggers. If ``None``, it is taken from the configuration. - console : bool + console Whether to set up console handlers for the loggers. Default: True. - log_file : str | None + log_file If not ``None``, a file handler for log messages will be installed for both loggers using this path. - reconfigure : bool + reconfigure Rebuild logging setup from scratch for this run/session. Especially useful in interactive environments like Jupyter notebooks to avoid duplicate log messages due to multiple logging handlers. - Default: False. + Default Returns ------- diff --git a/skyllh/core/math.py b/skyllh/core/math.py index 3fd7df6465..17fdb39df6 100644 --- a/skyllh/core/math.py +++ b/skyllh/core/math.py @@ -23,6 +23,7 @@ class properties. The tuple of parameter names is defined through the """ def __init__(self, **kwargs): + """Creates a new instance of MathFunction.""" super().__init__(**kwargs) self.param_names = () @@ -50,20 +51,20 @@ def param_names(self, names): @property @abc.abstractmethod - def math_function_str(self): + def math_function_str(self) -> 'str | None': """The string showing the mathematical function of this MathFunction.""" - def __str__(self): + def __str__(self) -> str: """Pretty string representation of this MathFunction instance.""" - return self.math_function_str + return self.math_function_str or '' - def copy(self, newparams=None): + def copy(self, newparams: dict | None = None): """Copies this MathFunction object by calling the copy.deepcopy function, and sets new parameters if requested. Parameters ---------- - newparams : dict | None + newparams The dictionary with the new parameter values to set, where the dictionary key is the parameter name and the dictionary value is the new value of the parameter. @@ -76,18 +77,18 @@ def copy(self, newparams=None): return f - def get_param(self, name): + def get_param(self, name: str) -> float: """Retrieves the value of the given parameter. It returns ``np.nan`` if the parameter does not exist. Parameters ---------- - name : str + name The name of the parameter. Returns ------- - value : float | np.nan + value The value of the parameter. """ if name not in self._param_names: @@ -97,19 +98,19 @@ def get_param(self, name): return value - def set_params(self, pdict): + def set_params(self, pdict: dict) -> bool: """Sets the parameters of the math function to the given parameter values. Parameters ---------- - pdict : dict (name: value) + pdict The dictionary holding the names of the parameters and their new values. Returns ------- - updated : bool + updated Flag if parameter values were actually updated. """ if not isinstance(pdict, dict): diff --git a/skyllh/core/minimizer.py b/skyllh/core/minimizer.py index 6b1866cd72..30b54ba74d 100644 --- a/skyllh/core/minimizer.py +++ b/skyllh/core/minimizer.py @@ -4,6 +4,7 @@ """ import abc +from collections.abc import Callable, Sequence import numpy as np import scipy.optimize @@ -20,6 +21,7 @@ from skyllh.core.py import ( classname, ) +from skyllh.core.random import RandomStateService logger = get_logger(__name__) @@ -33,35 +35,36 @@ class MinimizerImpl( """ def __init__(self, **kwargs): + """Creates a new instance of MinimizerImpl.""" super().__init__(**kwargs) @abc.abstractmethod def minimize( self, - initials, - bounds, - func, - func_args=None, + initials: np.ndarray, + bounds: np.ndarray, + func: Callable, + func_args: Sequence | None = None, **kwargs, - ): + ) -> tuple[np.ndarray, float, dict]: """This method is supposed to minimize the given function with the given initials. Parameters ---------- - initials : 1D (N_fitparams)-shaped numpy ndarray + initials The ndarray holding the initial values of all the fit parameters. - bounds : 2D (N_fitparams,2)-shaped numpy ndarray + bounds The ndarray holding the boundary values (vmin, vmax) of the fit parameters. - func : callable + func The function that should get minimized. The call signature must be ``__call__(x, *args)`` The return value of ``func`` is minimizer implementation dependent. - func_args : sequence | None + func_args Optional sequence of arguments for ``func``. Additional Keyword Arguments @@ -71,65 +74,65 @@ def minimize( Returns ------- - xmin : 1D ndarray + xmin The array containing the function parameter values at the function's minimum. - fmin : float + fmin The function value at its minimum. - status : dict + status The status dictionary with information about the minimization process. """ @abc.abstractmethod - def get_niter(self, status): + def get_niter(self, status: dict) -> int: """This method is supposed to return the number of iterations that were required to find the minimum. Parameters ---------- - status : dict + status The dictionary with the status information about the last minimization process. Returns ------- - niter : int + niter The number of iterations needed to find the minimum. """ @abc.abstractmethod - def has_converged(self, status): + def has_converged(self, status: dict) -> bool: """This method is supposed to analyze the status information dictionary if the last minimization process has converged. Parameters ---------- - status : dict + status The dictionary with the status information about the last minimization process. Returns ------- - converged : bool + converged The flag if the minimization has converged (True), or not (False). """ @abc.abstractmethod - def is_repeatable(self, status): + def is_repeatable(self, status: dict) -> bool: """This method is supposed to analyze the status information dictionary if the last minimization process can be repeated to obtain a better minimum. Parameters ---------- - status : dict + status The dictionary with the status information about the last minimization process. Returns ------- - repeatable : bool + repeatable The flag if the minimization process can be repeated to obtain a better minimum. """ @@ -147,7 +150,7 @@ def __init__( Parameters ---------- - method : str + method The minimizer method to use. See the documentation for the method argument of the :func:`scipy.optimize.minimize` function for possible values. @@ -158,23 +161,23 @@ def __init__( def minimize( self, - initials, - bounds, - func, - func_args=None, + initials: np.ndarray, + bounds: np.ndarray | None, + func: Callable, + func_args: Sequence | None = None, **kwargs, - ): + ) -> tuple[np.ndarray, float, dict]: """Minimizes the given function ``func`` with the given initial function argument values ``initials``. Parameters ---------- - initials : 1D numpy ndarray + initials The ndarray holding the initial values of all the fit parameters. - bounds : 2D (N_fitparams,2)-shaped numpy ndarray + bounds The ndarray holding the boundary values (vmin, vmax) of the fit parameters. - func : callable + func The function that should get minimized. The call signature must be @@ -185,7 +188,7 @@ def minimize( the function gradient for each fit parameter, if the ``func_provides_grads`` keyword argument option is set to True. If set to False, ``func`` must return only the function value. - func_args : sequence | None + func_args Optional sequence of arguments for ``func``. Additional Keyword Arguments @@ -204,12 +207,12 @@ def minimize( Returns ------- - xmin : instance of numpy.ndarray + xmin The 1D array containing the function arguments at the function's minimum. - fmin : float + fmin The function value at its minimum. - res : instance of scipy.optimize.OptimizeResult + res The scipy OptimizeResult. """ @@ -224,14 +227,14 @@ def minimize( elif self._method == 'COBYLA': # COBYLA doesn't allow for bounds, but we can convert bounds # to a linear constraint - - constraints = [] - for bound_num, bound in enumerate(bounds): - lower, upper = bound - lc = {'type': 'ineq', 'fun': lambda x, lb=lower, i=bound_num: x[i] - lb} - uc = {'type': 'ineq', 'fun': lambda x, ub=upper, i=bound_num: ub - x[i]} - constraints.append(lc) - constraints.append(uc) + if bounds is not None: + constraints = [] + for bound_num, bound in enumerate(bounds): + lower, upper = bound + lc = {'type': 'ineq', 'fun': lambda x, lb=lower, i=bound_num: x[i] - lb} + uc = {'type': 'ineq', 'fun': lambda x, ub=upper, i=bound_num: ub - x[i]} + constraints.append(lc) + constraints.append(uc) bounds = None if (bounds is not None) and (not method_supports_bounds): @@ -253,41 +256,41 @@ def minimize( return (res.x, res.fun, res) - def get_niter(self, status): + def get_niter(self, status: dict) -> int: """Returns the number of iterations needed to find the minimum. Parameters ---------- - status : dict + status The dictionary with the status information about the minimization process. Returns ------- - niter : int + niter The number of iterations needed to find the minimum. """ return status['nit'] - def has_converged(self, status): + def has_converged(self, status: dict) -> bool: """Analyzes the status information dictionary if the minimization process has converged. By definition the minimization process has converged if ``status['warnflag']`` equals 0. Parameters ---------- - status : dict + status The dictionary with the status information about the minimization process. Returns ------- - converged : bool + converged The flag if the minimization has converged (True), or not (False). """ return bool(status['success']) - def is_repeatable(self, status): + def is_repeatable(self, status: dict) -> bool: """Checks if the minimization process can be repeated to get a better result. @@ -296,13 +299,13 @@ def is_repeatable(self, status): Parameters ---------- - status : dict + status The dictionary with the status information about the last minimization process. Returns ------- - repeatable : bool + repeatable The flag if the minimization process can be repeated to obtain a better minimum. """ @@ -316,9 +319,9 @@ class LBFGSMinimizerImpl(MinimizerImpl): def __init__( self, - ftol=1e-6, - pgtol=1e-5, - maxls=100, + ftol: float = 1e-6, + pgtol: float = 1e-5, + maxls: int = 100, **kwargs, ): """Creates a new L-BGF-S minimizer instance to minimize the given @@ -326,11 +329,11 @@ def __init__( Parameters ---------- - ftol : float + ftol The function value tolerance. - pgtol : float + pgtol The gradient value tolerance. - maxls : int + maxls The maximum number of line search steps for an iteration. """ super().__init__(**kwargs) @@ -343,23 +346,23 @@ def __init__( def minimize( self, - initials, - bounds, - func, - func_args=None, + initials: np.ndarray, + bounds: np.ndarray, + func: Callable, + func_args: Sequence | None = None, **kwargs, - ): + ) -> tuple[np.ndarray, float, dict]: """Minimizes the given function ``func`` with the given initial function argument values ``initials``. Parameters ---------- - initials : 1D numpy ndarray + initials The ndarray holding the initial values of all the fit parameters. - bounds : 2D (N_fitparams,2)-shaped numpy ndarray + bounds The ndarray holding the boundary values (vmin, vmax) of the fit parameters. - func : callable + func The function that should get minimized. The call signature must be @@ -370,7 +373,7 @@ def minimize( the function gradient for each fit parameter, if the ``func_provides_grads`` keyword argument option is set to True. If set to False, ``func`` must return only the function value. - func_args : sequence | None + func_args Optional sequence of arguments for ``func``. Additional Keyword Arguments @@ -388,22 +391,22 @@ def minimize( Returns ------- - xmin : 1D ndarray + xmin The array containing the function arguments at the function's minimum. - fmin : float + fmin The function value at its minimum. - status : dict + status The status dictionary with information about the minimization process. The following information are provided: - niter : int + niter The number of iterations needed to find the minimum. - warnflag : int + warnflag The warning flag indicating if the minimization did converge. The possible values are: - 0: The minimization converged. + 0 """ if func_args is None: func_args = () @@ -430,48 +433,48 @@ def minimize( def get_niter( self, - status, - ): + status: dict, + ) -> int: """Returns the number of iterations needed to find the minimum. Parameters ---------- - status : dict + status The dictionary with the status information about the minimization process. Returns ------- - niter : int + niter The number of iterations needed to find the minimum. """ return status['nit'] def has_converged( self, - status, - ): + status: dict, + ) -> bool: """Analyzes the status information dictionary if the minimization process has converged. By definition the minimization process has converged if ``status['warnflag']`` equals 0. Parameters ---------- - status : dict + status The dictionary with the status information about the minimization process. Returns ------- - converged : bool + converged The flag if the minimization has converged (True), or not (False). """ return status['warnflag'] == 0 def is_repeatable( self, - status, - ): + status: dict, + ) -> bool: """Checks if the minimization process can be repeated to get a better result. It's repeatable if @@ -479,13 +482,13 @@ def is_repeatable( Parameters ---------- - status : dict + status The dictionary with the status information about the last minimization process. Returns ------- - repeatable : bool + repeatable The flag if the minimization process can be repeated to obtain a better minimum. """ @@ -509,8 +512,8 @@ class NR1dNsMinimizerImpl(MinimizerImpl): def __init__( self, - ns_tol=1e-3, - max_steps=100, + ns_tol: float = 1e-3, + max_steps: int = 100, **kwargs, ): """Creates a new NRNs minimizer instance to minimize the given @@ -518,9 +521,9 @@ def __init__( Parameters ---------- - ns_tol : float + ns_tol The tolerance / precision for the ns parameter value. - max_steps : int + max_steps The maximum number of NR steps. If max_step is reached, the fit is considered NOT converged. """ @@ -531,12 +534,12 @@ def __init__( def minimize( self, - initials, - bounds, - func, - func_args=None, + initials: np.ndarray, + bounds: np.ndarray, + func: Callable, + func_args: Sequence | None = None, **kwargs, - ): + ) -> tuple[np.ndarray, float, dict]: """Minimizes the given function ``func`` with the given initial function argument values ``initials``. This minimizer implementation will only vary the first parameter. All other parameters will be set to their @@ -544,12 +547,12 @@ def minimize( Parameters ---------- - initials : 1D numpy ndarray + initials The ndarray holding the initial values of all the fit parameters. - bounds : 2D (N_fitparams,2)-shaped numpy ndarray + bounds The ndarray holding the boundary values (vmin, vmax) of the fit parameters. - func : callable + func The function that should get minimized. The call signature must be @@ -559,7 +562,7 @@ def minimize( function value at the function arguments ``x``, the value of the function first derivative for the one fit parameter, and the value of the second derivative for the one fit parameter. - func_args : sequence | None + func_args Optional sequence of arguments for ``func``. Additional Keyword Arguments @@ -569,20 +572,20 @@ def minimize( Returns ------- - xmin : 1D ndarray + xmin The array containing the function parameter values at the function's minimum. - fmin : float + fmin The function value at its minimum. - status : dict + status The status dictionary with information about the minimization process. The following information are provided: - niter : int + niter The number of iterations needed to find the minimum. - last_nr_step : float + last_nr_step The Newton-Raphson step size of the last iteration. - warnflag : int + warnflag The warning flag indicating if the minimization did converge. The possible values are: @@ -590,12 +593,12 @@ def minimize( parameter value. Convergence forced at upper bound. -2: The function minimum is below the lower bound of the parameter value. Convergence forced at lower bound. - 0: The minimization converged with a iteration step size + 0 smaller than the specified precision. - 1: The minimization did NOT converge within self.max_steps + 1 number of steps - warnreason: str + warnreason The description for the set warn flag. """ @@ -682,44 +685,45 @@ def minimize( status['niter'] = niter status['last_nr_step'] = step + assert f is not None return (x, f, status) def get_niter( self, - status, - ): + status: dict, + ) -> int: """Returns the number of iterations needed to find the minimum. Parameters ---------- - status : dict + status The dictionary with the status information about the minimization process. Returns ------- - niter : int + niter The number of iterations needed to find the minimum. """ return status['niter'] def has_converged( self, - status, - ): + status: dict, + ) -> bool: """Analyzes the status information dictionary if the minimization process has converged. By definition the minimization process has converged if ``status['warnflag']`` is smaller or equal to 0. Parameters ---------- - status : dict + status The dictionary with the status information about the minimization process. Returns ------- - converged : bool + converged The flag if the minimization has converged (True), or not (False). """ return status['warnflag'] <= 0 @@ -740,18 +744,18 @@ class NRNsScan2dMinimizerImpl(NR1dNsMinimizerImpl): def __init__( self, - p2_scan_step, - ns_tol=1e-3, + p2_scan_step: float, + ns_tol: float = 1e-3, **kwargs, ): """Creates a new minimizer implementation instance. Parameters ---------- - p2_scan_step : float + p2_scan_step The step size for the scan of the second parameter of the function to minimize. - ns_tol : float + ns_tol The tolerance / precision for the ns parameter value. """ super().__init__(ns_tol=ns_tol, **kwargs) @@ -760,12 +764,12 @@ def __init__( def minimize( self, - initials, - bounds, - func, - func_args=None, + initials: np.ndarray, + bounds: np.ndarray, + func: Callable, + func_args: Sequence | None = None, **kwargs, - ): + ) -> tuple[np.ndarray, float, dict]: """Minimizes the given function ``func`` with the given initial function argument values ``initials``. This minimizer implementation will only vary the first two parameters. The first parameter is the number of @@ -775,12 +779,12 @@ def minimize( Parameters ---------- - initials : 1D numpy ndarray + initials The ndarray holding the initial values of all the fit parameters. - bounds : 2D (N_fitparams,2)-shaped numpy ndarray + bounds The ndarray holding the boundary values (vmin, vmax) of the fit parameters. - func : callable + func The function that should get minimized. The call signature must be @@ -790,7 +794,7 @@ def minimize( function value at the function arguments ``x``, the value of the function first derivative for the first fit parameter, and the value of the second derivative for the first fit parameter. - func_args : sequence | None + func_args Optional sequence of arguments for ``func``. Additional Keyword Arguments @@ -800,34 +804,34 @@ def minimize( Returns ------- - xmin : 1D ndarray + xmin The array containing the function parameter values at the function's minimum. - fmin : float + fmin The function value at its minimum. - status : dict + status The status dictionary with information about the minimization process. The following information are provided: - niter : int + niter The number of iterations needed to find the minimum. - last_nr_step : float + last_nr_step The Newton-Raphson step size of the last iteration. - p2_n_steps : int + p2_n_steps The number of scanning steps performed for the 2nd parameter. - warnflag : int + warnflag The warning flag indicating if the minimization did converge. The possible values are: - 0: The minimization converged with a iteration step size + 0 smaller than the specified precision. - 1: The function minimum is below the minimum bound of the + 1 parameter value. The last iteration's step size did not achieve the specified precision. - 2: The function minimum is above the maximum bound of the + 2 parameter value. The last iteration's step size did not achieve the specified precision. - warnreason: str + warnreason The description for the set warn flag. """ p2_low = bounds[1][0] @@ -853,6 +857,9 @@ def minimize( best_fmin = fmin best_status = status + assert best_xmin is not None + assert best_fmin is not None + assert best_status is not None best_status['p2_n_steps'] = len(p2_scan_values) best_status['niter'] = niter_total @@ -867,17 +874,17 @@ class Minimizer: def __init__( self, - minimizer_impl, - max_repetitions=100, + minimizer_impl: 'MinimizerImpl', + max_repetitions: int = 100, **kwargs, ): """Creates a new Minimizer instance. Parameters ---------- - minimizer_impl : instance of MinimizerImpl + minimizer_impl The minimizer implementation for a specific minimizer algorithm. - max_repetitions : int + max_repetitions In case the minimization process did not converge at the first time this option specifies the maximum number of repetitions with different initials. @@ -916,12 +923,12 @@ def max_repetitions(self, n): def minimize( self, - rss, - paramset, - func, + rss: RandomStateService, + paramset: ParameterSet, + func: Callable, args=None, - kwargs=None, - ): + kwargs: dict | None = None, + ) -> tuple[np.ndarray, float, dict]: """Minimizes the the given function ``func`` by calling the ``minimize`` method of the minimizer implementation. @@ -933,31 +940,31 @@ def minimize( Parameters ---------- - rss : RandomStateService instance + rss The RandomStateService instance to draw random numbers from. - paramset : instance of ParameterSet + paramset The ParameterSet instances holding the floating parameters of the function ``func``. - func : callable ``f(x, *args)`` + func The function to be minimized. It must have the call signature ``__call__(x, *args)`` The return value of ``func`` is minimizer implementation dependent. - args : sequence of arguments for ``func`` | None + args The optional sequence of arguments for ``func``. - kwargs : dict | None + kwargs The optional dictionary with keyword arguments for the minimizer implementation minimize method. Returns ------- - xmin : 1d numpy ndarray + xmin The array holding the parameter values for which the function has a minimum. - fmin : float + fmin The function value at its minimum. - status : dict + status The status dictionary with information about the minimization process. """ diff --git a/skyllh/core/minimizers/iminuit.py b/skyllh/core/minimizers/iminuit.py index 71906042e7..a433635a23 100644 --- a/skyllh/core/minimizers/iminuit.py +++ b/skyllh/core/minimizers/iminuit.py @@ -3,6 +3,9 @@ minimizer. """ +from collections.abc import Callable, Sequence +from typing import Any, cast + import numpy as np from skyllh.core import ( @@ -31,21 +34,21 @@ class FuncWithGradsFunctor( def __init__( self, - func, - func_args=None, + func: Callable, + func_args: Sequence | None = None, **kwargs, ): """Initializes a new functor instance for the given function ``func``. Parameters ---------- - func : callable + func The function with call signature ``__call__(x, *args)`` returning the a two-element tuple (f, grads). - func_args : tuple | None + func_args The optional positional arguments for the function ``func``. """ super().__init__(**kwargs) @@ -63,6 +66,9 @@ def __init__( self._cache_grads = None def get_f(self, x): + """Returns the function value at the given parameter values ``x``, using + a cached value if ``x`` matches the last evaluation. + """ tracing = self._tracing if self._cache_x is None: @@ -82,6 +88,9 @@ def get_f(self, x): return self._cache_f def get_grads(self, x): + """Returns the function gradients at the given parameter values ``x``, + using cached values if ``x`` matches the last evaluation. + """ tracing = self._tracing if self._cache_x is None: @@ -109,7 +118,7 @@ class IMinuitMinimizerImpl( @tool.requires('iminuit') def __init__( self, - ftol=1e-6, + ftol: float = 1e-6, **kwargs, ): """Creates a new IMinuit minimizer instance to minimize a given @@ -117,7 +126,7 @@ def __init__( Parameters ---------- - ftol : float + ftol The function value tolerance as absolute value. """ super().__init__(**kwargs) @@ -126,24 +135,24 @@ def __init__( def minimize( self, - initials, - bounds, - func, - func_args=None, + initials: np.ndarray, + bounds: np.ndarray, + func: Callable, + func_args: Sequence | None = None, **kwargs, - ): + ) -> tuple[np.ndarray, float, dict]: """Minimizes the given function ``func`` with the given initial function argument values ``initials`` and within the given parameter bounds ``bounds``. Parameters ---------- - initials : 1D numpy ndarray + initials The ndarray holding the initial values of all the fit parameters. - bounds : 2D (N_fitparams,2)-shaped numpy ndarray + bounds The ndarray holding the boundary values (vmin, vmax) of the fit parameters. - func : callable + func The function that should get minimized. The call signature must be @@ -154,7 +163,7 @@ def minimize( the function gradient for each fit parameter, if the ``func_provides_grads`` keyword argument option is set to True. If set to False, ``func`` must return only the function value. - func_args : sequence | None + func_args Optional sequence of arguments for ``func``. Additional Keyword Arguments @@ -172,12 +181,12 @@ def minimize( Returns ------- - xmin : 1D ndarray + xmin The array containing the function arguments at the function's minimum. - fmin : float + fmin The function value at its minimum. - res : iminuit.OptimizeResult + res The iminuit OptimizeResult dictionary with additional information. """ if func_args is None: @@ -185,7 +194,7 @@ def minimize( if kwargs is None: kwargs = {} - iminuit = tool.get('iminuit') + iminuit = cast(Any, tool.get('iminuit')) func_provides_grads = kwargs.pop('func_provides_grads', True) @@ -204,41 +213,41 @@ def minimize( return (res.x, res.fun, res) - def get_niter(self, status): + def get_niter(self, status: dict) -> int: """Returns the number of iterations needed to find the minimum. Parameters ---------- - status : dict + status The dictionary with the status information about the minimization process. Returns ------- - niter : int + niter The number of iterations needed to find the minimum. """ return status['nfev'] - def has_converged(self, status): + def has_converged(self, status: dict) -> bool: """Analyzes the status information dictionary if the minimization process has converged. By definition the minimization process has converged if ``status['is_valid']`` equals True. Parameters ---------- - status : dict + status The dictionary with the status information about the minimization process. Returns ------- - converged : bool + converged The flag if the minimization has converged (True), or not (False). """ return bool(status['success']) - def is_repeatable(self, status): + def is_repeatable(self, status: dict) -> bool: """Checks if the minimization process can be repeated to get a better result. @@ -247,13 +256,13 @@ def is_repeatable(self, status): Parameters ---------- - status : dict + status The dictionary with the status information about the last minimization process. Returns ------- - repeatable : bool + repeatable The flag if the minimization process can be repeated to obtain a better minimum. """ diff --git a/skyllh/core/model.py b/skyllh/core/model.py index 0f98f1f950..535a0ad862 100644 --- a/skyllh/core/model.py +++ b/skyllh/core/model.py @@ -1,5 +1,7 @@ """This module defines the base class for any model class used in SkyLLH.""" +from collections.abc import Sequence + from skyllh.core.py import ( NamedObjectCollection, issequenceof, @@ -13,19 +15,19 @@ class Model: Models could be for instance source models or background models. """ - def __init__(self, name=None, **kwargs): + def __init__(self, name: str | None = None, **kwargs): """Creates a new Model instance. Parameters ---------- - name : str | None + name The name of the model. If set to `None`, the id of the object is taken as name. """ super().__init__(**kwargs) if name is None: - name = self.id + name = str(self.id) self.name = name @@ -54,17 +56,18 @@ class ModelCollection(NamedObjectCollection): """ @staticmethod - def cast(obj, errmsg=None, **kwargs): + def cast( + obj: 'Model | ModelCollection | Sequence[Model] | None', errmsg: str | None = None, **kwargs + ) -> 'ModelCollection': """Casts the given object to a ModelCollection object. If the cast fails, a TypeError with the given error message is raised. Parameters ---------- - obj : Model instance | sequence of Model instances | - ModelCollection | None + obj The object that should be casted to ModelCollection. If set to None, an empty ModelCollection is created. - errmsg : str | None + errmsg The error message if the cast fails. If set to None, a generic error message will be used. @@ -80,7 +83,7 @@ def cast(obj, errmsg=None, **kwargs): Returns ------- - model_collection : instance of ModelCollection + model_collection The created ModelCollection instance. If `obj` is already a ModelCollection instance, it will be returned. """ @@ -100,15 +103,15 @@ def cast(obj, errmsg=None, **kwargs): errmsg = f'Cast of object "{obj!s}" of type "{typename(type(obj))}" to ModelCollection failed!' raise TypeError(errmsg) - def __init__(self, models=None, model_type=None, **kwargs): + def __init__(self, models=None, model_type: type | None = None, **kwargs): """Creates a new Model collection. The type of the model instances this collection holds can be restricted, by setting the model_type argument. Parameters ---------- - models : sequence of model_type instances | None + models The sequence of models this collection should be initalized with. - model_type : type | None + model_type The type of the model. It must be a subclass of class ``Model``. If set to None (default), Model will be used. """ @@ -135,12 +138,12 @@ class DetectorModel(Model): in combination with the ParameterModelMapper class. """ - def __init__(self, name, **kwargs): + def __init__(self, name: str, **kwargs): """Creates a new DetectorModel instance. Parameters ---------- - name : str + name The name of the detector model. """ super().__init__(name=name, **kwargs) diff --git a/skyllh/core/multiproc.py b/skyllh/core/multiproc.py index c93752d37f..dcc01a40b4 100644 --- a/skyllh/core/multiproc.py +++ b/skyllh/core/multiproc.py @@ -1,15 +1,15 @@ import multiprocessing as mp import queue import time +from collections.abc import Callable from logging.handlers import ( QueueHandler, ) +from typing import cast import numpy as np -from skyllh.core.config import ( - HasConfig, -) +from skyllh.core.config import Config, HasConfig from skyllh.core.logging import ( get_logger, ) @@ -28,22 +28,22 @@ def get_ncpu( - cfg, - local_ncpu, -): + cfg: Config, + local_ncpu: int | None, +) -> int: """Determines the number of CPUs to use for functions that support multi-processing. Parameters ---------- - cfg : instance of Config + cfg The instance of Config holding the local configuration. - local_ncpu : int | None + local_ncpu The local setting of the number of CPUs to use. Returns ------- - ncpu : int + ncpu The number of CPUs to use by functions that allow multi-processing. If ``local_ncpu`` is set to None, the global NCPU setting is returned. If the global NCPU setting is None as well, the default value 1 is @@ -65,39 +65,39 @@ def get_ncpu( def parallelize( - func, - args_list, - ncpu, - rss=None, - tl=None, - ppbar=None, -): + func: Callable, + args_list: list[tuple], + ncpu: int, + rss: RandomStateService | None = None, + tl: TimeLord | None = None, + ppbar: ProgressBar | None = None, +) -> list: """Parallelizes the execution of the given function for different arguments. Parameters ---------- - func : callable + func The function which should be called with different arguments, which are given through the args_list argument. If the `rss` argument is not None, `func` requires an argument named `rss`. - args_list : list of 2-element tuple + args_list The list of the different arguments for function ``func``. Each element of that list must be a 2-element tuple, where the first element is a tuple of the arguments of ``func``, and the second element is a dictionary with the keyword arguments of ``func``. If the `rss` argument is not None, `func` argument `rss` has to be omitted. - ncpu : int + ncpu The number of CPUs to use, i.e. the number of subprocesses to spawn. - rss : RandomStateService | None + rss The RandomStateService instance to use for generating random numbers. - tl : instance of TimeLord | None + tl The instance of TimeLord that should be used to time individual tasks. - ppbar : instance of ProgressBar | None + ppbar The possible parent ProgressBar instance. Returns ------- - result_list : list + result_list The list of the result values of ``func``, where each element of that list corresponds to the arguments element in ``args_list``. """ @@ -105,14 +105,14 @@ def parallelize( # Define a wrapper function for the multiprocessing module that evaluates # ``func`` for a subset of `args_list` on a worker process. def worker_wrapper( - func, - sub_args_list, - pid, + func: Callable, + sub_args_list: list[tuple], + pid: int, rqueue, lqueue, squeue=None, - rss=None, - tl=None, + rss: RandomStateService | None = None, + tl: TimeLord | None = None, ): """Wrapper function for the multiprocessing module that evaluates ``func`` for the subset ``sub_args_list`` of ``args_list`` on a worker @@ -120,38 +120,38 @@ def worker_wrapper( Parameters ---------- - func : callable + func The function which should be called with different arguments, which are given through the sub_args_list argument. If the `rss` argument is not None `func` requires an argument named `rss`. - sub_args_list : list of 2-element tuple + sub_args_list The list of the different arguments for function ``func``. Each element of that list must be a 2-element tuple, where the first element is a tuple of the arguments of ``func``, and the second element is a dictionary with the keyword arguments of ``func``. If the `rss` argument is not None, `func` argument `rss` has to be omitted. - pid : int + pid The process ID that identifies the process in order to sort the results to the initial order of the function arguments. - rqueue : multiprocessing.Queue + rqueue The Queue instance where to put the function result in. If set to None, the result list will be returned. - lqueue : multiprocessing.Queue + lqueue The queue to hold generated log records by a given function. - squeue : multiprocessing.Queue | None + squeue The Queue instance where to put in status information about finished tasks. Can be None to skip sending status information. - rss : RandomStateService | None + rss The RandomStateService instance to use for generating random numbers. - tl : instance of TimeLord | None + tl The instance of TimeLord that should be used to time individual tasks. """ # Get the `QueueHandler` and update its log records queue. logger = get_logger('skyllh') - queue_handler = list(logger.handlers)[0] # noqa: RUF015 + queue_handler = cast(QueueHandler, list(logger.handlers)[0]) # noqa: RUF015 queue_handler.queue = lqueue result_list = [] @@ -173,46 +173,46 @@ def worker_wrapper( # Define a wrapper function that evaluates ``func`` for a subset of # `args_list` on the master process. def master_wrapper( - pbar, - sarr, - func, - sub_args_list, + pbar: ProgressBar | None, + sarr: np.ndarray, + func: Callable, + sub_args_list: list[tuple], squeue=None, - rss=None, - tl=None, + rss: RandomStateService | None = None, + tl: TimeLord | None = None, ): """This is the wrapper function for the master process. Parameters ---------- - pbar : instance of ProgressBar | None + pbar The instance of ProgressBar that should be used to display the progress if the current session is interactive. - sarr : numpy record ndarray + sarr The status numpy record ndarray for all the processes. The length of that array must equal the number of processes, including the master process. Hence, the array index is the process id. The array must contain the following fields: - n_finished_tasks : int + n_finished_tasks The number of finished tasks. - func : callable + func The function which should be called with different arguments, which are given through the sub_args_list argument. If the `rss` argument is not None `func` requires an argument named `rss`. - sub_args_list : list of 2-element tuple + sub_args_list The list of the different arguments for function ``func``. Each element of that list must be a 2-element tuple, where the first element is a tuple of the arguments of ``func``, and the second element is a dictionary with the keyword arguments of ``func``. If the `rss` argument is not None, `func` argument `rss` has to be omitted. - squeue : multiprocessing.Queue | None + squeue The status queue for the worker processes that should be used to receive status information about finished tasks. - rss : RandomStateService | None + rss The RandomStateService instance to use for generating random numbers. - tl : instance of TimeLord | None + tl The instance of TimeLord that should be used to time individual tasks. """ @@ -226,7 +226,7 @@ def master_wrapper( # Skip the rest, if we are not in an interactive session, hence # there is not progress bar. - if not pbar.is_shown: + if pbar is None or not pbar.is_shown: continue sarr[0]['n_finished_tasks'] = master_task_idx + 1 @@ -296,7 +296,7 @@ def master_wrapper( orig_handlers = list(logger.handlers) for orig_handler in orig_handlers: logger.removeHandler(orig_handler) - queue_handler = QueueHandler(lqueue_list[0]) + queue_handler = QueueHandler(cast(mp.Queue, lqueue_list[0])) logger.addHandler(queue_handler) processes = [ @@ -321,7 +321,7 @@ def master_wrapper( # Compute the first chunk in the main process. sarr = np.zeros((len(processes) + 1,), dtype=[('n_finished_tasks', np.int64)]) result_list_0 = master_wrapper( - pbar, sarr, func, sub_args_list_list[0], squeue=squeue, rss=rss_list[0], tl=tl_list[0] + pbar, sarr, func, list(sub_args_list_list[0]), squeue=squeue, rss=rss_list[0], tl=tl_list[0] ) # Initialize logger. @@ -335,6 +335,9 @@ def master_wrapper( # Get the result record from the result queue. result_received = False proc_died = False + pid: int = -1 + result_list: list = [] + proc_tl: TimeLord | None = None while (result_received is False) and (proc_died is False): try: (pid, result_list, proc_tl) = rqueue.get(block=False) @@ -352,12 +355,14 @@ def master_wrapper( raise RuntimeError(f'Child process {proc.pid} did not return with 0! Exit code was {proc.exitcode}.') pid_result_list_map[pid] = result_list - if tl is not None: + if tl is not None and proc_tl is not None: tl.join(proc_tl) logger.debug(f'Beginning of worker process (pid={pid}) log records.') lqueue_end = False + lqueue_pid = lqueue_list[pid] + assert lqueue_pid is not None while not lqueue_end: - record = lqueue_list[pid].get() + record = lqueue_pid.get() if record is None: lqueue_end = True else: @@ -391,6 +396,14 @@ def __init__( ncpu=None, **kwargs, ): + """Creates a new instance of IsParallelizable. + + Parameters + ---------- + ncpu + The number of CPUs to utilize. If set to ``None``, the global + setting will be used. + """ super().__init__(*args, **kwargs) if not isinstance(self, HasConfig): @@ -404,7 +417,7 @@ def ncpu(self): utility function with this property as argument. Hence, if this property is set to None, the global NCPU setting will take precedence. """ - return get_ncpu(self._cfg, self._ncpu) + return get_ncpu(cast(HasConfig, self).cfg, self._ncpu) @ncpu.setter def ncpu(self, n): diff --git a/skyllh/core/parameters.py b/skyllh/core/parameters.py index 4b99156e52..80efad2115 100644 --- a/skyllh/core/parameters.py +++ b/skyllh/core/parameters.py @@ -1,11 +1,14 @@ import itertools +from collections.abc import Sequence from copy import deepcopy +from typing import cast import numpy as np from skyllh.core import ( display, ) +from skyllh.core.binning import BinningDefinition from skyllh.core.model import ( Model, ModelCollection, @@ -20,55 +23,56 @@ issequence, issequenceof, ) +from skyllh.core.random import RandomStateService from skyllh.core.source_model import ( SourceModel, ) -def make_linear_parameter_grid_1d(name, low, high, delta): +def make_linear_parameter_grid_1d(name: str, low: float, high: float, delta: float) -> 'ParameterGrid': """Utility function to create a ParameterGrid object for a 1-dimensional linear parameter grid. Parameters ---------- - name : str + name The name of the parameter. - low : float + low The lowest value of the parameter. - high : float + high The highest value of the parameter. - delta : float + delta The constant distance between the grid values. By definition this defines also the precision of the parameter values. Returns ------- - obj : ParameterGrid + obj The ParameterGrid object holding the discrete parameter grid values. """ grid = np.arange(low, high + delta, delta) return ParameterGrid(name, grid, delta) -def make_logarithmic_parameter_grid_1d(name, low, high, delta): +def make_logarithmic_parameter_grid_1d(name: str, low: float, high: float, delta: float) -> 'ParameterGrid': """Utility function to create a ParameterGrid object for a 1-dimensional logarithmic parameter grid. Parameters ---------- - name : str + name The name of the parameter. - low : float + low The lowest value of the parameter. - high : float + high The highest value of the parameter. - delta : float + delta The logarithmic distance between the grid values. By definition this defines also the precision of the parameter values. Returns ------- - obj : ParameterGrid + obj The ParameterGrid object holding the discrete parameter grid values. """ low = np.log10(low) @@ -84,22 +88,29 @@ class Parameter: parameter has a fixed value or not. """ - def __init__(self, name, initial, valmin=None, valmax=None, isfixed=None): + def __init__( + self, + name: str, + initial: float, + valmin: float | None = None, + valmax: float | None = None, + isfixed: bool | None = None, + ): """Creates a new Parameter instance. Parameters ---------- - name : str + name The name of the parameter. - initial : float + initial The initial value of the parameter. - valmin : float | None + valmin The minimum value of the parameter in case this parameter is mutable. - valmax : float | None + valmax The maximum value of the parameter in case this parameter is mutable. - isfixed : bool | None + isfixed Flag if the value of this parameter is mutable (False), or not (True). If set to `True`, the value of the parameter will always be the `initial` value. @@ -152,7 +163,7 @@ def isfixed(self, b): self._isfixed = b @property - def valmin(self): + def valmin(self) -> float | None: """The minimum bound value of the parameter.""" return self._valmin @@ -162,7 +173,7 @@ def valmin(self, v): self._valmin = v @property - def valmax(self): + def valmax(self) -> float | None: """The maximum bound value of the parameter.""" return self._valmax @@ -194,23 +205,26 @@ def value(self, v): ) self._value = v - def __eq__(self, other): + def __eq__(self, other: object) -> bool: """Implements the equal comparison operator (==). By definition two parameters are equal if there property values are equal. Parameters ---------- - other : Parameter instance + other The instance of Parameter which should be used to compare against this Parameter instance. Returns ------- - cmp : bool + cmp True, if this Parameter instance and the other Parameter instance have the same property values. """ + if not isinstance(other, Parameter): + return NotImplemented + if (self.name != other.name) or (self.value != other.value) or (self.isfixed != other.isfixed): return False @@ -238,19 +252,19 @@ def __str__(self): return s - def as_linear_grid(self, delta): + def as_linear_grid(self, delta: float) -> 'ParameterGrid': """Creates a ParameterGrid instance with a linear grid with constant grid value distances delta. Parameters ---------- - delta : float + delta The constant distance between the grid values. By definition this defines also the precision of the parameter values. Returns ------- - grid : ParameterGrid instance + grid The ParameterGrid instance holding the grid values. Raises @@ -265,23 +279,28 @@ def as_linear_grid(self, delta): grid = np.array([self.initial]) return ParameterGrid(self._name, grid, delta) - grid = make_linear_parameter_grid_1d(name=self._name, low=self._valmin, high=self._valmax, delta=delta) + grid = make_linear_parameter_grid_1d( + name=self._name, + low=cast(float, self._valmin), + high=cast(float, self._valmax), + delta=delta, + ) return grid - def as_logarithmic_grid(self, delta): + def as_logarithmic_grid(self, delta: float) -> 'ParameterGrid': """Creates a ParameterGrid instance with a linear grid with constant grid value distances delta. Parameters ---------- - delta : float + delta The constant distance between the grid values. By definition this defines also the precision of the parameter values. Returns ------- - grid : ParameterGrid instance + grid The ParameterGrid instance holding the grid values. Raises @@ -296,21 +315,26 @@ def as_logarithmic_grid(self, delta): delta = float_cast(delta, 'The delta argument must be castable to type float!') - grid = make_logarithmic_parameter_grid_1d(name=self._name, low=self._valmin, high=self._valmax, delta=delta) + grid = make_logarithmic_parameter_grid_1d( + name=self._name, + low=cast(float, self._valmin), + high=cast(float, self._valmax), + delta=delta, + ) return grid - def change_fixed_value(self, value): + def change_fixed_value(self, value: float) -> float: """Changes the value of this fixed parameter to the given value. Parameters ---------- - value : float + value The parameter's new value. Returns ------- - value : float + value The parameter's new value. Raises @@ -323,19 +347,20 @@ def change_fixed_value(self, value): self.initial = value self.value = value + return value - def make_fixed(self, initial=None): + def make_fixed(self, initial: float | None = None) -> float: """Fixes this parameter to the given initial value. Parameters ---------- - initial : float | None + initial The new fixed initial value of the Parameter. If set to None, the parameter's current value will be used as initial value. Returns ------- - value : float + value The parameter's new value. """ self._isfixed = True @@ -360,27 +385,29 @@ def make_fixed(self, initial=None): return self._value - def make_floating(self, initial=None, valmin=None, valmax=None): + def make_floating( + self, initial: float | None = None, valmin: float | None = None, valmax: float | None = None + ) -> float: """Defines this parameter as floating with the given initial, minimal, and maximal value. Parameters ---------- - initial : float | None + initial The initial value of the parameter. If set to `None`, the parameter's current value will be used as initial value. - valmin : float | None + valmin The minimal value the parameter's value can take. If set to `None`, the parameter's current minimal value will be used. - valmax : float | None + valmax The maximal value the parameter's value can take. If set to `None`, the parameter's current maximal value will be used. Returns ------- - value : float + value The parameter's new value. Raises @@ -421,7 +448,7 @@ class ParameterSet: """This class holds a set of Parameter instances.""" @staticmethod - def union(*paramsets): + def union(*paramsets) -> 'ParameterSet': """Creates a ParameterSet instance that is the union of the given ParameterSet instances. @@ -432,7 +459,7 @@ def union(*paramsets): Returns ------- - paramset : ParameterSet instance + paramset The newly created ParameterSet instance that holds the union of the parameters provided by all the ParameterSet instances. """ @@ -449,12 +476,12 @@ def union(*paramsets): return paramset - def __init__(self, params=None): + def __init__(self, params: 'Parameter | Sequence[Parameter] | None' = None): """Constructs a new ParameterSet instance. Parameters ---------- - params : instance of Parameter | sequence of Parameter instances | None + params The initial sequence of Parameter instances of this ParameterSet instance. """ @@ -605,18 +632,18 @@ def floating_param_bounds(self): return bounds - def __contains__(self, param_name): + def __contains__(self, param_name: str) -> bool: """Implements the ``param_name in self`` expression. It calls the :meth:`has_param` method of this class. Parameters ---------- - param_name : str + param_name The name of the parameter. Returns ------- - check : bool + check Returns ``True`` if the given parameter is part of this ParameterSet instance, ``False`` otherwise. """ @@ -648,17 +675,17 @@ def __str__(self): s += '\n}' return s - def get_fixed_pidx(self, param_name): + def get_fixed_pidx(self, param_name: str) -> int: """Returns the parameter index of the given fixed parameter. Parameters ---------- - param_name : str + param_name The name of the parameter. Returns ------- - pidx : int + pidx The index of the fixed parameter. Raises @@ -668,17 +695,17 @@ def get_fixed_pidx(self, param_name): """ return self._fixed_param_name_to_idx[param_name] - def get_floating_pidx(self, param_name): + def get_floating_pidx(self, param_name: str) -> int: """Returns the parameter index of the given floating parameter. Parameters ---------- - param_name : str + param_name The name of the parameter. Returns ------- - pidx : int + pidx The index of the floating parameter. Raises @@ -689,7 +716,7 @@ def get_floating_pidx(self, param_name): """ return self._floating_param_name_to_idx[param_name] - def generate_random_floating_param_initials(self, rss): + def generate_random_floating_param_initials(self, rss: RandomStateService) -> np.ndarray: """Generates a set of random initials for all floating parameters. A new random initial is defined as @@ -699,13 +726,13 @@ def generate_random_floating_param_initials(self, rss): Parameters ---------- - rss : RandomStateService instance + rss The RandomStateService instance that should be used for drawing random numbers from. Returns ------- - ri : (N_floating_params,)-shaped numpy ndarray + ri The numpy 1D ndarray holding the generated random initial values. """ vb = self.floating_param_bounds @@ -715,46 +742,46 @@ def generate_random_floating_param_initials(self, rss): return ri - def has_fixed_param(self, param_name): + def has_fixed_param(self, param_name: str) -> bool: """Checks if this ParameterSet instance has a fixed parameter named ``param_name``. Parameters ---------- - param_name : str + param_name The name of the parameter. Returns ------- - check : bool + check ``True`` if this ParameterSet instance has a fixed parameter of the given name, ``False`` otherwise. """ return param_name in self._fixed_param_name_list - def has_floating_param(self, param_name): + def has_floating_param(self, param_name: str) -> bool: """Checks if this ParameterSet instance has a floating parameter named ``param_name``. Parameters ---------- - param_name : str + param_name The name of the parameter. Returns ------- - check : bool + check ``True`` if this ParameterSet instance has a floating parameter of the given name, ``False`` otherwise. """ return param_name in self._floating_param_name_list - def make_params_fixed(self, fix_params): + def make_params_fixed(self, fix_params: dict): """Fixes the given parameters to the given values. Parameters ---------- - fix_params : dict + fix_params The dictionary defining the parameters that should get fixed to the given dictionary entry values. @@ -790,13 +817,13 @@ def make_params_fixed(self, fix_params): self._floating_param_name_list += [pname] self._floating_param_name_to_idx[pname] = len(self._floating_param_name_list) - 1 - def make_params_floating(self, float_params): + def make_params_floating(self, float_params: dict): """Makes the given parameters floating with the given initial value and within the given bounds. Parameters ---------- - float_params : dict + float_params The dictionary defining the parameters that should get set to be floating. The format of a dictionary's entry can be one of the following formats: @@ -804,7 +831,7 @@ def make_params_floating(self, float_params): ``None`` The parameter's initial, minimal and maximal value should be taken from the parameter's current settings. - initial : float + initial The parameter's initial value should be set to the given value. The minimal and maximal values of the parameter will be taken from the parameter's current settings. @@ -864,31 +891,31 @@ def update_fixed_param_value_cache(self): for i, param in enumerate(self.fixed_params): self._fixed_param_values[i] = param.value - def copy(self): + def copy(self) -> 'ParameterSet': """Creates a deep copy of this ParameterSet instance. Returns ------- - copy : ParameterSet instance + copy The copied instance of this ParameterSet instance. """ copy = deepcopy(self) return copy - def add_param(self, param, atfront=False): + def add_param(self, param: 'Parameter', atfront: bool = False) -> 'ParameterSet': """Adds the given Parameter instance to this set of parameters. Parameters ---------- - param : instance of Parameter + param The parameter, which should get added. - atfront : bool + atfront Flag if the parameter should be added at the front of the parameter list. If set to False (default), it will be added at the back. Returns ------- - self : instance of ParameterSet + self This ParameterSet instance so that multiple add_param calls can just be concatenated. @@ -912,7 +939,7 @@ def add_param(self, param, atfront=False): if atfront: # Add parameter at front of parameter list. - self._params = np.concatenate(([param], self._params)) + self._params = np.concatenate((np.array([param], dtype=object), self._params)) self._params_fixed_mask = np.concatenate(([param_fixed_mask], self._params_fixed_mask)) if param.isfixed: self._fixed_param_name_list = [param.name, *self._fixed_param_name_list] @@ -927,7 +954,7 @@ def add_param(self, param, atfront=False): self._floating_param_name_to_idx[param.name] = 0 else: # Add parameter at back of parameter list. - self._params = np.concatenate((self._params, [param])) + self._params = np.concatenate((self._params, np.array([param], dtype=object))) self._params_fixed_mask = np.concatenate((self._params_fixed_mask, [param_fixed_mask])) if param.isfixed: self._fixed_param_name_list = [*self._fixed_param_name_list, param.name] @@ -939,37 +966,39 @@ def add_param(self, param, atfront=False): return self - def has_param(self, param): + def has_param(self, param: 'str | Parameter') -> bool: """Checks if the given Parameter is already present in this ParameterSet instance. The check is performed based on the parameter name. Parameters ---------- - param : Parameter instance - The Parameter instance that should be checked. + param + The Parameter instance or parameter name string that should be + checked. Returns ------- - check : bool + check ``True`` if the given parameter is present in this parameter set, ``False`` otherwise. """ - return (param.name in self._floating_param_name_list) or (param.name in self._fixed_param_name_list) + name = param if isinstance(param, str) else param.name + return (name in self._floating_param_name_list) or (name in self._fixed_param_name_list) - def get_params_dict(self, floating_param_values): + def get_params_dict(self, floating_param_values: np.ndarray) -> dict: """Converts the given floating parameter values into a dictionary with the floating parameter names and values and also adds the fixed parameter names and their values to this dictionary. Parameters ---------- - floating_param_values : 1D ndarray + floating_param_values The ndarray holding the values of the floating parameters in the order that the floating parameters are defined. Returns ------- - params_dict : dict + params_dict The dictionary with the floating and fixed parameter names and values. """ @@ -980,19 +1009,19 @@ def get_params_dict(self, floating_param_values): return params_dict - def get_floating_params_dict(self, floating_param_values): + def get_floating_params_dict(self, floating_param_values: np.ndarray) -> dict: """Converts the given floating parameter values into a dictionary with the floating parameter names and values. Parameters ---------- - floating_param_values : 1D ndarray + floating_param_values The ndarray holding the values of the floating parameters in the order that the floating parameters are defined. Returns ------- - params_dict : dict + params_dict The dictionary with the floating and fixed parameter names and values. """ @@ -1008,19 +1037,21 @@ class ParameterGrid: """ @staticmethod - def from_BinningDefinition(binning, delta=None, decimals=None): + def from_BinningDefinition( + binning: BinningDefinition, delta: float | None = None, decimals: int | None = None + ) -> 'ParameterGrid': """Creates a ParameterGrid instance from a BinningDefinition instance. Parameters ---------- - binning : BinningDefinition instance + binning The BinningDefinition instance that should be used to create the ParameterGrid instance from. - delta : float | None + delta The width between the grid values. If set to ``None``, the width is taken from the equal-distant ``grid`` values. - decimals : int | None + decimals The number of decimals the grid values should get rounded to. The maximal number of decimals is 16. If set to None, the number of decimals will be the maximum of the @@ -1029,27 +1060,27 @@ def from_BinningDefinition(binning, delta=None, decimals=None): Returns ------- - param_grid : instance of ParameterGrid + param_grid The created ParameterGrid instance. """ return ParameterGrid(name=binning.name, grid=binning.binedges, delta=delta, decimals=decimals) @staticmethod - def from_range(name, start, stop, delta, decimals=None): + def from_range(name: str, start: float, stop: float, delta: float, decimals: int | None = None) -> 'ParameterGrid': """Creates a ParameterGrid instance from a range definition. The stop value will be the last grid point. Parameters ---------- - name : str + name The name of the parameter grid. - start : float + start The start value of the range. - stop : float + stop The end value of the range. - delta : float + delta The width between the grid values. - decimals : int | None + decimals The number of decimals the grid values should get rounded to. The maximal number of decimals is 16. If set to None, the number of decimals will be the maximum of the @@ -1058,7 +1089,7 @@ def from_range(name, start, stop, delta, decimals=None): Returns ------- - param_grid : instance of ParameterGrid + param_grid The created ParameterGrid instance. """ start = float_cast(start, 'The start argument must be castable to type float!') @@ -1070,21 +1101,23 @@ def from_range(name, start, stop, delta, decimals=None): return ParameterGrid(name=name, grid=grid, delta=delta, decimals=decimals) - def __init__(self, name, grid, delta=None, decimals=None): + def __init__( + self, name: str, grid: Sequence[float] | np.ndarray, delta: float | None = None, decimals: int | None = None + ): """Creates a new parameter grid. Parameters ---------- - name : str + name The name of the parameter. - grid : sequence of float + grid The sequence of float values defining the discrete grid values of the parameter. - delta : float | None + delta The width between the grid values. If set to ``None``, the width is taken from the equal-distant ``grid`` values. - decimals : int | None + decimals The number of decimals the grid values should get rounded to. The maximal number of decimals is 16. If set to None, the number of decimals will be the maximum of the @@ -1094,7 +1127,7 @@ def __init__(self, name, grid, delta=None, decimals=None): if delta is None: # We need to take the mean of all the "equal" differences in order # to smooth out unlucky rounding issues of a particular difference. - delta = np.mean(np.diff(grid)) + delta = float(np.mean(np.diff(grid))) delta = float_cast(delta, 'The delta argument must be castable to type float!') self._delta = np.float64(delta) @@ -1140,7 +1173,7 @@ def decimals(self): return self._decimals @property - def grid(self): + def grid(self) -> np.ndarray: """The numpy.ndarray with the grid values of the parameter.""" return self._grid @@ -1152,7 +1185,7 @@ def grid(self, arr): arr = np.array(arr, dtype=np.float64) if arr.ndim != 1: raise ValueError('The grid property must be a 1D numpy.ndarray!') - self._grid = self.round_to_nearest_grid_point(arr) + self._grid = np.asarray(self.round_to_nearest_grid_point(arr)) @property def delta(self): @@ -1205,17 +1238,17 @@ def copy(self): copy = deepcopy(self) return copy - def round_to_nearest_grid_point(self, value): + def round_to_nearest_grid_point(self, value: float | np.ndarray) -> float | np.ndarray: """Rounds the given value to the nearest grid point. Parameters ---------- - value : float | ndarray of float + value The value(s) to round. Returns ------- - grid_point : float | ndarray of float + grid_point The calculated grid point(s). """ scalar_input = np.isscalar(value) @@ -1229,7 +1262,7 @@ def round_to_nearest_grid_point(self, value): return gp - def round_to_lower_grid_point(self, value): + def round_to_lower_grid_point(self, value: float | np.ndarray) -> float | np.ndarray: """Rounds the given value to the nearest grid point that is lower than the given value. @@ -1238,12 +1271,12 @@ def round_to_lower_grid_point(self, value): Parameters ---------- - value : float | ndarray of float + value The value(s) to round. Returns ------- - grid_point : float | ndarray of float + grid_point The calculated grid point(s). """ scalar_input = np.isscalar(value) @@ -1257,7 +1290,7 @@ def round_to_lower_grid_point(self, value): return gp - def round_to_upper_grid_point(self, value): + def round_to_upper_grid_point(self, value: float | np.ndarray) -> float | np.ndarray: """Rounds the given value to the nearest grid point that is larger than the given value. @@ -1266,12 +1299,12 @@ def round_to_upper_grid_point(self, value): Parameters ---------- - value : float | ndarray of float + value The value(s) to round. Returns ------- - grid_point : ndarray of float + grid_point The calculated grid point(s). """ scalar_input = np.isscalar(value) @@ -1293,18 +1326,18 @@ class IrregularParameterGrid: """ @staticmethod - def from_BinningDefinition(binning): + def from_BinningDefinition(binning: BinningDefinition) -> 'IrregularParameterGrid': """Creates a IrregularParameterGrid instance from a BinningDefinition instance. Parameters ---------- - binning : BinningDefinition instance + binning The BinningDefinition instance that should be used to create the ParameterGrid instance from. Returns ------- - param_grid : instance of ParameterGrid + param_grid The created ParameterGrid instance. """ return IrregularParameterGrid( @@ -1312,14 +1345,14 @@ def from_BinningDefinition(binning): grid=binning.binedges, ) - def __init__(self, name, grid): + def __init__(self, name: str, grid: Sequence[float] | np.ndarray): """Creates a new parameter grid. Parameters ---------- - name : str + name The name of the parameter. - grid : sequence of float + grid The sequence of float values defining the discrete grid values of the parameter. """ @@ -1342,7 +1375,7 @@ def name(self, name): self._name = name @property - def grid(self): + def grid(self) -> np.ndarray: """The numpy.ndarray with the grid values of the parameter.""" return self._grid @@ -1378,7 +1411,7 @@ def copy(self): copy = deepcopy(self) return copy - def round_to_nearest_grid_point(self, value): + def round_to_nearest_grid_point(self, value: float | np.ndarray) -> float | np.ndarray: """Rounds the given value to the nearest grid point. Note: If the given value is precisely in between two nearest grid @@ -1386,12 +1419,12 @@ def round_to_nearest_grid_point(self, value): Parameters ---------- - value : float | ndarray of float + value The value(s) to round. Returns ------- - grid_point : float | ndarray of float + grid_point The calculated grid point(s). """ scalar_input = np.isscalar(value) @@ -1405,7 +1438,7 @@ def round_to_nearest_grid_point(self, value): return gp - def round_to_lower_grid_point(self, value): + def round_to_lower_grid_point(self, value: float | np.ndarray) -> float | np.ndarray: """Rounds the given value to the nearest grid point that is lower than the given value. @@ -1414,12 +1447,12 @@ def round_to_lower_grid_point(self, value): Parameters ---------- - value : float | ndarray of float + value The value(s) to round. Returns ------- - grid_point : float | ndarray of float + grid_point The calculated grid point(s). """ scalar_input = np.isscalar(value) @@ -1432,7 +1465,7 @@ def round_to_lower_grid_point(self, value): return gp - def round_to_upper_grid_point(self, value): + def round_to_upper_grid_point(self, value: float | np.ndarray) -> float | np.ndarray: """Rounds the given value to the nearest grid point that is larger than the given value. @@ -1441,12 +1474,12 @@ def round_to_upper_grid_point(self, value): Parameters ---------- - value : float | ndarray of float + value The value(s) to round. Returns ------- - grid_point : ndarray of float + grid_point The calculated grid point(s). """ scalar_input = np.isscalar(value) @@ -1463,17 +1496,17 @@ def round_to_upper_grid_point(self, value): class ParameterGridSet(NamedObjectCollection): """Describes a set of parameter grids.""" - def __init__(self, param_grids=None, **kwargs): + def __init__(self, param_grids: 'Sequence[ParameterGrid] | ParameterGrid | None' = None, **kwargs): """Constructs a new ParameterGridSet object. Parameters ---------- - param_grids : sequence of instance of ParameterGrid | instance of ParameterGrid | None + param_grids The ParameterGrid instances this instance of ParameterGridSet should get initialized with. """ # Infer `obj_type` from the `param_grids` argument. - if issequence(param_grids): # noqa: SIM108 + if isinstance(param_grids, (list, tuple)) and len(param_grids) > 0: obj_type = type(param_grids[0]) else: obj_type = type(param_grids) @@ -1532,30 +1565,32 @@ class ParameterModelMapper: """ @staticmethod - def is_global_fitparam_a_local_param(fitparam_id, params_recarray, local_param_names): + def is_global_fitparam_a_local_param( + fitparam_id: int, params_recarray: np.ndarray, local_param_names: list[str] + ) -> bool: """Determines if the given global fit parameter is a local parameter of the given list of local parameter names. Parameters ---------- - fitparam_id : int + fitparam_id The ID of the global fit parameter. - params_recarray : instance of numpy record ndarray + params_recarray The (N_models,)-shaped numpy record ndarray holding the local parameter names and values of the models. See the :meth:`skyllh.core.parameters.ParameterModelMapper.create_src_params_recarray` method for the format of this record array. - local_param_names : list of str + local_param_names The list of local parameters. Returns ------- - check : bool + check ``True`` if the global fit parameter translates to a local parameter contained in the ``local_param_names`` list, ``False`` otherwise. """ for pname in local_param_names: - if pname not in params_recarray.dtype.fields: + if params_recarray.dtype.fields is None or pname not in params_recarray.dtype.fields: continue if np.any(params_recarray[f'{pname}:gpidx'] == fitparam_id + 1): return True @@ -1563,14 +1598,14 @@ def is_global_fitparam_a_local_param(fitparam_id, params_recarray, local_param_n return False @staticmethod - def is_local_param_a_fitparam(local_param_name, params_recarray): + def is_local_param_a_fitparam(local_param_name: str, params_recarray: np.ndarray) -> bool: """Checks if the given local parameter is a (partly) a fit parameter. Parameters ---------- - local_param_name : str + local_param_name The name of the local parameter. - params_recarray : instance of numpy record ndarray + params_recarray The (N_models,)-shaped numpy record ndarray holding the local parameter names and values of the models. See the :meth:`skyllh.core.parameters.ParameterModelMapper.create_src_params_recarray` @@ -1578,24 +1613,26 @@ def is_local_param_a_fitparam(local_param_name, params_recarray): Returns ------- - check : bool + check ``True`` if the given local parameter is (partly) a fit parameter. """ - return np.any(params_recarray[f'{local_param_name}:gpidx'] > 0) + return bool(np.any(params_recarray[f'{local_param_name}:gpidx'] > 0)) - def __init__(self, models, **kwargs): + def __init__(self, models: Sequence[Model], **kwargs): """Constructor of the parameter mapper. Parameters ---------- - models : sequence of instance of Model. + models The sequence of Model instances the parameter mapper can map global parameters to. """ super().__init__(**kwargs) - models = ModelCollection.cast(models, 'The models property must be castable to an instance of ModelCollection!') - self._models = models + _models_coll = ModelCollection.cast( + models, 'The models property must be castable to an instance of ModelCollection!' + ) + self._models = _models_coll # Create the parameter set for the global parameters. self._global_paramset = ParameterSet() @@ -1709,20 +1746,20 @@ def __str__(self): return s - def get_model_param_name(self, model_idx, gp_idx): + def get_model_param_name(self, model_idx: int, gp_idx: int) -> str | None: """Retrieves the local parameter name of a given model and global parameter index. Parameters ---------- - model_idx : int + model_idx The index of the model. - gp_idx : int + gp_idx The index of the global parameter. Returns ------- - param_name : str | None + param_name The name of the local model parameter. It is ``None``, if the given global parameter is not mapped to the given model. """ @@ -1730,33 +1767,33 @@ def get_model_param_name(self, model_idx, gp_idx): return param_name - def get_gflp_idx(self, name): + def get_gflp_idx(self, name: str) -> int: """Gets the index of the global floating parameter of the given name. Parameters ---------- - name : str + name The global floating parameter's name. Returns ------- - idx : int + idx The index of the global floating parameter. """ return self._global_paramset.get_floating_pidx(param_name=name) - def get_model_idx_by_name(self, name): + def get_model_idx_by_name(self, name: str) -> int: """Determines the index within this ParameterModelMapper instance of the model with the given name. Parameters ---------- - name : str + name The model's name. Returns ------- - model_idx : int + model_idx The model's index within this ParameterModelMapper instance. Raises @@ -1770,19 +1807,19 @@ def get_model_idx_by_name(self, name): raise KeyError(f'The model with name "{name}" does not exist within the ParameterModelMapper instance!') - def get_src_model_idxs(self, sources=None): + def get_src_model_idxs(self, sources: SourceModel | Sequence[SourceModel] | None = None) -> np.ndarray: """Creates a numpy ndarray holding the indices of the requested source models. Parameters ---------- - sources : instance of SourceModel | sequence of SourceModel | None + sources The requested sequence of source models. If set to ``None``, all source models will be requested. Returns ------- - src_model_idxs : numpy ndarray + src_model_idxs The (N_sources,)-shaped 1D ndarray holding the indices of the requested source models. """ @@ -1812,20 +1849,25 @@ def get_src_model_idxs(self, sources=None): return src_model_idxs - def map_param(self, param, models=None, model_param_names=None): + def map_param( + self, + param: 'Parameter', + models: Model | Sequence[Model] | None = None, + model_param_names: str | Sequence[str] | None = None, + ) -> 'ParameterModelMapper': """Maps the given instance of Parameter to the given sequence of models this parameter model mapper knows about. Aliases for the given parameter can be specified for each individual model. Parameters ---------- - param : instance of Parameter + param The global parameter which should get mapped to one or more models. - models : sequence of Model instances + models The sequence of Model instances the parameter should get mapped to. The instances in the sequence must match Model instances specified at construction of this mapper. - model_param_names : str | sequence of str | None + model_param_names The name of the parameter of the model. Hence, the global parameter name can be different to the parameter name of the model. If `None`, the name of the global parameter will be used as model @@ -1833,7 +1875,7 @@ def map_param(self, param, models=None, model_param_names=None): Returns ------- - self : ParameterModelMapper + self The instance of this ParameterModelMapper, so that several `map_param` calls can be concatenated. @@ -1843,25 +1885,32 @@ def map_param(self, param, models=None, model_param_names=None): If there is already a model parameter of the same name defined for any of the given to-be-applied models. """ + _mpnames_arr: np.ndarray if model_param_names is None: - model_param_names = np.array([param.name] * len(self._models)) - if isinstance(model_param_names, str): - model_param_names = np.array([model_param_names] * len(self._models)) - if not issequenceof(model_param_names, str): - raise TypeError( - 'The model_param_names argument must be None, an instance of str, or a sequence of instances of str!' - ) + _mpnames_arr = np.array([param.name] * len(self._models)) + elif isinstance(model_param_names, str): + _mpnames_arr = np.array([model_param_names] * len(self._models)) + else: + if not issequenceof(model_param_names, str): + raise TypeError( + 'The model_param_names argument must be None, an instance of str, or a sequence of instances of str!' + ) + _mpnames_arr = np.asarray(model_param_names) + _models_coll: ModelCollection if models is None: - models = self._models - models = ModelCollection.cast(models, 'The models argument must be castable to an instance of ModelCollection!') + _models_coll = self._models + else: + _models_coll = ModelCollection.cast( + models, 'The models argument must be castable to an instance of ModelCollection!' + ) # Make sure that the user did not provide an empty sequence. - if len(models) == 0: + if len(_models_coll) == 0: raise ValueError('The sequence of models, to which the parameter maps, cannot be empty!') # Get the list of model indices to which the parameter maps. mask = np.zeros((self.n_models,), dtype=np.bool_) - for (midx, model), applied_model in itertools.product(enumerate(self._models), models): + for (midx, model), applied_model in itertools.product(enumerate(self._models), _models_coll): if applied_model.id == model.id: mask[midx] = True @@ -1869,35 +1918,35 @@ def map_param(self, param, models=None, model_param_names=None): # the given to-be-mapped models. for midx in np.arange(self.n_models)[mask]: mpnames = self._model_param_names[midx][self._model_param_names[midx] != np.array(None)] - if model_param_names[midx] in mpnames: + if _mpnames_arr[midx] in mpnames: raise KeyError( - f'The model parameter "{model_param_names[midx]}" is ' + f'The model parameter "{_mpnames_arr[midx]}" is ' f'already defined for model "{self._models[midx].name}"!' ) self._global_paramset.add_param(param) - entry = np.where(mask, model_param_names, None) + entry = np.where(mask, _mpnames_arr, np.full(self.n_models, None, dtype=object)) self._model_param_names = np.hstack((self._model_param_names, entry[np.newaxis, :].T)) return self - def create_model_params_dict(self, gflp_values, model): + def create_model_params_dict(self, gflp_values: np.ndarray, model: Model | str | int) -> dict: """Creates a dictionary with the fixed and floating parameter names and their values for the given model. Parameters ---------- - gflp_values : 1D ndarray of float + gflp_values The ndarray instance holding the current values of the global floating parameters. - model : instance of Model | str | int + model The index of the model as it was defined at construction time of this ParameterModelMapper instance. Returns ------- - model_param_dict : dict + model_param_dict The dictionary holding the fixed and floating parameter names and values of the specified model. """ @@ -1938,7 +1987,11 @@ def create_model_params_dict(self, gflp_values, model): return model_param_dict - def create_src_params_recarray(self, gflp_values=None, sources=None): + def create_src_params_recarray( + self, + gflp_values: np.ndarray | None = None, + sources: SourceModel | Sequence[SourceModel] | np.ndarray | None = None, + ) -> np.ndarray: """Creates a numpy record ndarray with a field for each local source parameter name and parameter's value. In addition each parameter field ```` has a field named ``<:gpidx>`` which holds the index @@ -1951,13 +2004,13 @@ def create_src_params_recarray(self, gflp_values=None, sources=None): Parameters ---------- - gflp_values : numpy ndarray | None + gflp_values The (N_global_floating_param,)-shaped 1D ndarray holding the global floating parameter values. The order must match the order of parameter definition in this ParameterModelMapper instance. If set to ``None``, the value ``numpy.nan`` will be used as parameter value for floating parameters. - sources : SourceModel | sequence of SourceModel | ndarray of int32 | None + sources The sources which should be considered. If a ndarray of type int is provides, it must contain the global source indices. @@ -1965,7 +2018,7 @@ def create_src_params_recarray(self, gflp_values=None, sources=None): Returns ------- - recarray : numpy structured ndarray + recarray The (N_sources,)-shaped numpy structured ndarray holding the local parameter names and their values for each requested source. It contains the following fields: @@ -1976,7 +2029,7 @@ def create_src_params_recarray(self, gflp_values=None, sources=None): The field holding the value for the local parameter . Not all local parameters apply to all sources. - Example: "gamma". + Example :gpidx The field holding the global parameter index plus one for the local parameter . Example: "gamma:gpidx". Indices @@ -2002,7 +2055,7 @@ def create_src_params_recarray(self, gflp_values=None, sources=None): smidxs = sources else: # Get the source indices of the requested sources. - smidxs = self.get_src_model_idxs(sources=sources) + smidxs = self.get_src_model_idxs(sources=cast(SourceModel | Sequence[SourceModel] | None, sources)) # Create the output record array with nan as default value. dtype = [(':model_idx', np.int32)] @@ -2056,19 +2109,19 @@ def create_src_params_recarray(self, gflp_values=None, sources=None): return recarray - def create_global_params_dict(self, gflp_values): + def create_global_params_dict(self, gflp_values: np.ndarray) -> dict: """Converts the given global floating parameter values into a dictionary holding the names and values of all floating and fixed parameters. Parameters ---------- - gflp_values : numpy ndarray + gflp_values The (n_global_floating_params,)-shaped 1D numpy ndarray holding the values of the global floating parameters. Returns ------- - params_dict : dict + params_dict The dictionary holding the parameter name and values of all floating and fixed parameters. """ @@ -2076,19 +2129,19 @@ def create_global_params_dict(self, gflp_values): return params_dict - def create_global_floating_params_dict(self, gflp_values): + def create_global_floating_params_dict(self, gflp_values: np.ndarray) -> dict: """Converts the given global floating parameter values into a dictionary holding the names and values of all floating parameters. Parameters ---------- - gflp_values : numpy ndarray + gflp_values The (n_global_floating_params,)-shaped 1D numpy ndarray holding the values of the global floating parameters. Returns ------- - params_dict : dict + params_dict The dictionary holding the parameter name and values of all floating parameters. """ @@ -2098,19 +2151,19 @@ def create_global_floating_params_dict(self, gflp_values): def get_local_param_is_global_floating_param_mask( self, - local_param_names, - ): + local_param_names: Sequence[str], + ) -> np.ndarray: """Checks which local parameter name is mapped to a global floating parameter. Parameters ---------- - local_param_names : sequence of str + local_param_names The sequence of the local parameter names to test. Returns ------- - mask : instance of ndarray + mask The (N_local_param_names,)-shaped numpy ndarray holding the mask for each local parameter name if it is mapped to a global floating parameter. diff --git a/skyllh/core/pdf.py b/skyllh/core/pdf.py index 8fc3192fbb..90d8f078a3 100644 --- a/skyllh/core/pdf.py +++ b/skyllh/core/pdf.py @@ -1,4 +1,6 @@ import abc +from collections.abc import Callable, Sequence +from typing import cast import numpy as np from scipy.interpolate import ( @@ -27,6 +29,7 @@ get_logger, ) from skyllh.core.parameters import ( + Parameter, ParameterGrid, ParameterGridSet, ParameterModelMapper, @@ -41,9 +44,8 @@ issequenceof, make_dict_hash, ) -from skyllh.core.timing import ( - TaskTimer, -) +from skyllh.core.timing import TaskTimer, TimeLord +from skyllh.core.trialdata import TrialDataManager logger = get_logger(__name__) @@ -54,16 +56,16 @@ class PDFAxis: plot a PDF or a PDF ratio. """ - def __init__(self, name, vmin, vmax, *args, **kwargs): + def __init__(self, name: str, vmin: float, vmax: float, *args, **kwargs): """Creates a new axis for a PDF. Parameters ---------- - name : str + name The name of the axis. - vmin : float + vmin The minimal value of the axis. - vmax : float + vmax The maximal value of the axis. """ super().__init__(*args, **kwargs) @@ -131,7 +133,7 @@ class PDFAxes(NamedObjectCollection): """ @staticmethod - def union(*axeses): + def union(*axeses) -> 'PDFAxes': """Creates a PDFAxes instance that is the union of the given PDFAxes instances. @@ -142,7 +144,7 @@ def union(*axeses): Returns ------- - axes : PDFAxes instance + axes The newly created PDFAxes instance that holds the union of the PDFAxis instances provided by all the PDFAxes instances. """ @@ -159,12 +161,12 @@ def union(*axeses): return axes - def __init__(self, axes=None, **kwargs): + def __init__(self, axes: 'Sequence[PDFAxis] | None' = None, **kwargs): """Creates a new PDFAxes instance. Parameters ---------- - axes : sequence of instance of PDFAxis | None + axes The sequence of instance of PDFAxis for this PDFAxes instance. If set to ``None``, the PDFAxes instance will be empty. """ @@ -174,19 +176,19 @@ def __str__(self): """Pretty string implementation for the PDFAxes instance.""" return '\n'.join(str(axis) for axis in self) - def is_same_as(self, axes): + def is_same_as(self, axes: 'PDFAxes | Sequence[PDFAxis]') -> bool: """Checks if this PDFAxes object has the same axes and range then the given PDFAxes object. Parameters ---------- - axes : instance of PDFAxes | sequence of PDFAxis + axes The instance of PDFAxes or the sequence of instance of PDFAxis that should be compared to the axes of this PDFAxes instance. Returns ------- - check : bool + check True, if this PDFAxes and the given PDFAxes have the same axes and ranges. False otherwise. """ @@ -210,19 +212,20 @@ def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) - def __mul__(self, other): + def __mul__(self, other: 'IsBackgroundPDF'): """Creates a BackgroundPDFProduct instance for the multiplication of this background PDF and another background PDF. Parameters ---------- - other : instance of IsBackgroundPDF + other The instance of IsBackgroundPDF, which is the other background PDF. """ if not isinstance(other, IsBackgroundPDF): raise TypeError('The other PDF must be an instance of IsBackgroundPDF!') - return BackgroundPDFProduct(self, other, cfg=self.cfg) + _self = cast('PDF', self) + return BackgroundPDFProduct(_self, cast('PDF', other), cfg=cast(HasConfig, self).cfg) class IsSignalPDF: @@ -239,19 +242,19 @@ def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) - def __mul__(self, other): + def __mul__(self, other: 'IsSignalPDF'): """Creates a SignalPDFProduct instance for the multiplication of this signal PDF and another signal PDF. Parameters ---------- - other : instance of IsSignalPDF + other The instance of IsSignalPDF, which is the other signal PDF. """ if not isinstance(other, IsSignalPDF): raise TypeError('The other PDF must be an instance of IsSignalPDF!') - return SignalPDFProduct(self, other, cfg=self.cfg) + return SignalPDFProduct(cast('PDF', self), cast('PDF', other), cfg=cast(HasConfig, self).cfg) class PDF( @@ -268,19 +271,19 @@ class PDF( def __init__( self, - pmm=None, - param_set=None, + pmm: ParameterModelMapper | None = None, + param_set: Parameter | Sequence[Parameter] | ParameterSet | None = None, **kwargs, ): """Creates a new PDF instance. Parameters ---------- - pmm : instance of ParameterModelMapper | None + pmm The instance of ParameterModelMapper defining the global parameters and their mapping to local model/source parameters. It can be ``None``, if the PDF does not depend on any parameters. - param_set : instance of Parameter | sequence of instance of Parameter | instance of ParameterSet | None + param_set If this PDF depends on parameters, this set of parameters defines them. If a single parameter instance is given a ParameterSet instance will be created holding this single parameter. @@ -364,7 +367,7 @@ def add_axis(self, axis): self._axes += axis @abc.abstractmethod - def assert_is_valid_for_trial_data(self, tdm, tl=None, **kwargs): + def assert_is_valid_for_trial_data(self, tdm: TrialDataManager, tl: TimeLord | None = None, **kwargs): """This method is supposed to check if this PDF is valid for all the given trial data. This means, it needs to check if there is a PDF value for each trial data event that will be used in the @@ -374,9 +377,9 @@ def assert_is_valid_for_trial_data(self, tdm, tl=None, **kwargs): Parameters ---------- - tdm : instance of TrialDataManager + tdm The instance of TrialDataManager holding the trial data events. - tl : instance of TimeLord | None + tl The optional instance of TimeLord for measuring timing information. Raises @@ -385,32 +388,34 @@ def assert_is_valid_for_trial_data(self, tdm, tl=None, **kwargs): If some of the trial data is outside the PDF's value space. """ - def initialize_for_new_trial(self, tdm, tl=None, **kwargs): + def initialize_for_new_trial(self, tdm: TrialDataManager, tl: TimeLord | None = None, **kwargs): """This method is called when a new trial is initialized. Derived classes can use this call hook to pre-compute time-expensive data, which do not depend on any fit parameters. Parameters ---------- - tdm : instance of TrialDataManager + tdm The instance of TrialDataManager holding the new trial data. - tl : instance of TimeLord | None + tl The optional instance of TimeLord to measure timing information. """ @abc.abstractmethod - def get_pd(self, tdm, params_recarray=None, tl=None): + def get_pd( + self, tdm: TrialDataManager, params_recarray: np.ndarray | None = None, tl: TimeLord | None = None + ) -> tuple[np.ndarray, dict]: """This abstract method is supposed to calculate the probability density for the specified events given the specified parameter values. Parameters ---------- - tdm : instance of TrialDataManager + tdm The instance of TrialDataManager holding the data events for which the probability density should be calculated. What data fields are required is defined by the derived PDF class and depends on the application. - params_recarray : numpy record ndarray | None + params_recarray The (N_models,)-shaped numpy structured ndarray holding the local parameter names and values of the models. The models are defined by the ParameterModelMapper instance. @@ -423,13 +428,13 @@ def get_pd(self, tdm, params_recarray=None, tl=None): source value. For values mapping to non-fit parameters, the index should be negative. This can be ``None`` for PDFs that do not depend on any parameters. - tl : instance of TimeLord | None + tl The optional instance of TimeLord that should be used to measure timing information. Returns ------- - pd : instance of numpy ndarray + pd The (N_values,)-shaped numpy ndarray holding the probability density for each event. The length of this 1D array depends on the number of sources and the events belonging to those sources. In the worst @@ -437,7 +442,7 @@ def get_pd(self, tdm, params_recarray=None, tl=None): The assignment of values to sources is given by the :py:attr:`~skyllh.core.trialdata.TrialDataManager.src_evt_idxs` property. - grads : dict + grads The dictionary holding the gradients of the probability density w.r.t. each global fit parameter. The key of the dictionary is the id of the global fit parameter. The value is a (N_values,)-shaped @@ -452,16 +457,16 @@ class PDFProduct( ``pdf1 * pdf2``. It is derived from the PDF class and hence is a PDF itself. """ - def __init__(self, pdf1, pdf2, **kwargs): + def __init__(self, pdf1: 'PDF', pdf2: 'PDF', **kwargs): """Creates a new PDFProduct instance, which implements the operation ``pdf1 * pdf2``. The axes of the two PDF instances will be merged. Parameters ---------- - pdf1 : instance of PDF + pdf1 The left-hand-side PDF in the operation ``pdf1 * pdf2``. - pdf2 : instance of PDF + pdf2 The right-hand-side PDF in the operation ``pdf1 * pdf2``. """ self.pdf1 = pdf1 @@ -506,16 +511,16 @@ def pdf2(self, pdf): raise TypeError('The pdf2 property must be an instance of PDF!') self._pdf2 = pdf - def assert_is_valid_for_trial_data(self, tdm, tl=None, **kwargs): + def assert_is_valid_for_trial_data(self, tdm: TrialDataManager, tl: TimeLord | None = None, **kwargs): """Calls the :meth:`assert_is_valid_for_trial_data` method of ``pdf1`` and ``pdf2``. Parameters ---------- - tdm : instance of TrialDataManager + tdm The instance of TrialDataManager that should be used to get the trial data from. - tl : instance of TimeLord | None + tl The optional instance of TimeLord for measuring timing information. Raises @@ -529,8 +534,8 @@ def assert_is_valid_for_trial_data(self, tdm, tl=None, **kwargs): def initialize_for_new_trial( self, - tdm, - tl=None, + tdm: TrialDataManager, + tl: TimeLord | None = None, **kwargs, ): """Calls the ``initialize_for_new_trial`` method of the two PDF @@ -538,16 +543,18 @@ def initialize_for_new_trial( Parameters ---------- - tdm : instance of TrialDataManager + tdm The instance of TrialDataManager holding the new trial event data. - tl : instance of TimeLord | None + tl The optional instance of TimeLord to use for measuring timing information. """ self._pdf1.initialize_for_new_trial(tdm=tdm, tl=tl, **kwargs) self._pdf2.initialize_for_new_trial(tdm=tdm, tl=tl, **kwargs) - def get_pd(self, tdm, params_recarray=None, tl=None): + def get_pd( + self, tdm: TrialDataManager, params_recarray: np.ndarray | None = None, tl: TimeLord | None = None + ) -> tuple[np.ndarray, dict]: """Calculates the probability density for the trial events given the specified parameters by calling the `get_pd` method of `pdf1` and `pdf2` and combining the two property densities by multiplication. @@ -556,25 +563,25 @@ def get_pd(self, tdm, params_recarray=None, tl=None): Parameters ---------- - tdm : instance of TrialDataManager + tdm The TrialDataManager instance holding the trial event data for which the PDF values should get calculated. - params_recarray : instance of numpy.ndarray | None + params_recarray The (N_models,)-shaped structured numpy ndarray holding the parameter values of the models. The the documentation of the :meth:`~skyllh.core.pdf.PDF.get_pd` method of the :class:`~skyllh.core.pdf.PDF` class for further information. - tl : instance of TimeLord | None + tl The optional instance of TimeLord to use for measuring timing information. Returns ------- - pd : instance of numpy.ndarray + pd The (N_events,)-shaped numpy ndarray holding the probability density for each event. In case of a signal PDF product the shape will be (N_sources,N_events). - grads : dict + grads The dictionary holding the gradients of the probability density w.r.t. each fit parameter. The key of the dictionary is the id of the global fit parameter. The value is the (N_events,)-shaped @@ -647,15 +654,15 @@ class SpatialPDF(PDF, metaclass=abc.ABCMeta): has two axes, right-ascension (ra) and declination (dec). """ - def __init__(self, ra_range, dec_range, **kwargs): + def __init__(self, ra_range: tuple, dec_range: tuple, **kwargs): """Constructor of a spatial PDF. It adds the PDF axes "ra" and "dec" with the specified ranges of coverage. Parameters ---------- - ra_range : 2-element tuple + ra_range The tuple specifying the right-ascension range this PDF covers. - dec_range : 2-element tuple + dec_range The tuple specifying the declination range this PDF covers. """ super().__init__(**kwargs) @@ -663,7 +670,7 @@ def __init__(self, ra_range, dec_range, **kwargs): self.add_axis(PDFAxis(name='ra', vmin=ra_range[0], vmax=ra_range[1])) self.add_axis(PDFAxis(name='dec', vmin=dec_range[0], vmax=dec_range[1])) - def assert_is_valid_for_trial_data(self, tdm, tl=None, **kwargs): + def assert_is_valid_for_trial_data(self, tdm: TrialDataManager, tl: TimeLord | None = None, **kwargs): """Checks if this spatial PDF is valid for all the given experimental data. It checks if all the data is within the right-ascension and declination @@ -671,7 +678,7 @@ def assert_is_valid_for_trial_data(self, tdm, tl=None, **kwargs): Parameters ---------- - tdm : instance of TrialDataManager + tdm The instance of TrialDataManager holding the trial data. The following data fields must exist: @@ -679,7 +686,7 @@ def assert_is_valid_for_trial_data(self, tdm, tl=None, **kwargs): The right-ascension of the data event. - 'dec' : float The declination of the data event. - tl : instance of TimeLord | None + tl The optional instance of TimeLord for measuring timing information. Raises @@ -709,6 +716,7 @@ class EnergyPDF(PDF, metaclass=abc.ABCMeta): """This is the abstract base class for an energy PDF.""" def __init__(self, *args, **kwargs): + """Creates a new instance of EnergyPDF.""" super().__init__(*args, **kwargs) @@ -722,7 +730,7 @@ class TimePDF(PDF, metaclass=abc.ABCMeta): def __init__( self, - livetime, + livetime: Livetime, time_flux_profile, **kwargs, ): @@ -731,10 +739,10 @@ def __init__( Parameters ---------- - livetime : instance of Livetime + livetime An instance of Livetime, which provides the detector live-time information. - time_profile : instance of TimeFluxProfile + time_flux_profile The signal's time flux profile. **kwargs Additional keyword arguments are passed to the constructor of the @@ -794,13 +802,13 @@ def __str__(self): return s - def _calculate_sum_of_ontime_time_flux_profile_integrals(self): + def _calculate_sum_of_ontime_time_flux_profile_integrals(self) -> float: """Calculates the sum, S, of the time flux profile integrals during the detector on-time intervals. Returns ------- - S : float + S The sum of the time flux profile integrals during the detector on-time intervals. """ @@ -812,21 +820,21 @@ def _calculate_sum_of_ontime_time_flux_profile_integrals(self): return S - def assert_is_valid_for_trial_data(self, tdm, tl=None, **kwargs): + def assert_is_valid_for_trial_data(self, tdm: TrialDataManager, tl: TimeLord | None = None, **kwargs): """Checks if the time PDF is valid for all the given trial data. It checks if the time of all events is within the defined time axis of the PDF. Parameters ---------- - tdm : instance of TrialDataManager + tdm The instance of TrialDataManager that holds the trial data. The following data fields must exist: ``'time'`` : float The time of the data event. - tl : instance of TimeLord | None + tl The optional instance of TimeLord for measuring timing information. Raises @@ -853,12 +861,12 @@ class MultiDimGridPDF( def __init__( self, - pmm, - axis_binnings, - path_to_pdf_splinetable=None, - pdf_grid_data=None, - norm_factor_func=None, - cache_pd_values=False, + pmm: ParameterModelMapper, + axis_binnings: BinningDefinition | Sequence[BinningDefinition], + path_to_pdf_splinetable: str | None = None, + pdf_grid_data: np.ndarray | None = None, + norm_factor_func: Callable | None = None, + cache_pd_values: bool = False, **kwargs, ): """Creates a new PDF instance for a multi-dimensional PDF given @@ -877,23 +885,23 @@ def __init__( Parameters ---------- - pmm : instance of ParameterModelMapper + pmm The instance of ParameterModelMapper that defines the mapping of the global parameters to local model parameters. - axis_binnings : instance of BinningDefinition | sequence of instance of BinningDefinition + axis_binnings The sequence of BinningDefinition instances defining the binning of the PDF axes. The name of each instance of BinningDefinition defines the event field name that should be used for querying the PDF. - path_to_pdf_splinetable : str | None + path_to_pdf_splinetable The path to the file containing the spline table, which contains a pre-computed fit to the grid data. If specified, ``pdf_grid_data`` must be ``None``. - pdf_grid_data : instance of numpy ndarray | None + pdf_grid_data The n-dimensional numpy ndarray holding the PDF values at given grid points. The grid points must match the bin edges of the given BinningDefinition instances of the ``axis_binnings`` argument. If specified, ``path_to_pdf_splinetable`` must be ``None``. - norm_factor_func : callable | None + norm_factor_func The function that calculates a possible required normalization factor for the PDF value based on the event properties. The call signature of this function must be @@ -909,7 +917,7 @@ def __init__( (N_values,)-shaped numpy ndarray holding the mask for the events, i.e. rows in ``eventdata``, which should be considered. If ``None``, all events should be considered. - cache_pd_values : bool + cache_pd_values Flag if the probability density values should be cached. The evaluation of the photospline fit might be slow and caching the probability density values might increase performance. @@ -1054,13 +1062,16 @@ def norm_factor_func(self, func): if func is None: # Define a normalization function that just returns 1 for each # event. - def func(pdf, tdm, params_recarray, eventdata, evt_mask=None): + def _unity_norm_func(pdf, tdm, params_recarray, eventdata, evt_mask=None): + """Returns a normalization factor of 1 for each event.""" if evt_mask is None: # noqa: SIM108 n_values = eventdata.shape[1] else: n_values = np.count_nonzero(evt_mask) return np.ones((n_values,), dtype=np.float64) + func = _unity_norm_func + if not callable(func): raise TypeError('The norm_factor_func property must be a callable object!') if not func_has_n_args(func, 5): @@ -1085,8 +1096,8 @@ def pdf(self): def assert_is_valid_for_trial_data( self, - tdm, - tl=None, + tdm: TrialDataManager, + tl: TimeLord | None = None, **kwargs, ): """Checks if the PDF is valid for all values of the given evaluation @@ -1095,10 +1106,10 @@ def assert_is_valid_for_trial_data( Parameters ---------- - tdm : instance of TrialDataManager + tdm The instance of TrialDataManager that holds the trial data for which the PDF should be valid. - tl : instance of TimeLord | None + tl The optional instance of TimeLord for measuring timing information. Raises @@ -1129,13 +1140,13 @@ def initialize_for_new_trial( def _initialize_cache( self, - tdm, + tdm: TrialDataManager, ): """Initializes the cache variables. Parameters ---------- - tdm : instance of TrialDataManager + tdm The instance of TrialDataManager that hold the trial data events. """ self._cache_tdm_trial_data_state_id = None @@ -1143,19 +1154,19 @@ def _initialize_cache( def _store_pd_values_to_cache( self, - tdm, - pd, - evt_mask=None, + tdm: TrialDataManager, + pd: np.ndarray, + evt_mask: np.ndarray | None = None, ): """Stores the given pd values into the pd array cache. Parameters ---------- - tdm : instance of TrialDataManager + tdm The instance of TrialDataManager that hold the trial data events. - pd : instance of numpy ndarray + pd The (N,)-shaped numpy ndarray holding the pd values to be stored. - evt_mask : instance of numpy ndarray | None + evt_mask The (N_values,)-shaped numpy ndarray defining the elements of the (N_values,)-shaped pd cache array where the given pd values should get stored. If set to ``None``, the the ``pd`` array must be of @@ -1172,14 +1183,14 @@ def _store_pd_values_to_cache( self._cache_pd[evt_mask] = pd - def _get_cached_pd_values(self, tdm, evt_mask=None): + def _get_cached_pd_values(self, tdm: TrialDataManager, evt_mask: np.ndarray | None = None) -> np.ndarray | None: """Retrieves cached pd values for the given events. Parameters ---------- - tdm : instance of TrialDataManager + tdm The instance of TrialDataManager that hold the trial data events. - evt_mask : instance of numpy ndarray | None + evt_mask The (N_values,)-shaped numpy ndarray defining the elements of the (N_values,)-shaped pd cache array for which pd values should get returned. @@ -1187,7 +1198,7 @@ def _get_cached_pd_values(self, tdm, evt_mask=None): Returns ------- - pd : instance of numpy ndarray | None + pd Returns ``None``, when no cached values are available. Otherwise the (N,)-shaped numpy ndarray holding the pd values where evt_mask evaluates to True. @@ -1202,6 +1213,7 @@ def _get_cached_pd_values(self, tdm, evt_mask=None): if evt_mask is None: pd = self._cache_pd else: + assert self._cache_pd is not None pd = self._cache_pd[evt_mask] # If this PDF is evaluated for different sources, i.e. a subset of # pd values, those values could still be NaN and still need to be @@ -1213,12 +1225,12 @@ def _get_cached_pd_values(self, tdm, evt_mask=None): def get_pd_with_eventdata( self, - tdm, - params_recarray, - eventdata, - evt_mask=None, - tl=None, - ): + tdm: TrialDataManager, + params_recarray: np.ndarray | None, + eventdata: np.ndarray, + evt_mask: np.ndarray | None = None, + tl: TimeLord | None = None, + ) -> np.ndarray: """Calculates the probability density value for the given ``eventdata``. This method is useful when PDF values for the same trial data need to @@ -1226,29 +1238,29 @@ def get_pd_with_eventdata( Parameters ---------- - tdm : instance of TrialDataManager + tdm The TrialDataManager instance holding the trial event data for which the PDF values should get calculated. - params_recarray : instance of numpy structured ndarray | None + params_recarray The (N_models,)-shaped numpy structured ndarray holding the local parameter names and values of the models. By definition, this PDF does not depend on any parameters. - eventdata : instance of numpy.ndarray + eventdata The (V,N_values)-shaped numpy ndarray holding the V data attributes for each of the N_values events needed for the evaluation of the PDF. - evt_mask : instance of numpy ndarray | None + evt_mask The (N_values,)-shaped numpy ndarray defining the elements of the N_values pd array for which pd values should get calculated. This is needed to determine if the requested pd values are already cached. - tl : instance of TimeLord | None + tl The optional instance of TimeLord that should be used to measure timing information. Returns ------- - pd : (N,)-shaped numpy ndarray + pd The (N,)-shaped numpy ndarray holding the probability density value for each model and event. The length of this array depends on the ``evt_mask`` argument. Only values are returned where @@ -1318,27 +1330,28 @@ def get_pd_with_eventdata( @staticmethod def create_eventdata_for_sigpdf( - tdm, - axes, - ): + tdm: TrialDataManager, + axes: 'PDFAxes', + ) -> np.ndarray: """Creates the (V,N_values)-shaped eventdata ndarray necessary for evaluating the signal PDF. Parameters ---------- - tdm : instance of TrialDataManager + tdm The instance of TrialDataManager holding the trial event data. - axes : instance of PDFAxes + axes The instance of PDFAxes defining the data field names for the PDF. Returns ------- - eventdata : instance of numpy.ndarray + eventdata The (V,N_values)-shaped numpy ndarray holding the event data for evaluating the signal PDF. """ eventdata_fields = [] + assert tdm.src_evt_idxs is not None (src_idxs, evt_idxs) = tdm.src_evt_idxs for axis in axes: name = axis.name @@ -1358,17 +1371,17 @@ def create_eventdata_for_sigpdf( @staticmethod def create_eventdata_for_bkgpdf( - tdm, - axes, + tdm: TrialDataManager, + axes: 'PDFAxes', ): """Creates the (V,N_values)-shaped eventdata ndarray necessary for evaluating the background PDF. Parameters ---------- - tdm : instance of TrialDataManager + tdm The instance of TrialDataManager holding the trial event data. - axes : instance of PDFAxes + axes The instance of PDFAxes defining the data field names for the PDF. """ eventdata_fields = [] @@ -1382,32 +1395,32 @@ def create_eventdata_for_bkgpdf( def get_pd( self, - tdm, - params_recarray=None, - tl=None, - ): + tdm: TrialDataManager, + params_recarray: np.ndarray | None = None, + tl: TimeLord | None = None, + ) -> tuple[np.ndarray, dict]: """Calculates the probability density for the given trial events given the specified local parameters. Parameters ---------- - tdm : instance of TrialDataManager + tdm The instance of TrialDataManager holding the trial event data for which the PDF values should get calculated. - params_recarray : instance of numpy structured ndarray | None + params_recarray The (N_models,)-shaped numpy structured ndarray holding the local parameter names and values of the models. By definition, this PDF does not depend on any parameters. - tl : instance of TimeLord | None + tl The optional instance of TimeLord that should be used to measure timing information. Returns ------- - pd : (N_values,)-shaped numpy ndarray + pd The (N_values,)-shaped numpy ndarray holding the probability density value for each source and event. - grads : dict + grads The dictionary holding the gradients of the probability density w.r.t. each global fit parameter. Since this PDF does not depend on any fit parameter, this is an empty dictionary. @@ -1418,10 +1431,11 @@ def get_pd( return (pd, {}) with TaskTimer(tl, 'Get PDF eventdata.'): + _axes = cast('PDFAxes', self._axes) if self.is_signal_pdf: - eventdata = self.create_eventdata_for_sigpdf(tdm=tdm, axes=self._axes) + eventdata = self.create_eventdata_for_sigpdf(tdm=tdm, axes=_axes) elif self.is_background_pdf: - eventdata = self.create_eventdata_for_bkgpdf(tdm=tdm, axes=self._axes) + eventdata = self.create_eventdata_for_bkgpdf(tdm=tdm, axes=_axes) else: raise TypeError('The PDF is neither a signal nor a background PDF!') @@ -1446,13 +1460,12 @@ class PDFSet( method and can be retrieved via the :meth:`get_pdf` method. """ - def __init__(self, param_grid_set, **kwargs): + def __init__(self, param_grid_set: ParameterGrid | ParameterGridSet, **kwargs): """Constructs a new PDFSet instance. Parameters ---------- - param_grid_set : instance of ParameterGrid | - instance of ParameterGridSet + param_grid_set The instance of ParameterGridSet with the parameter grids defining the discrete parameter values for which the PDFs of this PDF set are made for. @@ -1472,10 +1485,10 @@ def param_grid_set(self): return self._param_grid_set @param_grid_set.setter - def param_grid_set(self, obj): + def param_grid_set(self, obj: 'ParameterGrid | ParameterGridSet'): if isinstance(obj, ParameterGrid): obj = ParameterGridSet([obj]) - if obj is not None and not isinstance(obj, ParameterGridSet): + if not isinstance(obj, ParameterGridSet): raise TypeError('The params_grid_set property must be an instance of type ParameterGridSet!') self._param_grid_set = obj @@ -1499,12 +1512,12 @@ def axes(self): key = next(iter(self._gridparams_hash_pdf_dict.keys())) return self._gridparams_hash_pdf_dict[key].axes - def __contains__(self, key): + def __contains__(self, key: dict | int): """Checks if the given key exists in this PDFSet instance. Parameters ---------- - key : dict | int + key If a dictionary is provided, it must be the gridparams dictionary containing the grid parameter names and vales. If an integer is provided, it must be the hash of the gridparams @@ -1536,31 +1549,31 @@ def values(self): """Returns an iterator over the PDF instances of the PDFSet instance.""" return self._gridparams_hash_pdf_dict.values() - def make_key(self, gridparams): + def make_key(self, gridparams: dict) -> int: """Creates the key for the given grid parameter dictionary. Parameters ---------- - gridparams : dict + gridparams The dictionary holding the grid parameter names and values. Returns ------- - key : int + key The key for the given grid parameter dictionary. """ return make_dict_hash(gridparams) - def add_pdf(self, pdf, gridparams): + def add_pdf(self, pdf: 'PDF', gridparams: dict): """Adds the given PDF object for the given parameters to the internal registry. If this PDF set is not empty, the to-be-added PDF must have the same axes than the already added PDFs. Parameters ---------- - pdf : instance of PDF + pdf The PDF instance, that should be added - gridparams : dict + gridparams The dictionary with the grid parameter values, which identify the PDF object. @@ -1588,7 +1601,7 @@ def add_pdf(self, pdf, gridparams): # Check that the new PDF has the same axes than the already added PDFs. if len(self._gridparams_hash_pdf_dict) > 0: some_pdf = self._gridparams_hash_pdf_dict[next(iter(self._gridparams_hash_pdf_dict.keys()))] - if not pdf.axes.is_same_as(some_pdf.axes): + if not cast('PDFAxes', pdf.axes).is_same_as(cast('PDFAxes', some_pdf.axes)): raise ValueError( 'The given PDF does not have the same axes than the ' 'already added PDFs!\n' @@ -1601,19 +1614,19 @@ def add_pdf(self, pdf, gridparams): self._gridparams_hash_pdf_dict[gridparams_hash] = pdf - def get_pdf(self, gridparams): + def get_pdf(self, gridparams: dict | int): """Retrieves the PDF object for the given set of fit parameters. Parameters ---------- - gridparams : dict | int + gridparams The dictionary with the grid parameter names and values for which the PDF object should get retrieved. If an integer is given, it is assumed to be the PDF key. Returns ------- - pdf : instance if PDF + pdf The PDF instance for the given parameters. Raises @@ -1635,22 +1648,22 @@ def get_pdf(self, gridparams): return pdf - def initialize_for_new_trial(self, tdm, tl=None, **kwargs): + def initialize_for_new_trial(self, tdm: TrialDataManager, tl: TimeLord | None = None, **kwargs): """This method is called whenever a new trial data is available. It calls the :meth:`~skyllh.core.pdf.PDF.initialize_for_new_trial` method of each PDF. Parameters ---------- - tdm : instance of TrialDataManager + tdm The instance of TrialDataManager holding the new trial data events. - tl : instance of TimeLord | None + tl The optional instance of TimeLord for measuring timing information. """ for pdf in self._gridparams_hash_pdf_dict.values(): pdf.initialize_for_new_trial(tdm=tdm, tl=tl, **kwargs) - def assert_is_valid_for_trial_data(self, tdm, tl=None, **kwargs): + def assert_is_valid_for_trial_data(self, tdm: TrialDataManager, tl: TimeLord | None = None, **kwargs): """Checks if the PDFs of this PDFSet instance are valid for all the given trial data events. Since all PDFs should have the same axes, only the first PDF will be @@ -1660,9 +1673,9 @@ def assert_is_valid_for_trial_data(self, tdm, tl=None, **kwargs): Parameters ---------- - tdm : instance of TrialDataManager + tdm The instance of TrialDataManager holding the trial data events. - tl : instance of TimeLord | None + tl The optional instance of TimeLord for measuring timing information. Raises @@ -1674,35 +1687,39 @@ def assert_is_valid_for_trial_data(self, tdm, tl=None, **kwargs): pdf = self._gridparams_hash_pdf_dict[key] pdf.assert_is_valid_for_trial_data(tdm=tdm, tl=tl, **kwargs) - def get_pd(self, gridparams, tdm, params_recarray=None, tl=None): + def get_pd( + self, gridparams: dict, tdm: TrialDataManager, params_recarray: np.ndarray | None = None, tl=None + ) -> 'tuple[np.ndarray, dict]': """Calls the ``get_pd`` method of the PDF instance that belongs to the given grid parameter values ``gridparams``. Parameters ---------- - gridparams : dict + gridparams The dictionary holding the parameter values, which define PDF instance within this PDFSet instance. Note, that the parameter values must match a set of parameter grid values for which a PDF instance has been created and added to this PDFSet instance. - tdm : instance of TrialDataManager + tdm The TrialDataManager instance holding the data events for which the probability density of the events should be calculated. - params_recarray : instance of ndarray | None + params_recarray The numpy record ndarray holding the parameter name and values for each source model. + tl + The optional instance of TimeLord for measuring timing information. Returns ------- - pd : numpy ndarray + pd The 1D numpy ndarray holding the probability density values for each event and source. - See :meth:`skyllh.core.pdf.PDF.get_pd` for further information. - grads : dict + See + grads The dictionary holding the gradient values for each global fit parameter. - See :meth:`skyllh.core.pdf.PDF.get_pd` for further information. + See """ pdf = self.get_pdf(gridparams) diff --git a/skyllh/core/pdfratio.py b/skyllh/core/pdfratio.py index dbd9af0190..5ff60f788f 100644 --- a/skyllh/core/pdfratio.py +++ b/skyllh/core/pdfratio.py @@ -1,4 +1,6 @@ import abc +from collections.abc import Sequence +from typing import cast import numpy as np @@ -13,6 +15,7 @@ ParameterModelMapper, ) from skyllh.core.pdf import ( + PDF, IsBackgroundPDF, IsSignalPDF, PDFSet, @@ -27,9 +30,8 @@ from skyllh.core.services import ( SrcDetSigYieldWeightsService, ) -from skyllh.core.timing import ( - TaskTimer, -) +from skyllh.core.timing import TaskTimer, TimeLord +from skyllh.core.trialdata import TrialDataManager class PDFRatio( @@ -40,20 +42,23 @@ class PDFRatio( It defines the interface of a signal over background PDF ratio class. """ + _sig_param_names: list[str] + _bkg_param_names: list[str] + def __init__( self, - sig_param_names=None, - bkg_param_names=None, + sig_param_names: Sequence[str] | str | None = None, + bkg_param_names: Sequence[str] | str | None = None, **kwargs, ): """Creates a new PDFRatio instance. Parameters ---------- - sig_param_names : sequence of str | str | None + sig_param_names The sequence of signal parameter names this PDFRatio instance is a function of. - bkg_param_names : sequence of str | str | None + bkg_param_names The sequence of background parameter names this PDFRatio instance is a function of. """ @@ -92,7 +97,7 @@ def n_bkg_params(self): return len(self._bkg_param_names) @property - def sig_param_names(self): + def sig_param_names(self) -> list[str]: """The list of signal parameter names this PDF ratio is a function of.""" return self._sig_param_names @@ -104,7 +109,7 @@ def sig_param_names(self, names): names = [names] if not issequenceof(names, str): raise TypeError('The sig_param_names property must be a sequence of str instances!') - self._sig_param_names = names + self._sig_param_names = cast(list[str], names) @property def bkg_param_names(self): @@ -121,13 +126,13 @@ def bkg_param_names(self, names): names = [names] if not issequenceof(names, str): raise TypeError('The bkg_param_names property must be a sequence of str instances!') - self._bkg_param_names = names + self._bkg_param_names = cast(list[str], names) @abc.abstractmethod def initialize_for_new_trial( self, - tdm, - tl=None, + tdm: TrialDataManager, + tl: TimeLord | None = None, **kwargs, ): """Initializes the PDFRatio instance for a new trial. This method can @@ -136,40 +141,40 @@ def initialize_for_new_trial( Parameters ---------- - tdm : instance of TrialDataManager + tdm The instance of TrialDataManager that holds the trial data. - tl : instance of TimeLord + tl The optional instance of TimeLord to measure timing information. """ @abc.abstractmethod def get_ratio( self, - tdm, - src_params_recarray, - tl=None, - ): + tdm: TrialDataManager, + src_params_recarray: np.ndarray, + tl: TimeLord | None = None, + ) -> np.ndarray: """Retrieves the PDF ratio value for each given trial data events (and sources), given the given set of parameters. Parameters ---------- - tdm : instance of TrialDataManager + tdm The TrialDataManager instance holding the trial data events for which the PDF ratio values should get calculated. - src_params_recarray : instance of numpy record ndarray | None + src_params_recarray The (N_sources,)-shaped numpy record ndarray holding the parameter names and values of the sources. See the documentation of the :meth:`skyllh.core.parameters.ParameterModelMapper.create_src_params_recarray` method for more information. - tl : instance of TimeLord | None + tl The optional TimeLord instance that should be used to measure timing information. Returns ------- - ratios : instance of ndarray + ratios The (N_values,)-shaped 1d numpy ndarray of float holding the PDF ratio value for each trial event and source. """ @@ -177,36 +182,36 @@ def get_ratio( @abc.abstractmethod def get_gradient( self, - tdm, - src_params_recarray, - fitparam_id, - tl=None, - ): + tdm: TrialDataManager, + src_params_recarray: np.ndarray, + fitparam_id: int, + tl: TimeLord | None = None, + ) -> np.ndarray: """Retrieves the PDF ratio gradient for the global fit parameter ``fitparam_id`` for each trial data event and source, given the given set of parameters ``src_params_recarray`` for each source. Parameters ---------- - tdm : instance of TrialDataManager + tdm The TrialDataManager instance holding the trial data events for which the PDF ratio gradient values should get calculated. - src_params_recarray : instance of numpy structured ndarray + src_params_recarray The (N_sources,)-shaped numpy structured ndarray holding the parameter names and values of the sources. See the documentation of the :meth:`skyllh.core.parameters.ParameterModelMapper.create_src_params_recarray` method for more information. - fitparam_id : int + fitparam_id The ID of the global fit parameter for which the gradient should get calculated. - tl : instance of TimeLord | None + tl The optional TimeLord instance that should be used to measure timing information. Returns ------- - gradient : instance of ndarray | 0 + gradient The (N_values,)-shaped 1d numpy ndarray of float holding the PDF ratio gradient value for each source and trial data event. If the PDF ratio does not depend on the given global fit parameter, @@ -230,8 +235,8 @@ class PDFRatioProduct( def __init__( self, - pdfratio1, - pdfratio2, + pdfratio1: PDFRatio, + pdfratio2: PDFRatio, **kwargs, ): """Creates a new PDFRatioProduct instance representing the product of @@ -240,8 +245,8 @@ def __init__( self.pdfratio1 = pdfratio1 self.pdfratio2 = pdfratio2 - sig_param_names = set(list(pdfratio1.sig_param_names) + list(pdfratio2.sig_param_names)) - bkg_param_names = set(list(pdfratio1.bkg_param_names) + list(pdfratio2.bkg_param_names)) + sig_param_names = list(set(list(pdfratio1.sig_param_names) + list(pdfratio2.sig_param_names))) + bkg_param_names = list(set(list(pdfratio1.bkg_param_names) + list(pdfratio2.bkg_param_names))) super().__init__(sig_param_names=sig_param_names, bkg_param_names=bkg_param_names, **kwargs) @@ -271,42 +276,42 @@ def pdfratio2(self, pdfratio): raise TypeError('The pdfratio2 property must be an instance of PDFRatio!') self._pdfratio2 = pdfratio - def initialize_for_new_trial(self, **kwargs): + def initialize_for_new_trial(self, tdm: TrialDataManager, tl: TimeLord | None = None, **kwargs): """Initializes the PDFRatioProduct instance for a new trial. It calls the :meth:`~skyllh.core.pdfratio.PDFRatio.initialize_for_new_trial` method of each of the two :class:`~skyllh.core.pdfratio.PDFRatio` instances. """ - self._pdfratio1.initialize_for_new_trial(**kwargs) - self._pdfratio2.initialize_for_new_trial(**kwargs) + self._pdfratio1.initialize_for_new_trial(tdm=tdm, tl=tl, **kwargs) + self._pdfratio2.initialize_for_new_trial(tdm=tdm, tl=tl, **kwargs) def get_ratio( self, - tdm, - src_params_recarray, - tl=None, - ): + tdm: TrialDataManager, + src_params_recarray: np.ndarray, + tl: TimeLord | None = None, + ) -> np.ndarray: """Retrieves the PDF ratio product value for each trial data event and source, given the given set of parameters for all sources. Parameters ---------- - tdm : instance of TrialDataManager + tdm The TrialDataManager instance holding the trial data events for which the PDF ratio values should get calculated. - src_params_recarray : instance of numpy record ndarray + src_params_recarray The (N_sources,)-shaped numpy record ndarray holding the parameter names and values of the sources. See the documentation of the :meth:`skyllh.core.parameters.ParameterModelMapper.create_src_params_recarray` method for more information. - tl : TimeLord instance | None + tl The optional TimeLord instance that should be used to measure timing information. Returns ------- - ratios : instance of ndarray + ratios The (N_values,)-shaped 1d numpy ndarray of float holding the product of the PDF ratio values for each trial event and source. The PDF ratio product value for each trial event. @@ -319,41 +324,42 @@ def get_ratio( def get_gradient( self, - tdm, - src_params_recarray, - fitparam_id, - tl=None, - ): + tdm: TrialDataManager, + src_params_recarray: np.ndarray | None, + fitparam_id: int, + tl: TimeLord | None = None, + ) -> np.ndarray: """Retrieves the PDF ratio product gradient for the global fit parameter with parameter ID ``fitparam_id`` for each trial data event and source, given the set of parameters ``src_params_recarray`` for all sources. Parameters ---------- - tdm : instance of TrialDataManager + tdm The TrialDataManager instance holding the trial data events for which the PDF ratio values should get calculated. - src_params_recarray : instance of numpy record ndarray | None + src_params_recarray The (N_sources,)-shaped numpy record ndarray holding the parameter names and values of the sources. See the documentation of the :meth:`skyllh.core.parameters.ParameterModelMapper.create_src_params_recarray` method for more information. - fitparam_id : int + fitparam_id The ID of the global fit parameter for which the gradient should get calculated. - tl : instance of TimeLord | None + tl The optional TimeLord instance that should be used to measure timing information. Returns ------- - gradient : instance of ndarray | 0 + gradient The (N_values,)-shaped 1d numpy ndarray of float holding the PDF ratio gradient value for each trial event and source. If none of the two PDFRatio instances depend on the given global fit parameter, the scalar value ``0`` is returned. """ + assert src_params_recarray is not None r1_depends_on_fitparam = ParameterModelMapper.is_global_fitparam_a_local_param( fitparam_id=fitparam_id, params_recarray=src_params_recarray, local_param_names=self._pdfratio1.param_names ) @@ -362,29 +368,30 @@ def get_gradient( fitparam_id=fitparam_id, params_recarray=src_params_recarray, local_param_names=self._pdfratio2.param_names ) - if r1_depends_on_fitparam: + if r1_depends_on_fitparam and r2_depends_on_fitparam: r2 = self._pdfratio2.get_ratio(tdm=tdm, src_params_recarray=src_params_recarray, tl=tl) - r1_grad = self._pdfratio1.get_gradient( tdm=tdm, src_params_recarray=src_params_recarray, fitparam_id=fitparam_id, tl=tl ) - - if r2_depends_on_fitparam: r1 = self._pdfratio1.get_ratio(tdm=tdm, src_params_recarray=src_params_recarray, tl=tl) - r2_grad = self._pdfratio2.get_gradient( tdm=tdm, src_params_recarray=src_params_recarray, fitparam_id=fitparam_id, tl=tl ) - - if r1_depends_on_fitparam and r2_depends_on_fitparam: - gradient = r1 * r2_grad - gradient += r1_grad * r2 + gradient: np.ndarray = r1 * r2_grad + r1_grad * r2 elif r1_depends_on_fitparam: + r2 = self._pdfratio2.get_ratio(tdm=tdm, src_params_recarray=src_params_recarray, tl=tl) + r1_grad = self._pdfratio1.get_gradient( + tdm=tdm, src_params_recarray=src_params_recarray, fitparam_id=fitparam_id, tl=tl + ) gradient = r1_grad * r2 elif r2_depends_on_fitparam: + r1 = self._pdfratio1.get_ratio(tdm=tdm, src_params_recarray=src_params_recarray, tl=tl) + r2_grad = self._pdfratio2.get_gradient( + tdm=tdm, src_params_recarray=src_params_recarray, fitparam_id=fitparam_id, tl=tl + ) gradient = r1 * r2_grad else: - gradient = 0 + gradient = np.zeros(tdm.n_selected_events, dtype=np.float64) return gradient @@ -401,19 +408,25 @@ class SourceWeightedPDFRatio(PDFRatio): """ - def __init__(self, dataset_idx, src_detsigyield_weights_service, pdfratio, **kwargs): + def __init__( + self, + dataset_idx: int, + src_detsigyield_weights_service: SrcDetSigYieldWeightsService, + pdfratio: 'PDFRatio', + **kwargs, + ): """Creates a new SourceWeightedPDFRatio instance. Parameters ---------- - dataset_idx : int + dataset_idx The index of the dataset. It is used to access the source detector signal yield weight. - src_detsigyield_weights_service : instance of SrcDetSigYieldWeightsService + src_detsigyield_weights_service The instance of SrcDetSigYieldWeightsService providing the source detector signal yield weights, i.e. the product of the theoretical source weight with the detector signal yield. - pdfratio : instance of PDFRatio + pdfratio The instance of PDFRatio providing the PDF ratio values and derivatives. """ @@ -463,21 +476,23 @@ def pdfratio(self): """ return self._pdfratio - def initialize_for_new_trial(self, tdm, tl=None, **kwargs): + def initialize_for_new_trial(self, tdm: TrialDataManager, tl: TimeLord | None = None, **kwargs): """Initializes the PDFRatio instance for a new trial. It calls the :meth:`~skyllh.core.pdfratio.PDFRatio.initialize_for_new_trial` method of the :class:`~skyllh.core.pdfratio.PDFRatio` instance. Parameters ---------- - tdm : instance of TrialDataManager + tdm The instance of TrialDataManager that holds the trial data. - tl : instance of TimeLord + tl The optional instance of TimeLord to measure timing information. """ self._pdfratio.initialize_for_new_trial(tdm=tdm, tl=tl, **kwargs) - def get_ratio(self, tdm, src_params_recarray, tl=None): + def get_ratio( + self, tdm: TrialDataManager, src_params_recarray: np.ndarray, tl: TimeLord | None = None + ) -> np.ndarray: """Retrieves the PDF ratio value for each given trial data events (and sources), given the given set of parameters. @@ -490,26 +505,27 @@ def get_ratio(self, tdm, src_params_recarray, tl=None): Parameters ---------- - tdm : instance of TrialDataManager + tdm The TrialDataManager instance holding the trial data events for which the PDF ratio values should get calculated. - src_params_recarray : instance of numpy record ndarray + src_params_recarray The (N_sources,)-shaped numpy record ndarray holding the parameter names and values of the sources. See the documentation of the :meth:`skyllh.core.parameters.ParameterModelMapper.create_src_params_recarray` method for more information. - tl : instance of TimeLord | None + tl The optional TimeLord instance that should be used to measure timing information. Returns ------- - ratios : instance of ndarray + ratios The (N_selected_events,)-shaped 1d numpy ndarray of float holding the PDF ratio value for each selected trial data event. """ (a_jk, _) = self._src_detsigyield_weights_service.get_weights() + assert a_jk is not None a_k = a_jk[self._dataset_idx] n_sources = len(a_k) @@ -522,6 +538,7 @@ def get_ratio(self, tdm, src_params_recarray, tl=None): R_i = np.zeros((n_sel_events,), dtype=np.double) + assert tdm.src_evt_idxs is not None (src_idxs, evt_idxs) = tdm.src_evt_idxs for k in range(n_sources): src_mask = src_idxs == k @@ -533,7 +550,9 @@ def get_ratio(self, tdm, src_params_recarray, tl=None): return R_i - def get_gradient(self, tdm, src_params_recarray, fitparam_id, tl=None): + def get_gradient( + self, tdm: TrialDataManager, src_params_recarray: np.ndarray, fitparam_id: int, tl: TimeLord | None = None + ) -> np.ndarray: """Retrieves the PDF ratio gradient for the parameter ``fitparam_id`` for each trial data event, given the given set of parameters ``src_params_recarray`` for each source. @@ -545,31 +564,34 @@ def get_gradient(self, tdm, src_params_recarray, fitparam_id, tl=None): Parameters ---------- - tdm : instance of TrialDataManager + tdm The TrialDataManager instance holding the trial data events for which the PDF ratio gradient values should get calculated. - src_params_recarray : instance of numpy record ndarray + src_params_recarray The (N_sources,)-shaped numpy record ndarray holding the parameter names and values of the sources. See the documentation of the :meth:`skyllh.core.parameters.ParameterModelMapper.create_src_params_recarray` method for more information. - fitparam_id : int + fitparam_id The ID of the global fit parameter for which the gradient should get calculated. - tl : instance of TimeLord | None + tl The optional TimeLord instance that should be used to measure timing information. Returns ------- - gradient : instance of ndarray | 0 + gradient The (N_selected_events,)-shaped 1d numpy ndarray of float holding the PDF ratio gradient value for each trial data event. If the PDF ratio does not depend on the given global fit parameter, 0 will be returned. """ (a_jk, a_jk_grads) = self._src_detsigyield_weights_service.get_weights() + assert a_jk is not None + assert a_jk_grads is not None + assert self._cache_R_i is not None a_k = a_jk[self._dataset_idx] A = np.sum(a_k) @@ -590,17 +612,19 @@ def get_gradient(self, tdm, src_params_recarray, fitparam_id, tl=None): # R_ik_grad is a (N_values,)-shaped ndarray or 0. if (isinstance(a_k_grad, int) and a_k_grad == 0) and (isinstance(R_ik_grad, int) and R_ik_grad == 0): - return 0 + return np.zeros(n_sel_events, dtype=np.float64) R_i_grad = -self._cache_R_i * dAdp src_sum_i = np.zeros((n_sel_events,), dtype=np.double) + assert tdm.src_evt_idxs is not None (src_idxs, evt_idxs) = tdm.src_evt_idxs for k in range(n_sources): src_mask = src_idxs == k src_evt_idxs = evt_idxs[src_mask] if isinstance(a_k_grad, np.ndarray): + assert self._cache_R_ik is not None src_sum_i[src_evt_idxs] += a_k_grad[k] * self._cache_R_ik[src_mask] if isinstance(R_ik_grad, np.ndarray): src_sum_i[src_evt_idxs] += a_k[k] * R_ik_grad[src_mask] @@ -618,19 +642,19 @@ class SigOverBkgPDFRatio(PDFRatio): *pdf_type* and calculates the PDF ratio. """ - def __init__(self, sig_pdf, bkg_pdf, same_axes=True, zero_bkg_ratio_value=1.0, **kwargs): + def __init__(self, sig_pdf, bkg_pdf, same_axes: bool = True, zero_bkg_ratio_value: float = 1.0, **kwargs): """Creates a new signal-over-background PDF ratio instance. Parameters ---------- - sig_pdf : class instance derived from `pdf_type`, IsSignalPDF + sig_pdf The instance of the signal PDF. - bkg_pdf : class instance derived from `pdf_type`, IsBackgroundPDF + bkg_pdf The instance of the background PDF. - same_axes : bool + same_axes Flag if the signal and background PDFs are supposed to have the same axes. Default is True. - zero_bkg_ratio_value : float + zero_bkg_ratio_value The value of the PDF ratio to take when the background PDF value is zero. This is to avoid division by zero. Default is 1. """ @@ -697,42 +721,46 @@ def initialize_for_new_trial(self, tdm, tl=None, **kwargs): :meth:`~skyllh.core.pdf.PDF.assert_is_valid_for_trial_data` of the signal and background :class:`~skyllh.core.pdf.PDF` instances. """ - self._sig_pdf.initialize_for_new_trial(tdm=tdm, tl=tl, **kwargs) - self._sig_pdf.assert_is_valid_for_trial_data(tdm=tdm, tl=tl, **kwargs) + cast(PDF, self._sig_pdf).initialize_for_new_trial(tdm=tdm, tl=tl, **kwargs) + cast(PDF, self._sig_pdf).assert_is_valid_for_trial_data(tdm=tdm, tl=tl, **kwargs) - self._bkg_pdf.initialize_for_new_trial(tdm=tdm, tl=tl, **kwargs) - self._bkg_pdf.assert_is_valid_for_trial_data(tdm=tdm, tl=tl, **kwargs) + cast(PDF, self._bkg_pdf).initialize_for_new_trial(tdm=tdm, tl=tl, **kwargs) + cast(PDF, self._bkg_pdf).assert_is_valid_for_trial_data(tdm=tdm, tl=tl, **kwargs) - def get_ratio(self, tdm, src_params_recarray, tl=None): + def get_ratio( + self, tdm: TrialDataManager, src_params_recarray: np.ndarray, tl: TimeLord | None = None + ) -> np.ndarray: """Calculates the PDF ratio for the given trial events. Parameters ---------- - tdm : instance of TrialDataManager + tdm The TrialDataManager instance holding the trial data events for which the PDF ratio values should be calculated. - src_params_recarray : instance of numpy record ndarray + src_params_recarray The (N_sources,)-shaped numpy record ndarray holding the local parameter names and values of the sources. See the :meth:`skyllh.core.parameters.ParameterModelMapper.create_src_params_recarray` method for more information. - tl : instance of TimeLord | None + tl The optional TimeLord instance that should be used to measure timing information. Returns ------- - ratios : instance of ndarray + ratios The (N_values,)-shaped numpy ndarray holding the probability density ratio for each event and source. """ with TaskTimer(tl, 'Get sig probability densities and grads.'): - (self._cache_sig_pd, self._cache_sig_grads) = self._sig_pdf.get_pd( + (self._cache_sig_pd, self._cache_sig_grads) = cast(PDF, self._sig_pdf).get_pd( tdm=tdm, params_recarray=src_params_recarray, tl=tl ) with TaskTimer(tl, 'Get bkg probability densities and grads.'): - (self._cache_bkg_pd, self._cache_bkg_grads) = self._bkg_pdf.get_pd(tdm=tdm, params_recarray=None, tl=tl) + (self._cache_bkg_pd, self._cache_bkg_grads) = cast(PDF, self._bkg_pdf).get_pd( + tdm=tdm, params_recarray=None, tl=tl + ) with TaskTimer(tl, 'Calculate PDF ratios.'): # Select only the events, where the background pdf is greater than @@ -744,7 +772,13 @@ def get_ratio(self, tdm, src_params_recarray, tl=None): return ratios - def get_gradient(self, tdm, src_params_recarray, fitparam_id, tl=None): + def get_gradient( + self, + tdm: TrialDataManager, + src_params_recarray: np.ndarray | None, + fitparam_id: int, + tl: TimeLord | None = None, + ) -> np.ndarray: """Retrieves the gradient of the PDF ratio w.r.t. the given parameter. Note: @@ -755,27 +789,32 @@ def get_gradient(self, tdm, src_params_recarray, fitparam_id, tl=None): Parameters ---------- - tdm : instance of TrialDataManager + tdm The instance of TrialDataManager holding the trial data. - src_params_recarray : instance of numpy record ndarray | None + src_params_recarray The (N_models,)-shaped numpy record ndarray holding the parameter names and values of the models. - See :meth:`skyllh.core.pdf.PDF.get_pd` for more information. + See This can be ``None``, if the signal and background PDFs do not depend on any parameters. - fitparam_id : int + fitparam_id The ID of the global fit parameter for which the gradient should get calculated. - tl : instance of TimeLord | None + tl The optional TimeLord instance that should be used to measure timing information. Returns ------- - grad : instance of ndarray + grad The (N_values,)-shaped 1d numpy ndarray of float holding the PDF ratio gradient value for each source and trial event. """ + assert self._cache_sig_pd is not None + assert self._cache_sig_grads is not None + assert self._cache_bkg_pd is not None + assert self._cache_bkg_grads is not None + # Create the 1D return array for the gradient. grad = np.zeros_like(self._cache_sig_pd, dtype=np.float64) @@ -830,18 +869,20 @@ class SigSetOverBkgPDFRatio(PDFRatio): values. """ - def __init__(self, sig_pdf_set, bkg_pdf, interpolmethod_cls=None, **kwargs): + def __init__( + self, sig_pdf_set, bkg_pdf, interpolmethod_cls: type[GridManifoldInterpolationMethod] | None = None, **kwargs + ): """Constructor called by creating an instance of a class which is derived from this PDFRatio class. Parameters ---------- - sig_pdf_set : instance of PDFSet and instance of IsSignalPDF + sig_pdf_set The PDF set, which provides signal PDFs for a set of discrete signal parameter values. - bkg_pdf : instance of PDF and instance of IsBackgroundPDF + bkg_pdf The background PDF instance. - interpolmethod_cls : class of GridManifoldInterpolationMethod | None + interpolmethod_cls The class implementing the parameter interpolation method for the PDF ratio manifold grid. If set to ``None`` (default), the :class:`skyllh.core.interpolate.Parabola1DGridManifoldInterpolationMethod` @@ -917,7 +958,7 @@ def interpolmethod_cls(self, cls): ) self._interpolmethod_cls = cls - def initialize_for_new_trial(self, tdm, tl=None, **kwargs): + def initialize_for_new_trial(self, tdm: TrialDataManager, tl: TimeLord | None = None, **kwargs): """Initializes the PDFRatio instance for a new trial. It calls the :meth:`~skyllh.core.pdf.PDF.assert_is_valid_for_trial_data` of the signal :class:`~skyllh.core.pdf.PDFSet` instance and the background @@ -925,13 +966,13 @@ def initialize_for_new_trial(self, tdm, tl=None, **kwargs): Parameters ---------- - tdm : instance of TrialDataManager + tdm The instance of TrialDataManager holding the trial data. - tl : instance of TimeLord | None + tl The optional instance of TimeLord for measuring timing information. """ self._sig_pdf_set.initialize_for_new_trial(tdm=tdm, tl=tl, **kwargs) self._sig_pdf_set.assert_is_valid_for_trial_data(tdm=tdm, tl=tl, **kwargs) - self._bkg_pdf.initialize_for_new_trial(tdm=tdm, tl=tl, **kwargs) - self._bkg_pdf.assert_is_valid_for_trial_data(tdm=tdm, tl=tl, **kwargs) + cast(PDF, self._bkg_pdf).initialize_for_new_trial(tdm=tdm, tl=tl, **kwargs) + cast(PDF, self._bkg_pdf).assert_is_valid_for_trial_data(tdm=tdm, tl=tl, **kwargs) diff --git a/skyllh/core/pdfratio_fill.py b/skyllh/core/pdfratio_fill.py index d6066c0de8..7092630987 100644 --- a/skyllh/core/pdfratio_fill.py +++ b/skyllh/core/pdfratio_fill.py @@ -21,19 +21,20 @@ class PDFRatioFillMethod(metaclass=abc.ABCMeta): """ def __init__(self, *args, **kwargs): + """Creates a new instance of PDFRatioFillMethod.""" super().__init__(*args, **kwargs) @abc.abstractmethod def __call__( self, - ratios, - sig_pd_h, - bkg_pd_h, - sig_mask_mc_covered, - sig_mask_mc_covered_zero_physics, - bkg_mask_mc_covered, - bkg_mask_mc_covered_zero_physics, - ): + ratios: np.ndarray, + sig_pd_h: np.ndarray, + bkg_pd_h: np.ndarray, + sig_mask_mc_covered: np.ndarray, + sig_mask_mc_covered_zero_physics: np.ndarray, + bkg_mask_mc_covered: np.ndarray, + bkg_mask_mc_covered_zero_physics: np.ndarray, + ) -> np.ndarray: """The __call__ method is supposed to fill the ratio bins (array) with the signal / background ratio values. For bins (array elements), where the division is undefined, e.g. due to zero background, the fill @@ -45,27 +46,27 @@ def __call__( Parameters ---------- - ratios : ndarray of float + ratios The multi-dimensional array for the final ratio bins. The shape is the same as the sig_h and bkg_h ndarrays. - sig_pd_h : ndarray of float + sig_pd_h The multi-dimensional array (histogram) holding the signal probability densities. - bkg_pd_h : ndarray of float + bkg_pd_h The multi-dimensional array (histogram) holding the background probability densities. - sig_mask_mc_covered : ndarray of bool + sig_mask_mc_covered The mask array indicating which array elements of sig_pd_h have monte-carlo coverage. - sig_mask_mc_covered_zero_physics : ndarray of bool + sig_mask_mc_covered_zero_physics The mask array indicating which array elements of sig_pd_h have monte-carlo coverage but don't have physics contribution. - bkg_mask_mc_covered : ndarray of bool + bkg_mask_mc_covered The mask array indicating which array elements of bkg_pd_h have monte-carlo coverage. In case of experimental data as background, this mask indicate where (experimental data) background is available. - bkg_mask_mc_covered_zero_physics : ndarray of bool + bkg_mask_mc_covered_zero_physics The mask array ndicating which array elements of bkg_pd_h have monte-carlo coverage but don't have physics contribution. In case of experimental data as background, this mask contains only @@ -73,7 +74,7 @@ def __call__( Returns ------- - ratios : ndarray + ratios The array holding the final ratio values. """ return ratios @@ -88,6 +89,7 @@ class Skylab2SkylabPDFRatioFillMethod(PDFRatioFillMethod): """ def __init__(self, **kwargs): + """Creates a new instance of Skylab2SkylabPDFRatioFillMethod.""" super().__init__(**kwargs) self.signallike_percentile = 99.0 @@ -149,14 +151,14 @@ class MostSignalLikePDFRatioFillMethod(PDFRatioFillMethod): ratio for bins, where there is signal but no background coverage. """ - def __init__(self, signallike_percentile=99.0, **kwargs): + def __init__(self, signallike_percentile: float = 99.0, **kwargs): """Creates the PDF ratio fill method object for filling PDF ratio bins, where there is signal MC coverage but no background (MC) coverage with the most signal-like ratio value. Parameters ---------- - signallike_percentile : float in range [0., 100.], default 99. + signallike_percentile The percentile of signal-like ratios, which should be taken as the ratio value for ratios with no background probability. """ diff --git a/skyllh/core/progressbar.py b/skyllh/core/progressbar.py index 5208b99614..4281c4f447 100644 --- a/skyllh/core/progressbar.py +++ b/skyllh/core/progressbar.py @@ -16,16 +16,16 @@ class ProgressBar: are updated. """ - def __init__(self, maxval, startval=0, parent=None, **kwargs): + def __init__(self, maxval: int, startval: int = 0, parent: 'ProgressBar | None' = None, **kwargs): """Creates a new ProgressBar instance. Parameters ---------- - maxval : int + maxval The maximal value the progress can reach. - startval : int + startval The progress value to start with. Must be smaller than `maxval`. - parent : instance of ProgressBar | False | None + parent The parent instance of ProgressBar if this progress bar is a sub progress bar. If set to ``False``, this progress bar is deactivated and the @@ -129,12 +129,12 @@ def remove_sub_progress_bars(self): self._sub_pbar_list = [] - def get_progressbar_list(self): + def get_progressbar_list(self) -> 'list[ProgressBar]': """Retrieves the list of ProgressBar instances. Returns ------- - pbar_list : list of instance of ProgressBar + pbar_list The list of ProgressBar instances, which are part of this ProgressBar instance. """ @@ -161,6 +161,7 @@ def rerender(self): maxval += pbar.maxval val += pbar.val + assert self._tqdm is not None dval = val - self._tqdm.n self._tqdm.total = maxval self._tqdm.update(dval) @@ -185,6 +186,7 @@ def start(self): elif not self.is_shown: return self else: + assert self._tqdm is not None self._tqdm.initial = self._val self._tqdm.n = self._val self._tqdm.reset() @@ -204,27 +206,28 @@ def finish(self): self.trigger_rerendering() if (self._parent is None) and self.is_shown: + assert self._tqdm is not None self._tqdm.close() self.remove_sub_progress_bars() - def increment(self, dval=1): + def increment(self, dval: int = 1): """Updates the progress bar by incrementing the progress by the given integral amount. Parameters ---------- - dval : int + dval The amount of progress to increment the progress bar with. """ self.update(self._val + dval) - def update(self, val): + def update(self, val: int): """Updates the progress value to the given value. Parameters ---------- - val : int + val The new current progress value. """ self._val = val diff --git a/skyllh/core/py.py b/skyllh/core/py.py index b90bc9497c..3ec9b3bc2f 100644 --- a/skyllh/core/py.py +++ b/skyllh/core/py.py @@ -3,24 +3,26 @@ import inspect import sys from collections import OrderedDict +from collections.abc import Callable, Iterable, Sequence +from typing import Any, Literal, cast, overload import numpy as np from skyllh.core.display import INDENTATION_WIDTH -def get_class_of_func(f): +def get_class_of_func(f: Callable) -> type | None: """Determines the class object that defined the given method or function ``f``. Parameters ---------- - f : function | method + f The function or method whose parent class should be determined. Returns ------- - cls : class | None + cls The class object which defines the given function or method. ``None`` is returned when no class could be determined. """ @@ -63,33 +65,33 @@ def module_classname(obj): return f'{obj.__module__}.{classname(obj)}' -def module_class_method_name(obj, meth_name): +def module_class_method_name(obj: object, meth_name: str): """Returns the module, class, and method name of the given instance ``obj``. Parameters ---------- - obj : instance of object + obj The object instance. - meth_name : str + meth_name The name of the method. """ return f'{module_classname(obj)}.{meth_name}' -def get_byte_size_prefix(size): +def get_byte_size_prefix(size: int) -> tuple[float, str]: """Determines the biggest size prefix for the given size in bytes such that the new size is still greater one. Parameters ---------- - size : int + size The size in bytes. Returns ------- - newsize : float + newsize The new byte size accounting for the byte prefix. - prefix : str + prefix The biggest byte size prefix. """ prefix_factor_list = [('', 1), ('K', 1024), ('M', 1024**2), ('G', 1024**3), ('T', 1024**4)] @@ -106,7 +108,7 @@ def get_byte_size_prefix(size): return (newsize, prefix) -def getsizeof(objects): +def getsizeof(objects) -> int: """Determines the size in bytes the given objects have in memory. If an object is a sequence, the size of the elements of the sequence will be estimated as well and added to the result. This does not account for the @@ -114,11 +116,11 @@ def getsizeof(objects): Parameters ---------- - objects : sequence of instances of object | instance of object. + objects Returns ------- - memsize : int + memsize The memory size in bytes of the given objects. """ if not issequence(objects): @@ -134,7 +136,7 @@ def getsizeof(objects): return memsize -def issequence(obj): +def issequence(obj) -> bool: """Checks if the given object ``obj`` is a sequence or not. The definition of a sequence in this case is, that the function ``len`` is defined for the object. @@ -145,7 +147,7 @@ def issequence(obj): Returns ------- - check : bool + check ``True`` if the given object is a sequence. ``False`` if the given object is an instance of str or not a sequence. @@ -161,43 +163,43 @@ def issequence(obj): return True -def issequenceof(obj, T): +def issequenceof(obj: object, T) -> bool: """Checks if the given object ``obj`` is a sequence with items being instances of type ``T``. Parameters ---------- - obj : instance of object + obj The Python object to check. - T : type | tuple of types + T The type each item of the sequence should be. If a tuple of types is given, each item can be one of the given types. Returns ------- - check : bool + check The result of the check. """ if not issequence(obj): return False - return all(isinstance(item, T) for item in obj) + return all(isinstance(item, T) for item in cast(Iterable, obj)) -def issequenceofsubclass(obj, T): +def issequenceofsubclass(obj: object, T: type) -> bool: """Checks if the given object ``obj`` is a sequence with items being sub-classes of class T. Parameters ---------- - obj : instance of object + obj The object to check. - T : class + T The base class of the items of the given object. Returns ------- - check : bool + check ``True`` if the given object is a sequence of instances which are sub-classes of class ``T``. ``False`` if ``obj`` is not a sequence or any item is not a sub-class of class ``T``. @@ -205,23 +207,23 @@ def issequenceofsubclass(obj, T): if not issequence(obj): return False - return all(issubclass(item, T) for item in obj) + return all(issubclass(item, T) for item in cast(Any, obj)) -def isproperty(obj, name): +def isproperty(obj: object, name: str) -> bool: """Checks if the given attribute is of type property. The attribute must exist in ``obj``. Parameters ---------- - obj : object + obj The Python object whose attribute to check for being a property. - name : str + name The name of the attribute. Returns ------- - check : bool + check True if the given attribute is of type property, False otherwise. Raises @@ -233,19 +235,19 @@ def isproperty(obj, name): return isinstance(attr, property) -def func_has_n_args(func, n): +def func_has_n_args(func: Callable, n: int) -> bool: """Checks if the given function `func` has `n` arguments. Parameters ---------- - func : callable + func The function to check. - n : int + n The number of arguments the function must have. Returns ------- - check : bool + check True if the given function has `n` arguments. False otherwise. """ check = len(inspect.signature(func).parameters) == n @@ -264,7 +266,11 @@ def bool_cast(v, errmsg): return v -def int_cast(v, errmsg, allow_None=False): +@overload +def int_cast(v, errmsg: str, allow_None: Literal[False] = ...) -> int: ... +@overload +def int_cast(v, errmsg: str, allow_None: Literal[True]) -> int | None: ... +def int_cast(v, errmsg, allow_None: bool = False) -> int | None: """Casts the given value to an integer value. If the cast is impossible, a TypeError is raised with the given error message. If `allow_None` is set to `True` the value `v` can also be `None`. @@ -280,19 +286,27 @@ def int_cast(v, errmsg, allow_None=False): return v -def float_cast(v, errmsg, allow_None=False): +@overload +def float_cast(v: Sequence, errmsg: str, allow_None: Literal[False] = ...) -> list[float]: ... +@overload +def float_cast(v: Sequence, errmsg: str, allow_None: Literal[True]) -> list[float | None]: ... +@overload +def float_cast(v: float | str | bool, errmsg: str, allow_None: Literal[False] = ...) -> float: ... +@overload +def float_cast(v: float | str | bool | None, errmsg: str, allow_None: Literal[True]) -> float | None: ... +def float_cast(v, errmsg: str, allow_None: bool = False) -> float | list[float] | list[float | None] | None: """Casts the given value to a float. If the cast is impossible, a TypeError is raised with the given error message. If `allow_None` is set to `True` the value `v` can also be `None`. Parameters ---------- - v : to_float_castable object | sequence of to_float_castable objects + v The object that should get casted to a float. This can also be a sequence of objects that should get casted to floats. - errmsg : str + errmsg The error message in case the cast failed. - allow_None : bool + allow_None Flag if ``None`` is allowed as value for v. If yes, the casted result is ``None``. @@ -305,6 +319,9 @@ def float_cast(v, errmsg, allow_None=False): # Define cast function for a single object. def _obj_float_cast(v, errmsg, allow_None): + """Casts a single object ``v`` to a float. If ``allow_None`` is set to + ``True`` and ``v`` is ``None``, ``None`` is returned. + """ if allow_None and v is None: return v @@ -324,7 +341,15 @@ def _obj_float_cast(v, errmsg, allow_None): return _obj_float_cast(v, errmsg, allow_None) -def str_cast(v, errmsg, allow_None=False): +@overload +def str_cast(v, errmsg, allow_None: Literal[False] = ...) -> str: ... + + +@overload +def str_cast(v, errmsg, allow_None: Literal[True]) -> str | None: ... + + +def str_cast(v, errmsg, allow_None: bool = False): """Casts the given value to a str object. If the cast is impossible, a TypeError is raised with the given error message. @@ -356,25 +381,25 @@ def list_of_cast(t, v, errmsg): return v -def get_smallest_numpy_int_type(values): +def get_smallest_numpy_int_type(values: int | Sequence[int]) -> np.ndarray: """Returns the smallest numpy integer type that can represent the given integer values. Parameters ---------- - values : int | sequence of int + values The integer value(s) that need to be representable by the returned integer type. Returns ------- - inttype : numpy integer type + inttype The smallest numpy integer type that can represent the given values. """ - values = np.atleast_1d(values) + _values = np.atleast_1d(values) - vmin = np.min(values) - vmax = np.max(values) + vmin = np.min(_values) + vmax = np.max(_values) if vmin < 0: # noqa: SIM108 types = [np.int8, np.int16, np.int32, np.int64] @@ -389,19 +414,19 @@ def get_smallest_numpy_int_type(values): raise ValueError(f'No integer type spans [{vmin}, {vmax}]!') -def get_number_of_float_decimals(value): +def get_number_of_float_decimals(value: float) -> int: """Determines the number of significant decimals the given float number has. The maximum number of supported decimals is 16. Parameters ---------- - value : float + value The float value whose number of significant decimals should get determined. Returns ------- - decimals : int + decimals The number of decimals of value which are non-zero. Raises @@ -421,26 +446,23 @@ def get_number_of_float_decimals(value): return 0 -def make_dict_hash(d): +def make_dict_hash(d: dict | None) -> int: """Creates a hash value for the given dictionary. Parameters ---------- - d : dict | None + d The dictionary holding (name: value) pairs. If set to None, an empty dictionary is used. Returns ------- - hash : int + hash The hash of the dictionary. """ if d is None: d = {} - if not isinstance(d, dict): - raise TypeError('The d argument must be of type dict!') - # A note on the ordering of Python dictionary items: The items are ordered # internally according to the hash value of their keys. Hence, if we don't # insert more dictionary items, the order of the items won't change. Thus, @@ -457,16 +479,16 @@ class ObjectCollection: well. """ - def __init__(self, objs=None, obj_type=None): + def __init__(self, objs=None, obj_type: type | None = None): """Constructor of the ObjectCollection class. Must be called by the derived class. Parameters ---------- - objs : instance of obj_type | sequence of obj_type instances | None + objs The sequence of objects of type ``obj_type`` with which this collection should get initialized with. - obj_type : type | None + obj_type The type of the objects, which can be added to the collection. If set to None, the type will be determined from the given objects. If no objects are given, the object type will be `object`. @@ -515,7 +537,7 @@ def __getitem__(self, key): def __iter__(self): return iter(self._objects) - def __add__(self, other): + def __add__(self, other) -> 'ObjectCollection': """Implementation to support the operation ``oc = self + other``, where ``self`` is this ObjectCollection object and ``other`` something useful else. This creates a copy ``oc`` of ``self`` and adds ``other`` @@ -523,11 +545,11 @@ def __add__(self, other): Parameters ---------- - other : obj_type | ObjectCollection of obj_type + other Returns ------- - oc : ObjectCollection + oc The new ObjectCollection object with object from self and other. """ oc = self.copy() @@ -547,13 +569,13 @@ def copy(self): oc._objects = copy.copy(self._objects) return oc - def add(self, obj): + def add(self, obj) -> 'ObjectCollection': """Adds the given object, sequence of objects, or object collection to this object collection. Parameters ---------- - obj : obj_type instance | sequence of obj_type | + obj ObjectCollection of obj_type An instance of ``obj_type`` that should be added to the collection. If given an ObjectCollection for objects of type obj_type, it will @@ -561,7 +583,7 @@ def add(self, obj): Returns ------- - self : ObjectCollection + self The instance of this ObjectCollection, in order to be able to chain several ``add`` calls. """ @@ -590,35 +612,35 @@ def add(self, obj): __iadd__ = add - def index(self, obj): + def index(self, obj) -> int: """Gets the index of the given object instance within this object collection. Parameters ---------- - obj : obj_type instance + obj The instance of obj_type whose index should get retrieved. Returns ------- - idx : int + idx The index of the object within this object collection. """ return self._objects.index(obj) - def pop(self, index=None): + def pop(self, index: int | None = None): """Removes and returns the object at the given index (default last). Raises IndexError if the collection is empty or index is out of range. Parameters ---------- - index : int | None + index The index of the object to remove. If set to None, the index of the last object is used. Returns ------- - obj : obj_type + obj The removed object. """ if index is None: @@ -633,16 +655,16 @@ class NamedObjectCollection(ObjectCollection): tracked w.r.t. its name. """ - def __init__(self, objs=None, obj_type=None, **kwargs): + def __init__(self, objs=None, obj_type: type | None = None, **kwargs): """Creates a new NamedObjectCollection instance. Must be called by the derived class. Parameters ---------- - objs : instance of obj_type | sequence of instances of obj_type | None + objs The sequence of objects of type ``obj_type`` with which this collection should get initialized with. - obj_type : type + obj_type The type of the objects, which can be added to the collection. This type must have an attribute named ``name``. """ @@ -664,20 +686,20 @@ def name_list(self): """ return list(self._obj_name_to_idx.keys()) - def _create_obj_name_to_idx_dict(self, start=None, end=None): + def _create_obj_name_to_idx_dict(self, start: int | None = None, end: int | None = None) -> dict: """Creates the dictionary {obj.name: index} for object in the interval [`start`, `end`). Parameters ---------- - start : int | None + start The object start index position, which is inclusive. - end : int | None + end The object end index position, which is exclusive. Returns ------- - obj_name_to_idx : dict + obj_name_to_idx The dictionary {obj.name: index}. """ if start is None: @@ -685,35 +707,35 @@ def _create_obj_name_to_idx_dict(self, start=None, end=None): return OrderedDict([(o.name, start + idx) for (idx, o) in enumerate(self._objects[start:end])]) - def __contains__(self, name): + def __contains__(self, name: str) -> bool: """Returns ``True`` if an object of the given name exists in this NamedObjectCollection instance, ``False`` otherwise. Parameters ---------- - name : str + name The name of the object. Returns ------- - check : bool + check ``True`` if an object of name ``name`` exists in this NamedObjectCollection instance, ``False`` otherwise. """ return name in self._obj_name_to_idx - def __getitem__(self, key): + def __getitem__(self, key: str | int): """Returns an object based on its name or index. Parameters ---------- - key : str | int + key The object identification. Either its name or its index position within the object collection. Returns ------- - obj : instance of obj_type + obj The requested object. Raises @@ -725,12 +747,12 @@ def __getitem__(self, key): key = self.get_index_by_name(key) return super().__getitem__(key) - def add(self, obj): + def add(self, obj) -> 'NamedObjectCollection': """Adds the given object to this named object collection. Parameters ---------- - obj : obj_type instance | NamedObjectCollection of obj_type + obj An instance of ``obj_type`` that should be added to this named object collection. If a NamedObjectCollection instance for objects of type ``obj_type`` @@ -739,7 +761,7 @@ def add(self, obj): Returns ------- - self : NamedObjectCollection + self The instance of this NamedObjectCollection, in order to be able to chain several ``add`` calls. """ @@ -753,36 +775,36 @@ def add(self, obj): __iadd__ = add - def get_index_by_name(self, name): + def get_index_by_name(self, name: str) -> int: """Gets the index of the object with the given name within this named object collection. Parameters ---------- - name : str + name The name of the object whose index should get retrieved. Returns ------- - idx : int + idx The index of the object within this named object collection. """ return self._obj_name_to_idx[name] - def pop(self, index=None): + def pop(self, index: int | str | None = None): """Removes and returns the object at the given index (default last). Raises IndexError if the collection is empty or index is out of range. Parameters ---------- - index : int | str | None + index The index of the object to remove. If set to None, the index of the last object is used. If a str instance is given, it specifies the name of the object. Returns ------- - obj : obj_type instance + obj The removed object. """ if isinstance(index, str): diff --git a/skyllh/core/random.py b/skyllh/core/random.py index c8fd6a86bf..bd5a24fa97 100644 --- a/skyllh/core/random.py +++ b/skyllh/core/random.py @@ -15,7 +15,7 @@ class RandomStateService: def __init__( self, - seed=None, + seed: int | None = None, **kwargs, ): """Creates a new random state service. The ``random`` property can then @@ -23,7 +23,7 @@ def __init__( Parameters ---------- - seed : int | None + seed The seed to use. If None, the random number generator will be seeded randomly. See the numpy documentation for numpy.random.RandomState what that means. @@ -51,12 +51,12 @@ def random(self, random): raise TypeError('The random property must be of type numpy.random.RandomState!') self._random = random - def reseed(self, seed): + def reseed(self, seed: int | None): """Reseeds the random number generator with the given seed. Parameters ---------- - seed : int | None + seed The seed to use. If None, the random number generator will be seeded randomly. See the numpy documentation for numpy.random.RandomState what that means. @@ -73,8 +73,8 @@ class RandomChoice: def __init__( self, - items, - probabilities, + items: np.ndarray, + probabilities: np.ndarray, **kwargs, ): """Creates a new instance of RandomChoice holding the probabilities @@ -82,10 +82,10 @@ def __init__( Parameters ---------- - items : instance of numpy.ndarray + items The (N,)-shaped numpy.ndarray holding the items from which to choose. - probabilities : instance of numpy.ndarray + probabilities The (N,)-shaped numpy.ndarray holding the probability for each item. """ super().__init__(**kwargs) @@ -117,13 +117,13 @@ def probabilities(self): def _assert_items( self, - items, + items: np.ndarray, ): """Checks for the correct type and shape of the items. Parameters ---------- - items : The (N,)-shaped numpy.ndarray holding the items from which to + items choose. Raises @@ -141,16 +141,16 @@ def _assert_items( def _assert_probabilities( self, - p, - n_items, + p: np.ndarray, + n_items: int, ): """Checks for correct values of the probabilities. Parameters ---------- - p : instance of numpy.ndarray + p The (N,)-shaped numpy.ndarray holding the probability for each item. - n_items : int + n_items The number of items. Raises @@ -178,23 +178,23 @@ def _assert_probabilities( def __call__( self, - rss, - size, - ): + rss: 'RandomStateService', + size: int, + ) -> np.ndarray: """Chooses ``size`` random items from ``self.items`` according to ``self.probabilities``. Parameters ---------- - rss : instance of RandomStateService + rss The instance of RandomStateService from which random numbers are drawn from. - size : int + size The number of items to draw. Returns ------- - random_items : instance of numpy.ndarray + random_items The (size,)-shaped numpy.ndarray holding the randomly selected items from ``self.items``. """ diff --git a/skyllh/core/scrambling.py b/skyllh/core/scrambling.py index c73196d0bc..0d118a6042 100644 --- a/skyllh/core/scrambling.py +++ b/skyllh/core/scrambling.py @@ -1,7 +1,11 @@ import abc +from collections.abc import Callable import numpy as np +from skyllh.core.dataset import Dataset +from skyllh.core.random import RandomStateService +from skyllh.core.storage import DataFieldRecordArray from skyllh.core.times import ( TimeGenerator, ) @@ -13,32 +17,33 @@ class DataScramblingMethod( """Base class for implementing a data scrambling method.""" def __init__(self, **kwargs): + """Creates a new instance of DataScramblingMethod.""" super().__init__(**kwargs) @abc.abstractmethod def scramble( self, - rss, - dataset, - data, - ): + rss: RandomStateService, + dataset: Dataset, + data: DataFieldRecordArray, + ) -> DataFieldRecordArray: """The scramble method implements the actual scrambling of the given data, which is method dependent. The scrambling must be performed in-place, i.e. it alters the data inside the given data array. Parameters ---------- - rss : instance of RandomStateService + rss The random state service providing the random number generator (RNG). - dataset : instance of Dataset + dataset The instance of Dataset for which the data should get scrambled. - data : instance of DataFieldRecordArray + data The DataFieldRecordArray containing the to be scrambled data. Returns ------- - data : instance of DataFieldRecordArray + data The given DataFieldRecordArray holding the scrambled data. """ @@ -57,7 +62,7 @@ class UniformRAScramblingMethod( def __init__( self, - ra_range=None, + ra_range: tuple | None = None, **kwargs, ): r"""Initializes a new RAScramblingMethod instance. @@ -92,26 +97,26 @@ def ra_range(self, ra_range): def scramble( self, - rss, - dataset, - data, - ): + rss: RandomStateService, + dataset: Dataset, + data: DataFieldRecordArray, + ) -> DataFieldRecordArray: """Scrambles the given data uniformly in right-ascention. Parameters ---------- - rss : instance of RandomStateService + rss The random state service providing the random number generator (RNG). - dataset : instance of Dataset + dataset The instance of Dataset for which the data should get scrambled. - data : instance of DataFieldRecordArray + data The DataFieldRecordArray instance containing the to be scrambled data. Returns ------- - data : instance of DataFieldRecordArray + data The given DataFieldRecordArray holding the scrambled data. """ dt = data['ra'].dtype @@ -131,17 +136,17 @@ class TimeScramblingMethod(DataScramblingMethod): def __init__( self, - timegen, - hor_to_equ_transform, + timegen: TimeGenerator, + hor_to_equ_transform: Callable, **kwargs, ): """Initializes a new time scramling method instance. Parameters ---------- - timegen : instance of TimeGenerator + timegen The time generator that should be used to generate random MJD times. - hor_to_equ_transform : callable + hor_to_equ_transform The transformation function to transform coordinates from the horizontal system into the equatorial system. @@ -171,24 +176,29 @@ def timegen(self, timegen): self._timegen = timegen @property - def hor_to_equ_transform(self): + def hor_to_equ_transform( + self, + ) -> Callable[[np.ndarray, np.ndarray, np.ndarray], tuple[np.ndarray, np.ndarray]]: """The transformation function to transform coordinates from the horizontal system into the equatorial system. """ return self._hor_to_equ_transform @hor_to_equ_transform.setter - def hor_to_equ_transform(self, transform): + def hor_to_equ_transform( + self, + transform: Callable[[np.ndarray, np.ndarray, np.ndarray], tuple[np.ndarray, np.ndarray]], + ): if not callable(transform): raise TypeError('The hor_to_equ_transform property must be a callable object!') self._hor_to_equ_transform = transform def scramble( self, - rss, - dataset, - data, - ): + rss: RandomStateService, + dataset: Dataset, + data: DataFieldRecordArray, + ) -> DataFieldRecordArray: """Scrambles the given data based on random MJD times, which are generated from a TimeGenerator instance. The event's right-ascention and declination coordinates are calculated via a horizontal-to-equatorial @@ -196,18 +206,18 @@ def scramble( Parameters ---------- - rss : instance of RandomStateService + rss The random state service providing the random number generator (RNG). - dataset : instance of Dataset + dataset The instance of Dataset for which the data should get scrambled. - data : instance of DataFieldRecordArray + data The DataFieldRecordArray instance containing the to be scrambled data. Returns ------- - data : instance of DataFieldRecordArray + data The given DataFieldRecordArray holding the scrambled data. """ mjds = self.timegen.generate_times(rss, len(data)) @@ -220,9 +230,13 @@ def scramble( class DataScrambler: + """This class provides a data scrambler that scrambles data using a defined + data scrambling method. + """ + def __init__( self, - method, + method: 'DataScramblingMethod', **kwargs, ): """Creates a data scrambler instance with a given defined scrambling @@ -230,7 +244,7 @@ def __init__( Parameters ---------- - method : instance of DataScramblingMethod + method The instance of DataScramblingMethod that defines the method of the data scrambling. """ @@ -253,11 +267,11 @@ def method(self, method): def scramble_data( self, - rss, - dataset, - data, - copy=False, - ): + rss: RandomStateService, + dataset: Dataset, + data: DataFieldRecordArray, + copy: bool = False, + ) -> DataFieldRecordArray: """Scrambles the given data by calling the scramble method of the scrambling method class, that was configured for the data scrambler. If the ``inplace_scrambling`` property is set to False, a copy of the @@ -265,21 +279,21 @@ def scramble_data( Parameters ---------- - rss : instance of RandomStateService + rss The random state service providing the random number generator (RNG). - dataset : instance of Dataset + dataset The instance of Dataset for which the data should get scrambled. - data : instance of DataFieldRecordArray + data The instance of DataFieldRecordArray holding the data, which should get scrambled. - copy : bool + copy Flag if a copy of the given data should be made before scrambling the data. The default is False. Returns ------- - data : instance of DataFieldRecordArray + data The given DataFieldRecordArray instance with the scrambled data. If the ``inplace_scrambling`` property is set to True, this output array is the same array as the input array, otherwise it's a new diff --git a/skyllh/core/services.py b/skyllh/core/services.py index e0e4df3cf0..8ceed0ce86 100644 --- a/skyllh/core/services.py +++ b/skyllh/core/services.py @@ -123,8 +123,8 @@ def change_shg_mgr( def get_builder_to_shgidxs_dict( self, - ds_idx, - ): + ds_idx: int, + ) -> dict: """Creates a dictionary with the builder instance as key and the list of source hypo group indices to which the builder applies as value. Hence, SHGs using the same builder instance can be grouped for @@ -132,12 +132,12 @@ def get_builder_to_shgidxs_dict( Parameters ---------- - ds_idx : int + ds_idx The index of the dataset for which the same builders apply. Returns ------- - builder_shgidxs_dict : dict + builder_shgidxs_dict The dictionary with the builder instance as key and the list of source hypo group indices to which the builder applies as value. """ @@ -165,8 +165,8 @@ def get_builder_to_shgidxs_dict( def construct_detsigyield_array( self, - ppbar=None, - ): + ppbar: ProgressBar | None = None, + ) -> np.ndarray: """Creates a (N_datasets, N_source_hypo_groups)-shaped numpy ndarray of object holding the constructed DetSigYield instances. @@ -178,12 +178,12 @@ def construct_detsigyield_array( Parameters ---------- - ppbar : instance of ProgressBar | None + ppbar The instance of ProgressBar of the optional parent progress bar. Returns ------- - detsigyield_arr : instance of numpy.ndarray + detsigyield_arr The (N_datasets, N_source_hypo_groups)-shaped numpy ndarray of object holding the constructed DetSigYield instances. """ @@ -247,15 +247,15 @@ class SrcDetSigYieldWeightsService: @staticmethod def create_src_recarray_list_list( - detsigyield_service, - ): + detsigyield_service: 'DetSigYieldService', + ) -> list[list[np.ndarray]]: """Creates a list of numpy record ndarrays, one for each source hypothesis group suited for evaluating the detector signal yield instance of that source hypothesis group. Parameters ---------- - detsigyield_service : instance of DetSigYieldService + detsigyield_service The instance of DetSigYieldService providing the (N_datasets, N_source_hypo_groups)-shaped 2D ndarray of DetSigYield instances, one for each dataset and source hypothesis @@ -263,7 +263,7 @@ def create_src_recarray_list_list( Returns ------- - src_recarray_list_list : list of list of numpy record ndarrays + src_recarray_list_list The (N_datasets,N_source_hypo_groups)-shaped list of list of the source numpy record ndarrays, one for each dataset and source hypothesis group combination, which is needed for @@ -286,20 +286,20 @@ def create_src_recarray_list_list( @staticmethod def create_src_weight_array_list( - shg_mgr, - ): + shg_mgr: SourceHypoGroupManager, + ) -> list[np.ndarray]: """Creates a list of numpy 1D ndarrays holding the source weights, one for each source hypothesis group. Parameters ---------- - shg_mgr : instance of SourceHypoGroupManager + shg_mgr The instance of SourceHypoGroupManager defining the source hypothesis groups with their sources. Returns ------- - src_weight_array_list : list of numpy 1D ndarrays + src_weight_array_list The list of 1D numpy ndarrays holding the source weights, one for each source hypothesis group. """ @@ -308,14 +308,14 @@ def create_src_weight_array_list( def __init__( self, - detsigyield_service, + detsigyield_service: 'DetSigYieldService', **kwargs, ): """Creates a new SrcDetSigYieldWeightsService instance. Parameters ---------- - detsigyield_service : instance of DetSigYieldService + detsigyield_service The instance of DetSigYieldService providing the (N_datasets, N_source_hypo_groups)-shaped array of DetSigYield instances, one instance for each combination of dataset and source @@ -394,14 +394,14 @@ def src_recarray_list_list(self): def change_shg_mgr( self, - shg_mgr, + shg_mgr: SourceHypoGroupManager, ): """Re-creates the internal source numpy record arrays needed for the detector signal yield calculation. Parameters ---------- - shg_mgr : instance of SourceHypoGroupManager + shg_mgr The new SourceHypoGroupManager instance. """ if id(shg_mgr) != id(self._detsigyield_service.shg_mgr): @@ -416,7 +416,7 @@ def change_shg_mgr( self._src_weight_array_list = type(self).create_src_weight_array_list(shg_mgr=self._detsigyield_service.shg_mgr) - def calculate(self, src_params_recarray): + def calculate(self, src_params_recarray: np.ndarray): """Calculates the source detector signal yield weights for each source and their derivative w.r.t. each global floating parameter. The result is stored internally as: @@ -433,7 +433,7 @@ def calculate(self, src_params_recarray): Parameters ---------- - src_params_recarray : instance of numpy record ndarray + src_params_recarray The numpy record ndarray of length N_sources holding the local source parameters. See the documentation of :meth:`skyllh.core.parameters.ParameterModelMapper.create_src_params_recarray` @@ -474,17 +474,17 @@ def calculate(self, src_params_recarray): sidx += shg_n_src - def get_weights(self): + def get_weights(self) -> tuple[np.ndarray | None, defaultdict | None]: """Returns the source detector signal yield weights and their derivatives w.r.t. the global fit parameters. Returns ------- - a_jk : instance of ndarray + a_jk The (N_datasets, N_sources)-shaped numpy ndarray holding the source detector signal yield weight for each combination of dataset and source. - a_jk_grads : dict + a_jk_grads The dictionary holding the (N_datasets, N_sources)-shaped numpy ndarray with the derivatives w.r.t. the global fit parameter the SrcDetSigYieldWeightsService depend on. The dictionary's key @@ -501,7 +501,7 @@ class DatasetSignalWeightFactorsService: :class:`~SrcDetSigYieldWeightsService` class. """ - def __init__(self, src_detsigyield_weights_service): + def __init__(self, src_detsigyield_weights_service: 'SrcDetSigYieldWeightsService'): r"""Creates a new DatasetSignalWeightFactors instance. Parameters @@ -548,6 +548,9 @@ def calculate(self): """ (a_jk, a_jk_grads) = self._src_detsigyield_weights_service.get_weights() + assert a_jk is not None, 'SrcDetSigYieldWeightsService.calculate() must be called before get_weights()' + assert a_jk_grads is not None, 'SrcDetSigYieldWeightsService.calculate() must be called before get_weights()' + a_j = np.sum(a_jk, axis=1) a = np.sum(a_jk) @@ -567,15 +570,15 @@ def calculate(self): a_grads = np.sum(a_jk_grads[gpidx]) self._f_j_grads[gpidx] = (a_j_grads * a - a_j * a_grads) / a**2 - def get_weights(self): + def get_weights(self) -> tuple[np.ndarray | None, dict]: """Returns the Returns ------- - f_j : instance of ndarray + f_j The (N_datasets,)-shaped 1D numpy ndarray holding the dataset signal weight factor for each dataset. - f_j_grads : dict + f_j_grads The dictionary holding the (N_datasets,)-shaped numpy ndarray with the derivatives w.r.t. the global fit parameter the DatasetSignalWeightFactorsService depend on. diff --git a/skyllh/core/session.py b/skyllh/core/session.py index 2ded3fbe12..9386793e86 100644 --- a/skyllh/core/session.py +++ b/skyllh/core/session.py @@ -22,23 +22,23 @@ def disable_interactive_session(): IS_INTERACTIVE_SESSION = False -def is_interactive_session(): +def is_interactive_session() -> bool: """Checks whether the current session is interactive (True) or not (False). Returns ------- - check : bool + check True if the current SkyLLH session is interactive, False otherwise. """ return IS_INTERACTIVE_SESSION -def is_python_interpreter_in_interactive_mode(): +def is_python_interpreter_in_interactive_mode() -> bool: """Checks if the Python interpreter is in interactive mode. Returns ------- - check : bool + check True if the Python interpreter is in interactive mode, False otherwise. """ return bool(getattr(sys, 'ps1', sys.flags.interactive)) diff --git a/skyllh/core/signal_generation.py b/skyllh/core/signal_generation.py index 09020f27af..f137c67a08 100644 --- a/skyllh/core/signal_generation.py +++ b/skyllh/core/signal_generation.py @@ -4,6 +4,7 @@ float_cast, issequence, ) +from skyllh.core.types import SourceHypoGroup_t class HasEnergyRange(metaclass=abc.ABCMeta): @@ -11,7 +12,7 @@ class HasEnergyRange(metaclass=abc.ABCMeta): @property @abc.abstractmethod - def energy_range(self): + def energy_range(self) -> 'tuple[float, float] | None': """Configured true-energy range as a 2-element tuple in GeV, or None.""" @energy_range.setter @@ -28,14 +29,14 @@ class SignalGenerationMethod(HasEnergyRange, metaclass=abc.ABCMeta): def __init__( self, - energy_range, + energy_range: tuple[float, float] | None, **kwargs, ): """Constructs a new signal generation method instance. Parameters ---------- - energy_range : 2-element tuple of float | None + energy_range The energy range from which to take MC events into account for signal event generation, specified in true neutrino energy (GeV). If set to None, the entire energy range [0, +inf] is used. @@ -45,7 +46,7 @@ def __init__( self.energy_range = energy_range @property - def energy_range(self): + def energy_range(self) -> 'tuple[float, float] | None': """The 2-element tuple of floats holding the energy range from which to take MC events into account for signal event generation, in GeV. """ @@ -72,7 +73,7 @@ def energy_range(self, r): def calc_source_signal_mc_event_flux( self, data_mc, - shg, + shg: SourceHypoGroup_t, ): """This method is supposed to calculate the signal flux of each given MC event for each source hypothesis of the given source hypothesis @@ -80,29 +81,29 @@ def calc_source_signal_mc_event_flux( Parameters ---------- - data_mc : numpy record ndarray + data_mc The numpy record array holding all the MC events. - shg : instance of SourceHypoGroup + shg The source hypothesis group instance, which defines the list of sources, and their flux model. Returns ------- - ev_idx_arr : ndarray + ev_idx_arr The (N_selected_signal_events,)-shaped 1D ndarray holding the index of the MC event. - shg_src_idx_arr : ndarray + shg_src_idx_arr The (N_selected_signal_events,)-shaped 1D ndarray holding the index of the source within the given source hypothesis group for each signal candidate event. - flux_arr : ndarray + flux_arr The (N_selected_signal_events,)-shaped 1D ndarray holding the flux value of each signal candidate event. """ def signal_event_post_sampling_processing( self, - shg, + shg: SourceHypoGroup_t, shg_sig_events_meta, shg_sig_events, ): @@ -112,26 +113,26 @@ def signal_event_post_sampling_processing( Parameters ---------- - shg : SourceHypoGroup instance + shg The source hypothesis group instance holding the sources and their locations. - shg_sig_events_meta : numpy record ndarray + shg_sig_events_meta The numpy record ndarray holding meta information about the generated signal events for the given source hypothesis group. The length of this array must be the same as shg_sig_events. It needs to contain the following data fields: - shg_src_idx : int + shg_src_idx The source index within the source hypothesis group. - shg_sig_events : numpy record ndarray + shg_sig_events The numpy record ndarray holding the generated signal events for the given source hypothesis group and in the format of the original MC events. Returns ------- - shg_sig_events : numpy record array + shg_sig_events The processed signal events. In the default implementation of this method this is just the shg_sig_events input array. """ diff --git a/skyllh/core/signal_generator.py b/skyllh/core/signal_generator.py index a75012245b..6df78e1aec 100644 --- a/skyllh/core/signal_generator.py +++ b/skyllh/core/signal_generator.py @@ -26,12 +26,8 @@ int_cast, issequenceof, ) -from skyllh.core.random import ( - RandomChoice, -) -from skyllh.core.services import ( - DatasetSignalWeightFactorsService, -) +from skyllh.core.random import RandomChoice, RandomStateService +from skyllh.core.services import DatasetSignalWeightFactorsService, SrcDetSigYieldWeightsService from skyllh.core.source_hypo_grouping import ( SourceHypoGroupManager, ) @@ -52,14 +48,14 @@ class SignalGenerator( def __init__( self, - shg_mgr, + shg_mgr: SourceHypoGroupManager, **kwargs, ): """Constructs a new signal generator instance. Parameters ---------- - shg_mgr : instance of SourceHypoGroupManager + shg_mgr The SourceHypoGroupManager instance defining the source hypothesis groups. """ @@ -80,20 +76,20 @@ def shg_mgr(self, manager): raise TypeError('The shg_mgr property must be an instance of SourceHypoGroupManager!') self._shg_mgr = manager - def create_src_params_recarray(self, src_detsigyield_weights_service): + def create_src_params_recarray(self, src_detsigyield_weights_service: SrcDetSigYieldWeightsService) -> np.ndarray: """Creates the src_params_recarray structured ndarray of length N_sources holding the local source parameter names and values needed for the calculation of the detector signal yields. Parameters ---------- - src_detsigyield_weights_service : instance of SrcDetSigYieldWeightsService + src_detsigyield_weights_service The instance of SrcDetSigYieldWeightsService providing the product of the source weights with the detector signal yield. Returns ------- - src_params_recarray : instance of numpy structured ndarray + src_params_recarray The structured numpy ndarray of length N_sources, holding the local parameter names and values of each source needed to calculate the detector signal yield. @@ -135,36 +131,42 @@ def change_shg_mgr(self, shg_mgr): self.shg_mgr = shg_mgr @abc.abstractmethod - def generate_signal_events(self, rss, mean, poisson=True, src_detsigyield_weights_service=None): + def generate_signal_events( + self, + rss: RandomStateService, + mean: float, + poisson: bool = True, + src_detsigyield_weights_service: SrcDetSigYieldWeightsService | None = None, + ) -> tuple[int, dict[int, DataFieldRecordArray]]: """This abstract method must be implemented by the derived class to generate a given number of signal events. Parameters ---------- - rss : instance of RandomStateService + rss The instance of RandomStateService providing the random number generator state. - mean : int | float + mean The mean number of signal events. If the ``poisson`` argument is set to True, the actual number of generated signal events will be drawn from a Poisson distribution with this given mean value of signal events. - poisson : bool + poisson If set to True, the actual number of generated signal events will be drawn from a Poisson distribution with the given mean value of signal events. If set to False, the argument ``mean`` specifies the actual number of generated signal events. - src_detsigyield_weights_service : instance of SrcDetSigYieldWeightsService | None + src_detsigyield_weights_service The instance of SrcDetSigYieldWeightsService providing the weighting of the sources within the detector. This can be ``None`` if this signal generator does not need this information. Returns ------- - n_signal : int + n_signal The number of generated signal events. - signal_events_dict : dict of DataFieldRecordArray + signal_events_dict The dictionary holding the DataFieldRecordArray instances with the generated signal events. Each key of this dictionary represents the dataset index for which the signal events have been generated. @@ -181,27 +183,33 @@ class MultiDatasetSignalGenerator( """ def __init__( - self, shg_mgr, dataset_list, data_list, sig_generator_list=None, ds_sig_weight_factors_service=None, **kwargs + self, + shg_mgr: SourceHypoGroupManager, + dataset_list: list[Dataset], + data_list: list[DatasetData], + sig_generator_list: 'list[SignalGenerator] | None' = None, + ds_sig_weight_factors_service: DatasetSignalWeightFactorsService | None = None, + **kwargs, ): """Constructs a new signal generator handling multiple datasets. Parameters ---------- - shg_mgr : instance of SourceHypoGroupManager + shg_mgr The instance of SourceHypoGroupManager that defines the list of source hypothesis groups, i.e. the list of sources. - dataset_list : list of instance of Dataset + dataset_list The list of instance of Dataset for which signal events should get generated. - data_list : list of instance of DatasetData + data_list The list of instance of DatasetData holding the actual data of each dataset. The order must match the order of ``dataset_list``. - sig_generator_list : list of instance of SignalGenerator | None + sig_generator_list The optional list of instance of SignalGenerator holding signal generator instances for each individual dataset. This can be ``None`` if this signal generator does not require individual signal generators for each dataset. - ds_sig_weight_factors_service : instance of DatasetSignalWeightFactorsService + ds_sig_weight_factors_service The instance of DatasetSignalWeightFactorsService providing the dataset signal weight factor service for calculating the dataset signal weights. @@ -297,7 +305,7 @@ def change_shg_mgr(self, shg_mgr): for sig_generator in filter(None, self.sig_generator_list): sig_generator.change_shg_mgr(shg_mgr=shg_mgr) - def fluxmodel_scaling_factor(self, src_params_recarray=None, per_source=False): + def fluxmodel_scaling_factor(self, src_params_recarray=None, per_source: bool = False): """Returns the scaling factor to convert a mean number of detected signal events into a flux normalization:: @@ -312,7 +320,7 @@ def fluxmodel_scaling_factor(self, src_params_recarray=None, per_source=False): the detector signal yield calculation. When provided, this overrides the internally cached reference parameter array for this call only (the cache is not updated). If ``None``, the cached reference parameter values built from the flux model defaults are used. - per_source : bool + per_source If set to True, return per-source scaling factors that sum to the global scaling factor. If set to False, return the global scaling factor. @@ -358,7 +366,11 @@ def fluxmodel_scaling_factor(self, src_params_recarray=None, per_source=False): f'user configuration between dataset {ref_energy_range_ds_idx} ({ref_energy_range}) ' f'and dataset {ds_idx} ({configured_energy_range}).' ) - if ref_energy_range is not None and tuple(ref_energy_range) != tuple(configured_energy_range): + if ( + ref_energy_range is not None + and configured_energy_range is not None + and tuple(ref_energy_range) != tuple(configured_energy_range) + ): raise ValueError( 'A single shared energy_range must be used across all datasets. Found inconsistent ' f'user configuration between dataset {ref_energy_range_ds_idx} ({ref_energy_range}) ' @@ -388,21 +400,27 @@ def fluxmodel_scaling_factor(self, src_params_recarray=None, per_source=False): return scaling_factor - def generate_signal_events(self, rss, mean, poisson=True): + def generate_signal_events( + self, + rss: RandomStateService, + mean: float, + poisson: bool = True, + src_detsigyield_weights_service: SrcDetSigYieldWeightsService | None = None, + ) -> tuple[int, dict[int, DataFieldRecordArray]]: """Generates a given number of signal events distributed across the individual datasets. Parameters ---------- - rss : instance of RandomStateService + rss The instance of RandomStateService providing the random number generator state. - mean : float | int + mean The mean number of signal events. If the ``poisson`` argument is set to True, the actual number of generated signal events will be drawn from a Poisson distribution with this given mean value of signal events. - poisson : bool + poisson If set to True, the actual number of generated signal events will be drawn from a Poisson distribution with the given mean value of signal events. @@ -411,9 +429,9 @@ def generate_signal_events(self, rss, mean, poisson=True): Returns ------- - n_signal : int + n_signal The number of actual generated signal events. - signal_events_dict : dict of DataFieldRecordArray + signal_events_dict The dictionary holding the DataFieldRecordArray instances with the generated signal events. Each key of this dictionary represents the dataset index for which the signal events have been generated. @@ -436,6 +454,7 @@ def generate_signal_events(self, rss, mean, poisson=True): self._ds_sig_weight_factors_service.calculate() (ds_weights, _) = self._ds_sig_weight_factors_service.get_weights() + assert ds_weights is not None # Calculate the number of events that need to be generated for each # individual dataset. Due to rounding errors, it could happen that the @@ -462,6 +481,7 @@ def generate_signal_events(self, rss, mean, poisson=True): n_signal = 0 signal_events_dict = {} + assert self._sig_generator_list is not None for n_events, ds_sig_generator in zip(n_events_arr, self._sig_generator_list, strict=True): (ds_n_signal, ds_sig_events_dict) = ds_sig_generator.generate_signal_events( rss=rss, @@ -495,26 +515,26 @@ class MCMultiDatasetSignalGenerator( def __init__( self, - shg_mgr, - dataset_list, - data_list, - valid_event_field_ranges_dict_list=None, + shg_mgr: SourceHypoGroupManager, + dataset_list: list[Dataset], + data_list: list[DatasetData], + valid_event_field_ranges_dict_list: list[dict] | None = None, **kwargs, ): """Constructs a new signal generator instance. Parameters ---------- - shg_mgr : instance of SourceHypoGroupManager + shg_mgr The SourceHypoGroupManager instance defining the source hypothesis groups. - dataset_list : list of Dataset instances + dataset_list The list of Dataset instances for which signal events should get generated for. - data_list : list of DatasetData instances + data_list The list of DatasetData instances holding the actual data of each dataset. The order must match the order of ``dataset_list``. - valid_event_field_ranges_dict_list : list of dict | None + valid_event_field_ranges_dict_list If not ``None``, it specifies for each dataset event fields (key) and their valid value range as a 2-element tuple (value). If a generated signal event does not fall into a given field range, the @@ -632,18 +652,18 @@ def _construct_signal_candidates(self): def _get_invalid_events_mask( self, - events, - valid_event_field_ranges_dict, - ): + events: DataFieldRecordArray, + valid_event_field_ranges_dict: dict, + ) -> np.ndarray: """Determines a boolean mask to select invalid events, which do not fulfill the given valid event field ranges. Parameters ---------- - events : instance of DataFieldRecordArray + events The instance of DataFieldRecordArray of length N_events holding the events to check. - valid_event_field_ranges_dict : dict + valid_event_field_ranges_dict The dictionary holding the data field names (key) and their valid value ranges (value). @@ -654,7 +674,7 @@ def _get_invalid_events_mask( Returns ------- - mask : instance of numpy.ndarray + mask The (N_events,)-shaped numpy.ndarray of bool, holding the mask of the invalid events. """ @@ -674,14 +694,14 @@ def _get_invalid_events_mask( def _draw_valid_sig_events_for_dataset_and_shg( self, - rss, - mc, - n_signal, - ds_idx, - valid_event_field_ranges_dict, + rss: RandomStateService, + mc: DataFieldRecordArray, + n_signal: int, + ds_idx: int, + valid_event_field_ranges_dict: dict, shg, - shg_idx, - ): + shg_idx: int, + ) -> DataFieldRecordArray | None: """Draws n_signal valid signal events for the given dataset and source hypothesis group. @@ -690,27 +710,27 @@ def _draw_valid_sig_events_for_dataset_and_shg( Parameters ---------- - rss : instance of RandomStateService + rss The instance of RandomStateService which should be used to draw random numbers from. - mc : instance of DataFieldRecordArray + mc The instance of DataFieldRecordArray holding the monte-carlo events. - n_signal : int + n_signal The number of signal events to draw. - ds_idx : int + ds_idx The index of the dataset. - valid_event_field_ranges_dict : dict + valid_event_field_ranges_dict The dictionary holding the data field names (key) and their valid value ranges (value) for the requested dataset. - shg : instance of SourceHypothesisGroup + shg The instance of SourceHypothesisGroup for which signal events should get drawn. - shg_idx : int + shg_idx The index of the source hypothesis group. Returns ------- - sig_events : instance of DataFieldRecordArray + sig_events The instance of DataFieldRecordArray holding the drawn valid signal events. """ @@ -752,7 +772,9 @@ def change_shg_mgr(self, shg_mgr): self._construct_signal_candidates() - def fluxmodel_scaling_factor(self, src_params_recarray=None, per_source=False): + def fluxmodel_scaling_factor( + self, src_params_recarray: np.ndarray | None = None, per_source: bool = False + ) -> float | np.ndarray: """Scaling factor to convert a mean number of signal events (mu) into a flux normalization as per the definition of the flux models of the source hypothesis groups:: @@ -762,18 +784,18 @@ def fluxmodel_scaling_factor(self, src_params_recarray=None, per_source=False): Parameters ---------- - src_params_recarray : numpy structured ndarray | None + src_params_recarray Must be ``None``. Providing a non-``None`` value raises :class:`NotImplementedError` because this generator derives its scaling factor from precomputed MC signal candidates rather than from a detector signal yield calculation and therefore cannot be re-evaluated at arbitrary parameter values without regenerating those candidates. - per_source : bool + per_source Flag if the scaling factors should be returned for each source individually (True), or as the sum of all these factors (False). The default is False. Returns ------- - scaling_factor : float | (n_sources,)-shaped numpy ndarray + scaling_factor Dimensionless conversion factor (summed for all sources if `per_source = False`) to convert one detected signal event into a flux normalization for the given signal hypothesis. If `per_source` is set to True, a numpy ndarray is returned that contains the flux for each individual source. @@ -820,21 +842,27 @@ def fluxmodel_scaling_factor(self, src_params_recarray=None, per_source=False): return scaling_factor - def generate_signal_events(self, rss, mean, poisson=True, **kwargs): + def generate_signal_events( + self, + rss: RandomStateService, + mean: float, + poisson: bool = True, + src_detsigyield_weights_service: SrcDetSigYieldWeightsService | None = None, + ) -> tuple[int, dict[int, DataFieldRecordArray]]: """Generates a given number of signal events from the signal candidate monte-carlo events. Parameters ---------- - rss : instance of RandomStateService + rss The instance of RandomStateService providing the random number generator state. - mean : float | int + mean The mean number of signal events. If the ``poisson`` argument is set to True, the actual number of generated signal events will be drawn from a Poisson distribution with this given mean value of signal events. - poisson : bool + poisson If set to True, the actual number of generated signal events will be drawn from a Poisson distribution with the given mean value of signal events. @@ -843,9 +871,9 @@ def generate_signal_events(self, rss, mean, poisson=True, **kwargs): Returns ------- - n_signal : int + n_signal The number of actual generated signal events. - signal_events_dict : dict of DataFieldRecordArray + signal_events_dict The dictionary holding the DataFieldRecordArray instances with the generated signal events. Each key of this dictionary represents the dataset index for which the signal events have been generated. @@ -912,7 +940,7 @@ def generate_signal_events(self, rss, mean, poisson=True, **kwargs): redrawn_shg_sig_events = self._draw_valid_sig_events_for_dataset_and_shg( rss=rss, mc=mc, - n_signal=n_redraw_events, + n_signal=int(n_redraw_events), ds_idx=ds_idx, valid_event_field_ranges_dict=valid_event_field_ranges_dict, shg=shg, diff --git a/skyllh/core/signalpdf.py b/skyllh/core/signalpdf.py index f5e298b8fe..43e482505a 100644 --- a/skyllh/core/signalpdf.py +++ b/skyllh/core/signalpdf.py @@ -2,18 +2,29 @@ likelihood function. """ +from collections.abc import Sequence + import numpy as np from skyllh.core import ( tool, ) +from skyllh.core.flux_model import TimeFluxProfile from skyllh.core.interpolate import ( GridManifoldInterpolationMethod, Linear1DGridManifoldInterpolationMethod, ) +from skyllh.core.livetime import Livetime from skyllh.core.logging import ( get_logger, ) +from skyllh.core.parameters import ( + Parameter, + ParameterGrid, + ParameterGridSet, + ParameterModelMapper, + ParameterSet, +) from skyllh.core.pdf import ( PDF, IsSignalPDF, @@ -29,9 +40,8 @@ from skyllh.core.source_hypo_grouping import ( SourceHypoGroupManager, ) -from skyllh.core.timing import ( - TaskTimer, -) +from skyllh.core.timing import TaskTimer, TimeLord +from skyllh.core.trialdata import TrialDataManager from skyllh.core.utils.coords import ( angular_separation, ) @@ -55,19 +65,25 @@ class GaussianPSFPointLikeSourceSignalSpatialPDF(SpatialPDF, IsSignalPDF): declination of the point-like sources, respectively. """ - def __init__(self, ra_range=None, dec_range=None, pd_event_data_field_name=None, **kwargs): + def __init__( + self, + ra_range: tuple | None = None, + dec_range: tuple | None = None, + pd_event_data_field_name: str | None = None, + **kwargs, + ): """Creates a new spatial signal PDF for point-like sources with a gaussian point-spread-function (PSF). Parameters ---------- - ra_range : 2-element tuple | None + ra_range The range in right-ascension this spatial PDF is valid for. If set to None, the range (0, 2pi) is used. - dec_range : 2-element tuple | None + dec_range The range in declination this spatial PDF is valid for. If set to None, the range (-pi/2, +pi/2) is used. - pd_event_data_field_name : str | None + pd_event_data_field_name The probability density values can be pre-calculated by the user. This specifies the name of the event data field, where these values are stored. @@ -97,35 +113,35 @@ def pd_event_data_field_name(self, name): ) self._pd_event_data_field_name = name - def calculate_pd(self, tdm): + def calculate_pd(self, tdm: TrialDataManager) -> np.ndarray: """Calculates the gaussian PSF probability density values for all events and sources. Parameters ---------- - tdm : instance of TrialDataManager + tdm The instance of TrialDataManager holding the trial event data for which to calculate the PDF values. The following data fields need to be present: - src_array : numpy record ndarray + src_array The numpy record ndarray with the following data fields: - ra : float + ra The right-ascension of the point-like source. - dec : float + dec The declination of the point-like source. - ra : float + ra The right-ascension in radian of the data event. - dec : float + dec The declination in radian of the data event. - ang_err: float + ang_err The reconstruction uncertainty in radian of the data event. Returns ------- - pd : instance of numpy ndarray + pd The (N_values,)-shaped numpy ndarray holding the probability density for each event. The length of this 1D array depends on the number of sources and the events belonging to those sources. In the worst @@ -138,7 +154,9 @@ def calculate_pd(self, tdm): dec = get_data('dec') sigma = get_data('ang_err') - (src_idxs, evt_idxs) = tdm.src_evt_idxs + _src_evt_idxs = tdm.src_evt_idxs + assert _src_evt_idxs is not None + (src_idxs, evt_idxs) = _src_evt_idxs src_ra = np.take(src_array['ra'], src_idxs) src_dec = np.take(src_array['dec'], src_idxs) @@ -152,47 +170,49 @@ def calculate_pd(self, tdm): return pd - def get_pd(self, tdm, params_recarray=None, tl=None): + def get_pd( + self, tdm: TrialDataManager, params_recarray: np.ndarray | None = None, tl: TimeLord | None = None + ) -> tuple[np.ndarray, dict]: """Calculates the spatial signal probability density of each event for all sources. Parameters ---------- - tdm : instance of TrialDataManager + tdm The instance of TrialDataManager holding the trial event data for which to calculate the PDF values. The following data fields need to be present: - src_array : numpy record ndarray + src_array The numpy record ndarray with the following data fields: - ra : float + ra The right-ascension of the point-like source. - dec : float + dec The declination of the point-like source. - ra : float + ra The right-ascension in radian of the data event. - dec : float + dec The declination in radian of the data event. - ang_err: float + ang_err The reconstruction uncertainty in radian of the data event. In case the probability density values were pre-calculated, - params_recarray : None + params_recarray Unused interface argument. - tl : TimeLord instance | None + tl The optional TimeLord instance to use for measuring timing information. Returns ------- - pd : instance of numpy ndarray + pd The (N_values,)-shaped numpy ndarray holding the probability density for each event. The length of this 1D array depends on the number of sources and the events belonging to those sources. In the worst case the length is N_sources * N_trial_events. - grads : dict + grads The dictionary holding the gradients of the probability density w.r.t. each fit parameter. By definition this PDF does not depend on any fit parameters and hence, this dictionary is empty. @@ -201,7 +221,7 @@ def get_pd(self, tdm, params_recarray=None, tl=None): logger = get_logger(f'{__name__}.{classname(self)}.get_pd') # Check if the probability density was pre-calculated. - if self._pd_event_data_field_name in tdm: + if self._pd_event_data_field_name is not None and self._pd_event_data_field_name in tdm: if self._cfg.is_tracing_enabled: logger.debug( 'Retrieve precalculated probability density values from ' @@ -236,7 +256,7 @@ class RayleighPSFPointSourceSignalSpatialPDF(SpatialPDF, IsSignalPDF): right-ascension and declination of the point-like sources, respectively. """ - def __init__(self, ra_range=None, dec_range=None, **kwargs): + def __init__(self, ra_range: tuple | None = None, dec_range: tuple | None = None, **kwargs): r"""Creates a new spatial signal PDF for point-like sources with a Rayleigh point-spread-function (PSF). @@ -264,6 +284,7 @@ def initialize_for_new_trial(self, tdm, tl=None, **kwargs): """ get_data = tdm.get_data + assert tdm.src_evt_idxs is not None (_, evt_idxs) = tdm.src_evt_idxs psi = get_data('psi') @@ -272,35 +293,37 @@ def initialize_for_new_trial(self, tdm, tl=None, **kwargs): self._pd = 0.5 / (np.pi * np.sin(psi)) * (psi / sigma_sq) * np.exp(-0.5 * (psi**2 / sigma_sq)) - def get_pd(self, tdm, params_recarray=None, tl=None): + def get_pd( + self, tdm: TrialDataManager, params_recarray: np.ndarray | None = None, tl: TimeLord | None = None + ) -> tuple[np.ndarray, dict]: """Calculates the spatial signal probability density of each event for all sources. Parameters ---------- - tdm : instance of TrialDataManager + tdm The instance of TrialDataManager holding the trial event data for which to calculate the PDF values. The following data fields need to be present: - psi : float + psi The opening angle in radian between the source direction and the reconstructed muon direction. - ang_err: float + ang_err The reconstruction uncertainty in radian of the data event. - params_recarray : None + params_recarray Unused interface argument. - tl : instance of TimeLord | None + tl The optional instance of TimeLord to use for measuring timing information. Returns ------- - pd : (N_values,)-shaped numpy ndarray + pd The (N_values,)-shaped 1D numpy ndarray holding the probability density value for each event and source in unit 1/rad. - grads : dict + grads The dictionary holding the gradients of the probability density w.r.t. each global fit parameter. By definition this PDF does not depend on any global fit parameters and hence, this dictionary is @@ -325,16 +348,16 @@ class SignalTimePDF( into account. """ - def __init__(self, livetime, time_flux_profile, **kwargs): + def __init__(self, livetime: Livetime, time_flux_profile: TimeFluxProfile, **kwargs): """Creates a new signal time PDF instance for a given time flux profile and detector live time. Parameters ---------- - livetime : instance of Livetime + livetime An instance of Livetime, which provides the detector live-time information. - time_flux_profile : instance of TimeFluxProfile + time_flux_profile The signal's time flux profile. .. note:: @@ -349,42 +372,45 @@ class when calculating the probability density values! def _calculate_pd( self, - tdm, - params_recarray, - tl=None, - ): + tdm: TrialDataManager, + params_recarray: np.ndarray, + tl: TimeLord | None = None, + ) -> np.ndarray: """Calculates the probability density values for the given trial data and source parameters. Parameters ---------- - tdm : instance of TrialDataManager + tdm The instance of TrialDataManager holding the trial data. The following data fields must exist: ``'time'`` : float The time of the event. - params_recarray : instance of structured ndarray + params_recarray The structured numpy ndarray of length N_sources holding the local parameter names and values of the sources. - tl : instance of TimeLord | None + tl The optional instance of TimeLord that should be used to measure timing information. Returns ------- - pd : instance of ndarray + pd The (N_values,)-shaped numpy ndarray holding the probability density values for each trial data event and source. """ - (src_idxs, evt_idxs) = tdm.src_evt_idxs + _src_evt_idxs = tdm.src_evt_idxs + assert _src_evt_idxs is not None + (src_idxs, evt_idxs) = _src_evt_idxs n_values = len(evt_idxs) pd = np.zeros((n_values,), dtype=np.float64) events_time = tdm.get_data('time') for src_idx, src_params_row in enumerate(params_recarray): + assert params_recarray.dtype.fields is not None params = dict(zip(params_recarray.dtype.fields.keys(), src_params_row, strict=True)) # Update the time flux profile if its parameter values have changed @@ -415,6 +441,18 @@ def initialize_for_new_trial( tl=None, **kwargs, ): + """Initializes this time PDF for a new trial. If the time PDF does not + depend on any global floating parameters, the PDF values are + pre-calculated for the trial data. + + Parameters + ---------- + tdm + The instance of TrialDataManager holding the trial event data. + tl + The optional instance of TimeLord that should be used to collect + timing information about this method. + """ # Check if this time PDF is not constant and does depend on any global # floating parameters. If that's not the case we can pre-calculate the # PDF values. @@ -434,6 +472,7 @@ def initialize_for_new_trial( # pre-calculate the PDF values. if self.pmm is None: + assert tdm.n_sources is not None params_recarray = np.empty((tdm.n_sources,), dtype=[]) else: params_recarray = self.pmm.create_src_params_recarray() @@ -442,16 +481,16 @@ def initialize_for_new_trial( def get_pd( self, - tdm, - params_recarray, - tl=None, - ): + tdm: TrialDataManager, + params_recarray: np.ndarray | None = None, + tl: TimeLord | None = None, + ) -> tuple[np.ndarray, dict]: """Calculates the signal time probability density of each event for the given set of time parameter values for each source. Parameters ---------- - tdm : instance of TrialDataManager + tdm The instance of TrialDataManager holding the trial event data for which to calculate the PDF value. The following data fields must exist: @@ -459,19 +498,19 @@ def get_pd( ``'time'`` : float The time of the event. - params_recarray : instance of numpy structured ndarray + params_recarray The numpy structured ndarray holding the local parameter values for each source. - tl : instance of TimeLord | None + tl The optional TimeLord instance that should be used to measure timing information. Returns ------- - pd : instance of numpy ndarray + pd The (N_values,)-shaped 1D numpy ndarray holding the probability density value for each trial event and source. - grads : dict + grads The dictionary holding the gradients of the probability density w.r.t. each global fit parameter. """ @@ -479,6 +518,7 @@ def get_pd( if self._pd is not None: return (self._pd, {}) + assert params_recarray is not None pd = self._calculate_pd(tdm=tdm, params_recarray=params_recarray, tl=tl) return (pd, {}) @@ -512,12 +552,12 @@ class SignalMultiDimGridPDFSet( def __init__( self, - pmm, - param_set, - param_grid_set, + pmm: ParameterModelMapper, + param_set: Parameter | Sequence[Parameter] | ParameterSet, + param_grid_set: ParameterGrid | ParameterGridSet, gridparams_pdfs, - interpol_method_cls=None, - use_same_photospline_bfi_for_all_pdfs=False, + interpol_method_cls: type[GridManifoldInterpolationMethod] | None = None, + use_same_photospline_bfi_for_all_pdfs: bool = False, **kwargs, ): """Creates a new MultiDimGridPDFSet instance, which holds a set of @@ -525,23 +565,23 @@ def __init__( Parameters ---------- - pmm : instance of ParameterModelMapper + pmm The instance of ParameterModelMapper that defines the mapping of the global parameters to local model parameters. - param_set : instance of Parameter | sequence of instance of Parameter | instance of ParameterSet + param_set The set of parameters defining the parameters of this PDF. - param_grid_set : ParameterGrid instance | ParameterGridSet instance + param_grid_set The set of ParameterGrid instances, which define the grid values of the model parameters, the given MultiDimGridPDF instances belong to. - gridparams_pdfs : sequence of (dict, MultiDimGridPDF) tuples + gridparams_pdfs The sequence of 2-element tuples which define the mapping of grid values to PDF instances. - interpol_method_cls : subclass of GridManifoldInterpolationMethod + interpol_method_cls The class specifying the interpolation method. This must be a subclass of ``GridManifoldInterpolationMethod``. If set to None, the default grid manifold interpolation method ``Linear1DGridManifoldInterpolationMethod`` will be used. - use_same_photospline_bfi_for_all_pdfs : bool + use_same_photospline_bfi_for_all_pdfs Flag if the same basis function indices (bfi) should be used for all PDFs when photospline tables are used. Default is ``False``. This should be set to ``True`` if all photospline tables share the @@ -589,20 +629,20 @@ def interpol_method_cls(self, cls): def _get_pdf_for_interpol_param_values( self, - interpol_param_values, - ): + interpol_param_values: np.ndarray, + ) -> MultiDimGridPDF: """Retrieves the PDF for the given set of interpolation parameter values. Parameters ---------- - interpol_param_values : instance of numpy ndarray + interpol_param_values The (N_interpol_params,)-shaped numpy ndarray holding the values of the interpolation parameters. Returns ------- - pdf : instance of MultiDimGridPDF + pdf The requested PDF instance. """ gridparams = dict(zip(self._interpol_param_names, interpol_param_values, strict=True)) @@ -613,34 +653,34 @@ def _get_pdf_for_interpol_param_values( def _evaluate_pdfs( self, - tdm, - eventdata, - gridparams_recarray, - n_values, + tdm: TrialDataManager, + eventdata: np.ndarray, + gridparams_recarray: np.ndarray, + n_values: int, tl=None, - ): + ) -> np.ndarray: """Evaluates the PDFs for the given event data. The particular PDF is selected based on the grid parameter values for each model. This method is called by the interpolation method. Parameters ---------- - tdm : instance of TrialDataManager + tdm The instance of TrialDataManager holding the trial event data. - eventdata : instance of numpy ndarray + eventdata The (V,N_values)-shaped numpy ndarray holding the event data for the PDF evaluation. - gridparams_recarray : instance of numpy structured ndarray + gridparams_recarray The numpy structured ndarray of length N_sources with the parameter names and values needed for the interpolation on the grid for all sources. If the length of this structured array is 1, the set of parameters will be used for all sources. - n_values : int + n_values The size of the output array. Returns ------- - pd : instance of ndarray + pd The (N_values,)-shaped numpy ndarray holding the probability density values for each event. """ @@ -658,7 +698,9 @@ def _evaluate_pdfs( pd = np.empty(n_values, dtype=np.float64) - (src_idxs, _) = tdm.src_evt_idxs + _src_evt_idxs = tdm.src_evt_idxs + assert _src_evt_idxs is not None + (src_idxs, _) = _src_evt_idxs v_start = 0 for sidx, interpol_param_values in enumerate(gridparams_recarray): @@ -679,8 +721,8 @@ def _evaluate_pdfs( def assert_is_valid_for_trial_data( self, - tdm, - tl=None, + tdm: TrialDataManager, + tl: TimeLord | None = None, **kwargs, ): """Checks if the PDFs of this PDFSet instance are valid for all the @@ -694,9 +736,9 @@ def assert_is_valid_for_trial_data( Parameters ---------- - tdm : instance of TrialDataManager + tdm The instance of TrialDataManager holding the trial data events. - tl : instance of TimeLord | None + tl The optional instance of TimeLord for measuring timing information. Raises @@ -708,8 +750,8 @@ def assert_is_valid_for_trial_data( def initialize_for_new_trial( self, - tdm, - tl=None, + tdm: TrialDataManager, + tl: TimeLord | None = None, **kwargs, ): """This method is called whenever a new trial data is initialized. It @@ -720,9 +762,9 @@ def initialize_for_new_trial( Parameters ---------- - tdm : instance of TrialDataManager + tdm The instance of TrialDataManager holding the new trial data events. - tl : instance of TimeLord | None + tl The optional instance of TimeLord for measuring timing information. """ super().initialize_for_new_trial(tdm=tdm, tl=tl, **kwargs) @@ -744,33 +786,33 @@ def initialize_for_new_trial( logger = get_logger(f'{__name__}.{classname(self)}.initialize_for_new_trial') logger.info('Falling back to the slower photospline evaluation.') - def get_pd( + def get_pd( # pyright: ignore[reportIncompatibleMethodOverride] self, - tdm, - params_recarray, - tl=None, - ): + tdm: TrialDataManager, + params_recarray: np.ndarray | None, + tl: TimeLord | None = None, + ) -> tuple[np.ndarray, dict]: """Calculates the probability density for each event, given the given parameter values. Parameters ---------- - tdm : instance of TrialDataManager + tdm The instance of TrialDataManager that will be used to get the data from the trial events. - params_recarray : instance of structured ndarray | None + params_recarray The numpy record ndarray holding the parameter name and values for each source model. - tl : instance of TimeLord | None + tl The optional instance of TimeLord to use for measuring timing information. Returns ------- - pd : instance of numpy ndarray + pd The (N_values,)-shaped numpy ndarray holding the probability density value for each source and event. - grads : dict + grads The dictionary holding the PDF gradient value for each event w.r.t. each global fit parameter. The key of the dictionary is the ID of the global fit parameter. @@ -779,6 +821,9 @@ def get_pd( """ logger = get_logger(f'{__name__}.{classname(self)}.get_pd') + assert params_recarray is not None + assert self._cache_eventdata is not None + # Get the interpolated PDF values for the arbitrary parameter values. # The (D,N_events)-shaped grads_arr ndarray contains the gradient of the # probability density w.r.t. each of the D parameters, which are defined @@ -786,6 +831,7 @@ def get_pd( # the parameter grids. with TaskTimer(tl, 'Call interpolate method to get probability densities for all events.'): if self._cfg.is_tracing_enabled: + assert params_recarray.dtype.fields is not None logger.debug( 'Call interpol_method with ' f'params_recarray={params_recarray} of fields ' @@ -803,6 +849,7 @@ def get_pd( grads = {} tdm_n_sources = tdm.n_sources + assert self.pmm is not None for fitparam_id in range(self.pmm.n_global_floating_params): grad = np.zeros((tdm.get_n_values(),), dtype=np.float64) @@ -849,10 +896,10 @@ class SignalSHGMappedMultiDimGridPDFSet( def __init__( self, - shg_mgr, - pmm, + shg_mgr: SourceHypoGroupManager, + pmm: ParameterModelMapper, shgidxs_pdf_list, - use_same_photospline_bfi_for_all_pdfs=False, + use_same_photospline_bfi_for_all_pdfs: bool = False, **kwargs, ): """Creates a new SignalSHGMappedMultiDimGridPDFSet instance, which holds @@ -861,16 +908,16 @@ def __init__( Parameters ---------- - shg_mgr : instance of SourceHypoGroupManager + shg_mgr The instance of SourceHypoGroupManager that defines the source hypothesis groups and their sources. - pmm : instance of ParameterModelMapper + pmm The instance of ParameterModelMapper which defines the mapping of global parameters to local source parameters. - shgidxs_pdf_list : sequence of (shg_idxs, MultiDimGridPDF) tuples + shgidxs_pdf_list The sequence of 2-element tuples which define the mapping of the source hypothesis groups to a PDF instance. - use_same_photospline_bfi_for_all_pdfs : bool + use_same_photospline_bfi_for_all_pdfs Flag if the same basis function indices (bfi) should be used for all PDFs when photospline tables are used. Default is ``False``. This should be set to ``True`` if all photospline tables share the @@ -911,8 +958,8 @@ def shg_mgr(self): def initialize_for_new_trial( self, - tdm, - tl=None, + tdm: TrialDataManager, + tl: TimeLord | None = None, **kwargs, ): """This method is called whenever a new trial data is initialized. It @@ -923,9 +970,9 @@ def initialize_for_new_trial( Parameters ---------- - tdm : instance of TrialDataManager + tdm The instance of TrialDataManager holding the new trial data events. - tl : instance of TimeLord | None + tl The optional instance of TimeLord for measuring timing information. """ super().initialize_for_new_trial(tdm=tdm, tl=tl, **kwargs) @@ -949,33 +996,33 @@ def initialize_for_new_trial( logger = get_logger(f'{__name__}.{classname(self)}.initialize_for_new_trial') logger.info('Falling back to the slower photospline evaluation.') - def get_pd( + def get_pd( # pyright: ignore[reportIncompatibleMethodOverride] self, - tdm, - params_recarray, - tl=None, - ): + tdm: TrialDataManager, + params_recarray: np.ndarray | None, + tl: TimeLord | None = None, + ) -> tuple[np.ndarray, dict]: """Calculates the probability density for each event, given the given parameter values. Parameters ---------- - tdm : instance of TrialDataManager + tdm The instance of TrialDataManager that will be used to get the data from the trial events. - params_recarray : instance of structured ndarray | None + params_recarray The numpy record ndarray holding the parameter name and values for each source model. - tl : instance of TimeLord | None + tl The optional instance of TimeLord to use for measuring timing information. Returns ------- - pd : instance of numpy ndarray + pd The (N_values,)-shaped numpy ndarray holding the probability density value for each event. - grads : dict + grads The dictionary holding the PDF gradient value for each event w.r.t. each global fit parameter. The key of the dictionary is the ID of the global fit parameter. @@ -986,7 +1033,9 @@ def get_pd( """ pd = np.zeros((tdm.get_n_values(),), dtype=np.float64) - src_idxs = tdm.src_evt_idxs[0] + _src_evt_idxs = tdm.src_evt_idxs + assert _src_evt_idxs is not None + src_idxs = _src_evt_idxs[0] src_idxs_arr = np.arange(self._shg_mgr.n_sources) # Loop over the individual PDFs (via their key). diff --git a/skyllh/core/smoothing.py b/skyllh/core/smoothing.py index 8bcc5bd266..daf5f83525 100644 --- a/skyllh/core/smoothing.py +++ b/skyllh/core/smoothing.py @@ -1,4 +1,5 @@ import abc +from collections.abc import Sequence import numpy as np import scipy.signal @@ -17,20 +18,21 @@ class HistSmoothingMethod( """Abstract base class for implementing a histogram smoothing method.""" def __init__(self, **kwargs): + """Creates a new instance of HistSmoothingMethod.""" super().__init__(**kwargs) @abc.abstractmethod - def smooth(self, h): + def smooth(self, h: np.ndarray) -> np.ndarray: """This method is supposed to smooth the given histogram h. Parameters ---------- - h : N-dimensional ndarray + h The ndarray holding histogram bin values. Returns ------- - smoothed_h : N-dimensional ndarray + smoothed_h The array holding the smoothed histogram bin values. """ @@ -41,19 +43,20 @@ class NoHistSmoothingMethod( """This class implements a no-shoothing histogram method.""" def __init__(self, **kwargs): + """Creates a new instance of NoHistSmoothingMethod.""" super().__init__(**kwargs) - def smooth(self, h): + def smooth(self, h: np.ndarray) -> np.ndarray: """Does not perform any smoothing and just returns the input histogram. Parameters ---------- - h : N-dimensional ndarray + h The ndarray holding histogram bin values. Returns ------- - h : N-dimensional ndarray + h The input histogram array. """ return h @@ -68,14 +71,14 @@ class NeighboringBinHistSmoothingMethod( def __init__( self, - axis_kernel_arrays, + axis_kernel_arrays: Sequence[np.ndarray], **kwargs, ): """Constructs a new neighboring bin histogram smoothing method. Parameters ---------- - axis_kernel_arrays: sequence of 1D ndarrays + axis_kernel_arrays The sequence of smoothing kernel arrays, one for each axis. If an axis should not get smoothed, the UNSMOOTH_AXIS constant should be used for that axis' smoothing kernel array. @@ -99,19 +102,19 @@ def ndim(self): """ return self._ndim - def smooth(self, h): + def smooth(self, h: np.ndarray) -> np.ndarray: """Smoothes the given histogram array h with the internal kernel array k. Both arrays must have the same dimensionality. The shape values of k must be smaller than or equal to the shape values of h. Parameters ---------- - h : N-dimensional ndarray + h The ndarray holding histogram bin values. Returns ------- - smoothed_h : N-dimensional ndarray. + smoothed_h """ if h.ndim != self._ndim: raise ValueError( @@ -141,6 +144,14 @@ class SmoothingFilter: """ def __init__(self, axis_kernel_array, **kwargs): + """Creates a new instance of SmoothingFilter. + + Parameters + ---------- + axis_kernel_array + The kernel array defining how many neighboring bins of a histogram + bin should be used to smooth that histogram bin. + """ super().__init__(**kwargs) self.axis_kernel_array = axis_kernel_array @@ -165,12 +176,12 @@ class BlockSmoothingFilter( block is specified via the nbins argument. """ - def __init__(self, nbins, **kwargs): + def __init__(self, nbins: int, **kwargs): """Creates a new BlockSmoothingFilter instance. Parameters ---------- - nbins : int + nbins The number of neighboring bins into one direction of a histogram bin, which should be used to smooth that histogram bin. """ @@ -195,14 +206,14 @@ class GaussianSmoothingFilter( def __init__( self, - nbins, + nbins: int, **kwargs, ): """Creates a new GaussianSmoothingFilter instance. Parameters ---------- - nbins : int + nbins The number of neighboring bins into one direction of a histogram bin, which should be used to smooth that histogram bin. """ diff --git a/skyllh/core/source_hypo_grouping.py b/skyllh/core/source_hypo_grouping.py index 006d06e316..3fde46d7b8 100644 --- a/skyllh/core/source_hypo_grouping.py +++ b/skyllh/core/source_hypo_grouping.py @@ -3,6 +3,8 @@ analysis. """ +from collections.abc import Sequence + import numpy as np from skyllh.core.detsigyield import ( @@ -36,24 +38,31 @@ class SourceHypoGroup(SourceHypoGroup_t): and signal generation methods. """ - def __init__(self, sources, fluxmodel, detsigyield_builders, sig_gen_method=None, **kwargs): + def __init__( + self, + sources: SourceModel | Sequence[SourceModel], + fluxmodel: FluxModel, + detsigyield_builders: DetSigYieldBuilder | Sequence[DetSigYieldBuilder], + sig_gen_method: SignalGenerationMethod | None = None, + **kwargs, + ): """Constructs a new source hypothesis group. Parameters ---------- - sources : SourceModel | sequence of SourceModel + sources The source or sequence of sources that define the source group. - fluxmodel : instance of FluxModel + fluxmodel The FluxModel instance that applies to the list of sources of the group. - detsigyield_builders : sequence of DetSigYieldBuilder instances + detsigyield_builders The sequence of detector signal yield builder instances, which should be used to create the detector signal yield for the sources of this group. Each element is the detector signal yield builder for the particular dataset, if several datasets are used. If this list contains only one builder, it should be used for all datasets. - sig_gen_method : SignalGenerationMethod instance | None + sig_gen_method The instance of SignalGenerationMethod that implements the signal generation for the specific detector and source hypothesis. It can be set to None, which means, no signal can be generated. Useful for @@ -168,12 +177,12 @@ def __str__(self): return s - def get_source_weights(self): + def get_source_weights(self) -> np.ndarray | None: """Gets the weight from each source of this source hypothesis group. Returns ------- - weights : numpy ndarray | None + weights The (N_sources,)-shaped numpy ndarray holding the theoretical weight of each source. It is ``None`` if any of the individual source weights is None. @@ -196,13 +205,12 @@ class SourceHypoGroupManager: way. """ - def __init__(self, src_hypo_groups=None, **kwargs): + def __init__(self, src_hypo_groups: 'SourceHypoGroup | Sequence[SourceHypoGroup] | None' = None, **kwargs): """Creates a new source hypothesis group manager instance. Parameters ---------- - src_hypo_groups : SourceHypoGroup instance | - sequence of SourceHypoGroup instances | None + src_hypo_groups The SourceHypoGroup instances to initialize the manager with. """ super().__init__(**kwargs) @@ -272,13 +280,13 @@ def __str__(self): return s - def _extend_sidx_to_gidx_gsidx_map_arr(self, shg): + def _extend_sidx_to_gidx_gsidx_map_arr(self, shg: 'SourceHypoGroup'): """Extends the source index to (group index, group source index) map array by one source hypo group. Parameters ---------- - shg : SourceHypoGroup instance + shg The SourceHypoGroup instance for which the map array should get extented. """ @@ -287,7 +295,13 @@ def _extend_sidx_to_gidx_gsidx_map_arr(self, shg): arr[:, 1] = np.arange(shg.n_sources) # Group source index. self._sidx_to_gidx_gsidx_map_arr = np.vstack((self._sidx_to_gidx_gsidx_map_arr, arr)) - def create_source_hypo_group(self, sources, fluxmodel, detsigyield_builders, sig_gen_method=None): + def create_source_hypo_group( + self, + sources: SourceModel | Sequence[SourceModel], + fluxmodel: FluxModel, + detsigyield_builders: Sequence[DetSigYieldBuilder], + sig_gen_method: SignalGenerationMethod | None = None, + ): """Creates and adds a source hypothesis group to this source hypothesis group manager. A source hypothesis group shares sources of the same source model with the same flux model and hence the same detector signal @@ -295,19 +309,19 @@ def create_source_hypo_group(self, sources, fluxmodel, detsigyield_builders, sig Parameters ---------- - sources : SourceModel | sequence of SourceModel + sources The source or sequence of sources that define the source group. - fluxmodel : instance of FluxModel + fluxmodel The FluxModel instance that applies to the list of sources of the group. - detsigyield_builders : sequence of DetSigYieldBuilder instances + detsigyield_builders The sequence of detector signal yield builder instances, which should be used to create the detector signal yield for the sources of this group. Each element is the detector signal yield builder for the particular dataset, if several datasets are used. If this list contains only one builder, it should be used for all datasets. - sig_gen_method : instance of SignalGenerationMethod | None + sig_gen_method The SignalGenerationMethod instance that implements the detector and source hypothesis specific signal generation. It can be set to None which means no signal can be generated. @@ -327,72 +341,72 @@ def create_source_hypo_group(self, sources, fluxmodel, detsigyield_builders, sig # array. self._extend_sidx_to_gidx_gsidx_map_arr(group) - def get_fluxmodel_by_src_idx(self, src_idx): + def get_fluxmodel_by_src_idx(self, src_idx: int) -> FluxModel: """Retrieves the FluxModel instance for the source specified by its source index. Parameters ---------- - src_idx : int + src_idx The index of the source, which must be in the range [0, N_sources-1]. Returns ------- - fluxmodel : instance of FluxModel + fluxmodel The FluxModel instance that applies to the specified source. """ gidx = self._sidx_to_gidx_gsidx_map_arr[src_idx, 0] return self._shg_list[gidx]._fluxmodel - def get_detsigyield_builder_list_by_src_idx(self, src_idx): + def get_detsigyield_builder_list_by_src_idx(self, src_idx: int) -> list[DetSigYieldBuilder]: """Retrieves the list of DetSigYieldBuilder instances for the source specified by its source index. Parameters ---------- - src_idx : int + src_idx The index of the source, which must be in the range [0, N_sources-1]. Returns ------- - detsigyield_builder_list : list of DetSigYieldBuilder instances + detsigyield_builder_list The list of DetSigYieldBuilder instances that apply to the specified source. """ gidx = self._sidx_to_gidx_gsidx_map_arr[src_idx, 0] return self._shg_list[gidx]._detsigyield_builder_list - def get_src_mask_of_shg(self, shg_idx): + def get_src_mask_of_shg(self, shg_idx: int) -> np.ndarray: """Creates a source mask for the sources of the ``shg_idx`` th source hypothesis group. Parameters ---------- - shg_idx : int + shg_idx The index of the source hypothesis group. Returns ------- - src_mask : instance of numpy ndarray + src_mask The (N_sources,)-shaped numpy ndarray of bool holding the mask for selecting the sources of the given source hypothesis group. """ return self._sidx_to_gidx_gsidx_map_arr[:, 0] == shg_idx - def get_src_idxs_of_shg(self, shg_idx): + def get_src_idxs_of_shg(self, shg_idx: int) -> np.ndarray: """Creates an array of indices of sources that belong to the given source hypothesis group. Parameters ---------- - shg_idx : int + shg_idx The index of the source hypothesis group. Returns ------- - src_idxs : instance of numpy ndarray + src_idxs The numpy ndarray of int holding the indices of the sources that belong to the given source hypothesis group. """ diff --git a/skyllh/core/source_model.py b/skyllh/core/source_model.py index 2c696a63a3..27b9e05010 100644 --- a/skyllh/core/source_model.py +++ b/skyllh/core/source_model.py @@ -4,6 +4,8 @@ model for a point-like source at a given location in the sky. """ +from collections.abc import Callable, Sequence + import numpy as np from skyllh.core.model import ( @@ -26,16 +28,18 @@ class SourceModel( relative weight w.r.t. other sources. """ - def __init__(self, name=None, classification=None, weight=None, **kwargs): + def __init__( + self, name: str | None = None, classification: str | None = None, weight: float | None = None, **kwargs + ): """Creates a new source model instance. Parameters ---------- - name : str | None + name The name of the source model. - classification : str | None + classification The astronomical classification of the source. - weight : float | None + weight The relative weight of the source w.r.t. other sources. If set to None, unity will be used. """ @@ -74,17 +78,18 @@ class SourceModelCollection( """ @staticmethod - def cast(obj, errmsg=None, **kwargs): + def cast( # pyright: ignore[reportIncompatibleMethodOverride] + obj: 'SourceModel | Sequence[SourceModel] | SourceModelCollection | None', errmsg: str | None = None, **kwargs + ) -> 'SourceModelCollection': """Casts the given object to a SourceModelCollection object. If the cast fails, a TypeError with the given error message is raised. Parameters ---------- - obj : SourceModel | sequence of SourceModel | SourceModelCollection | - None + obj The object that should be casted to SourceModelCollection. If set to None, an empty SourceModelCollection is created. - errmsg : str | None + errmsg The error message if the cast fails. If set to None, a generic error message will be used. @@ -114,15 +119,15 @@ def cast(obj, errmsg=None, **kwargs): errmsg = f'Cast of object "{obj!s}" of type "{typename(obj)}" to SourceModelCollection failed!' raise TypeError(errmsg) - def __init__(self, sources=None, source_type=None, **kwargs): + def __init__(self, sources=None, source_type: type | None = None, **kwargs): """Creates a new source collection. Parameters ---------- - sources : sequence of source_type instances | None + sources The sequence of sources this collection should be initalized with. If set to None, an empty SourceModelCollection instance is created. - source_type : type | None + source_type The type of the source. If set to None (default), SourceModel will be used. """ @@ -151,12 +156,12 @@ class IsPointlike: def __init__( self, - ra_func_instance=None, - get_ra_func=None, - set_ra_func=None, - dec_func_instance=None, - get_dec_func=None, - set_dec_func=None, + ra_func_instance: object | None = None, + get_ra_func: Callable | None = None, + set_ra_func: Callable | None = None, + dec_func_instance: object | None = None, + get_dec_func: Callable | None = None, + set_dec_func: Callable | None = None, **kwargs, ): """Constructor method. Gets called when the an instance of a class is @@ -164,25 +169,25 @@ def __init__( Parameters ---------- - ra_func_instance : object + ra_func_instance The instance object the right-ascention property's getter and setter functions are defined in. - get_ra_func : callable + get_ra_func The callable object of the getter function of the right-ascention property. It must have the call signature `__call__(ra_func_instance)`. - set_ra_func : callable + set_ra_func The callable object of the setter function of the right-ascention property. It must have the call signature `__call__(ra_func_instance, value)`. - dec_func_instance : object + dec_func_instance The instance object the declination property's getter and setter functions are defined in. - get_dec_func : object + get_dec_func The callable object of the getter function of the declination property. It must have the call signature `__call__(dec_func_instance)`. - set_dec_func : object + set_dec_func The callable object of the setter function of the declination property. It must have the call signature `__call__(dec_func_instance, value)`. @@ -200,21 +205,25 @@ def __init__( @property def ra(self): """The right-ascention coordinate of the point-like source.""" + assert self._get_ra_func is not None return self._get_ra_func(self._ra_func_instance) @ra.setter def ra(self, v): v = float_cast(v, 'The ra property must be castable to type float!') + assert self._set_ra_func is not None self._set_ra_func(self._ra_func_instance, v) @property def dec(self): """The declination coordinate of the point-like source.""" + assert self._get_dec_func is not None return self._get_dec_func(self._dec_func_instance) @dec.setter def dec(self, v): v = float_cast(v, 'The dec property must be castable to type float!') + assert self._set_dec_func is not None self._set_dec_func(self._dec_func_instance, v) @@ -223,19 +232,19 @@ class PointLikeSource(SourceModel, IsPointlike): object in the sky at a given location (right-ascention and declination). """ - def __init__(self, ra, dec, name=None, weight=None, **kwargs): + def __init__(self, ra: float, dec: float, name: str | None = None, weight: float | None = None, **kwargs): """Creates a new PointLikeSource instance for defining a point-like source. Parameters ---------- - ra : float + ra The right-ascention coordinate of the source in radians. - dec : float + dec The declination coordinate of the source in radians. - name : str | None + name The name of the source. - weight : float | None + weight The relative weight of the source w.r.t. other sources. If set to None, unity will be used. """ @@ -255,15 +264,19 @@ def __init__(self, ra, dec, name=None, weight=None, **kwargs): self.dec = dec def _get_ra(self): + """Returns the right-ascention of the source in radians.""" return self._ra def _set_ra(self, ra): + """Sets the right-ascention of the source in radians.""" self._ra = ra def _get_dec(self): + """Returns the declination of the source in radians.""" return self._dec def _set_dec(self, dec): + """Sets the declination of the source in radians.""" self._dec = dec def __str__(self): diff --git a/skyllh/core/storage.py b/skyllh/core/storage.py index 1f90834a03..7d74d35f53 100644 --- a/skyllh/core/storage.py +++ b/skyllh/core/storage.py @@ -3,6 +3,8 @@ import os.path import pickle import sys +from collections.abc import Sequence +from typing import Any, cast, overload import numpy as np @@ -25,16 +27,16 @@ _FILE_LOADER_REG = {} -def register_FileLoader(formats, fileloader_cls): +def register_FileLoader(formats: str | Sequence[str], fileloader_cls: type['FileLoader']): """Registers the given file formats (file extensions) to the given FileLoader class. Parameters ---------- - formats : str | list of str + formats The list of file name extensions that should be mapped to the FileLoader class. - fileloader_cls : instance of FileLoader + fileloader_cls The subclass of FileLoader that should be used for the given file formats. """ @@ -51,14 +53,14 @@ def register_FileLoader(formats, fileloader_cls): _FILE_LOADER_REG[fmt] = fileloader_cls -def create_FileLoader(pathfilenames, **kwargs): +def create_FileLoader(pathfilenames: str | Sequence[str], **kwargs) -> 'FileLoader': """Creates the appropriate FileLoader object for the given file names. It looks up the FileLoader class from the FileLoader registry for the file name extension of the first file name in the given list. Parameters ---------- - pathfilenames : str | sequence of str + pathfilenames The sequence of fully qualified file names of the files that should be loaded. @@ -69,7 +71,7 @@ def create_FileLoader(pathfilenames, **kwargs): Returns ------- - fileloader : FileLoader + fileloader The appropriate FileLoader instance for the given type of data files. """ if isinstance(pathfilenames, str): @@ -100,13 +102,13 @@ def assert_file_exists(pathfilename): class FileLoader(metaclass=abc.ABCMeta): """Abstract base class for a FileLoader class.""" - def __init__(self, pathfilenames, **kwargs): + def __init__(self, pathfilenames: str | Sequence[str], **kwargs): """Creates a new FileLoader instance. Parameters ---------- - pathfilenames : str | sequence of str - The sequence of fully qualified file names of the data files that + pathfilenames + The sequence of fully qualified file name(s) of the data file(s) that need to be loaded. """ super().__init__(**kwargs) @@ -127,7 +129,7 @@ def pathfilename_list(self, pathfilenames): self._pathfilename_list = list(pathfilenames) @abc.abstractmethod - def load_data(self, **kwargs): + def load_data(self, **kwargs) -> Any: """This method is supposed to load the data from the file.""" @@ -138,32 +140,32 @@ class NPYFileLoader(FileLoader): several data files. """ - def __init__(self, pathfilenames, **kwargs): + def __init__(self, pathfilenames: str | Sequence[str], **kwargs): """Creates a new NPYFileLoader instance. Parameters ---------- - pathfilenames : str | sequence of str + pathfilenames The sequence of fully qualified file names of the data files that need to be loaded. """ super().__init__(pathfilenames=pathfilenames, **kwargs) def _load_file_memory_efficiently( - self, pathfilename, keep_fields, dtype_conversions, dtype_conversion_except_fields - ): + self, pathfilename: str, keep_fields: list[str] | None, dtype_conversions, dtype_conversion_except_fields + ) -> 'DataFieldRecordArray': """Loads a single file in a memory efficient way. Parameters ---------- - pathfilename : str + pathfilename The fully qualified file name of the to-be-loaded file. - keep_fields : list of str | None + keep_fields The list of field names which should be kept. Returns ------- - data : DataFieldRecordArray instance + data An instance of DataFieldRecordArray holding the data. """ assert_file_exists(pathfilename) @@ -238,25 +240,29 @@ def _load_file_time_efficiently(self, pathfilename, keep_fields, dtype_conversio return data - def load_data( - self, keep_fields=None, dtype_conversions=None, dtype_conversion_except_fields=None, efficiency_mode=None - ): + def load_data( # pyright: ignore[reportIncompatibleMethodOverride] + self, + keep_fields: str | Sequence[str] | None = None, + dtype_conversions: dict | None = None, + dtype_conversion_except_fields: str | Sequence[str] | None = None, + efficiency_mode: str | None = None, + ) -> 'DataFieldRecordArray': """Loads the data from the files specified through their fully qualified file names. Parameters ---------- - keep_fields : str | sequence of str | None + keep_fields Load the data into memory only for these data fields. If set to ``None``, all in-file-present data fields are loaded into memory. - dtype_conversions : dict | None + dtype_conversions If not None, this dictionary defines how data fields of specific data types get converted into the specified data types. This can be used to use less memory. - dtype_conversion_except_fields : str | sequence of str | None + dtype_conversion_except_fields The sequence of field names whose data type should not get converted. - efficiency_mode : str | None + efficiency_mode The efficiency mode the data should get loaded with. Possible values are: @@ -274,7 +280,7 @@ def load_data( Returns ------- - data : instance of DataFieldRecordArray + data The DataFieldRecordArray holding the loaded data. Raises @@ -341,12 +347,12 @@ class ParquetFileLoader(FileLoader): """ @tool.requires('pyarrow', 'pyarrow.parquet') - def __init__(self, pathfilenames, **kwargs): + def __init__(self, pathfilenames: str | Sequence[str], **kwargs): """Creates a new file loader instance for parquet data files. Parameters ---------- - pathfilenames : str | sequence of str + pathfilenames The sequence of fully qualified file names of the data files that need to be loaded. """ @@ -357,35 +363,35 @@ def __init__(self, pathfilenames, **kwargs): def load_data( self, - keep_fields=None, - dtype_conversions=None, - dtype_conversion_except_fields=None, - copy=False, + keep_fields: str | Sequence[str] | None = None, + dtype_conversions: dict | None = None, + dtype_conversion_except_fields: str | Sequence[str] | None = None, + copy: bool = False, **kwargs, - ): + ) -> 'DataFieldRecordArray': """Loads the data from the files specified through their fully qualified file names. Parameters ---------- - keep_fields : str | sequence of str | None + keep_fields Load the data into memory only for these data fields. If set to ``None``, all in-file-present data fields are loaded into memory. - dtype_conversions : dict | None + dtype_conversions If not ``None``, this dictionary defines how data fields of specific data types get converted into the specified data types. This can be used to use less memory. - dtype_conversion_except_fields : str | sequence of str | None + dtype_conversion_except_fields The sequence of field names whose data type should not get converted. - copy : bool + copy If set to ``True``, the column data from the pyarrow.Table instance will be copied into the DataFieldRecordArray. This should not be necessary. Returns ------- - data : instance of DataFieldRecordArray + data The DataFieldRecordArray holding the loaded data. """ assert_file_exists(self.pathfilename_list[0]) @@ -413,15 +419,15 @@ class PKLFileLoader(FileLoader): `pickle.load` function for loading the data from the file. """ - def __init__(self, pathfilenames, pkl_encoding=None, **kwargs): + def __init__(self, pathfilenames: str | Sequence[str], pkl_encoding: str | None = None, **kwargs): """Creates a new file loader instance for a pickled data file. Parameters ---------- - pathfilenames : str | sequence of str + pathfilenames The sequence of fully qualified file names of the data files that need to be loaded. - pkl_encoding : str | None + pkl_encoding The encoding of the pickled data files. If None, the default encodings 'ASCII' and 'latin1' will be tried to load the data. """ @@ -449,7 +455,7 @@ def load_data(self, **kwargs): Returns ------- - data : Python object | list of Python objects + data The de-pickled Python object. If more than one file was specified, this is a list of Python objects, i.e. one object for each file. The file <-> object mapping order is preserved. @@ -500,17 +506,23 @@ class TextFileLoader(FileLoader): first line of the text file for a table header. """ - def __init__(self, pathfilenames, header_comment='#', header_separator=None, **kwargs): + def __init__( + self, + pathfilenames: str | Sequence[str], + header_comment: str = '#', + header_separator: str | None = None, + **kwargs, + ): """Creates a new file loader instance for a text data file. Parameters ---------- - pathfilenames : str | sequence of str + pathfilenames The sequence of fully qualified file names of the data files that need to be loaded. - header_comment : str + header_comment The character that defines a comment line in the text file. - header_separator : str | None + header_separator The separator of the header field names. If None, it assumes whitespaces. """ @@ -543,18 +555,18 @@ def header_separator(self, s): raise TypeError('The header_separator property must be None or of type str!') self._header_separator = s - def _extract_column_names(self, line): + def _extract_column_names(self, line: str) -> list[str] | None: """Tries to extract the column names of the data table based on the given line. Parameters ---------- - line : str + line The text line containing the column names. Returns ------- - names : list of str | None + names The column names. It returns None, if the column names cannot be extracted. """ @@ -577,28 +589,34 @@ def _extract_column_names(self, line): return names - def _load_file(self, pathfilename, keep_fields, dtype_conversions, dtype_conversion_except_fields): + def _load_file( + self, + pathfilename: str, + keep_fields: str | Sequence[str] | None, + dtype_conversions: dict | None, + dtype_conversion_except_fields: str | Sequence[str] | None, + ) -> 'DataFieldRecordArray': """Loads the given file. Parameters ---------- - pathfilename : str + pathfilename The fully qualified file name of the data file that need to be loaded. - keep_fields : str | sequence of str | None + keep_fields Load the data into memory only for these data fields. If set to ``None``, all in-file-present data fields are loaded into memory. - dtype_conversions : dict | None + dtype_conversions If not None, this dictionary defines how data fields of specific data types get converted into the specified data types. This can be used to use less memory. - dtype_conversion_except_fields : str | sequence of str | None + dtype_conversion_except_fields The sequence of field names whose data type should not get converted. Returns ------- - data : DataFieldRecordArray instance + data The DataFieldRecordArray instance holding the loaded data. """ assert_file_exists(pathfilename) @@ -636,26 +654,32 @@ def _load_file(self, pathfilename, keep_fields, dtype_conversions, dtype_convers return data - def load_data(self, keep_fields=None, dtype_conversions=None, dtype_conversion_except_fields=None, **kwargs): + def load_data( + self, + keep_fields: str | Sequence[str] | None = None, + dtype_conversions: dict | None = None, + dtype_conversion_except_fields: str | Sequence[str] | None = None, + **kwargs, + ) -> 'DataFieldRecordArray': """Loads the data from the data files specified through their fully qualified file names. Parameters ---------- - keep_fields : str | sequence of str | None + keep_fields Load the data into memory only for these data fields. If set to ``None``, all in-file-present data fields are loaded into memory. - dtype_conversions : dict | None + dtype_conversions If not None, this dictionary defines how data fields of specific data types get converted into the specified data types. This can be used to use less memory. - dtype_conversion_except_fields : str | sequence of str | None + dtype_conversion_except_fields The sequence of field names whose data type should not get converted. Returns ------- - data : instance of DataFieldRecordArray + data The DataFieldRecordArray holding the loaded data. Raises @@ -708,72 +732,71 @@ def load_data(self, keep_fields=None, dtype_conversions=None, dtype_conversion_e return data -class DataTableAccessor( - metaclass=abc.ABCMeta, -): +class DataTableAccessor(metaclass=abc.ABCMeta): """This class provides an interface wrapper to access the data table of a particular format in a unified way. """ def __init__(self, **kwargs): + """Creates a new instance of DataTableAccessor.""" super().__init__(**kwargs) @abc.abstractmethod - def get_column(self, data, name): + def get_column(self, data: Any, name: str) -> np.ndarray: """This method is supposed to return a numpy.ndarray holding the data of the column with name ``name``. Parameters ---------- - data : any + data The data table. - name : str + name The name of the column. Returns ------- - arr : instance of numpy.ndarray + arr The column data as numpy ndarray. """ @abc.abstractmethod - def get_field_names(self, data): + def get_field_names(self, data) -> list: """This method is supposed to return a list of field names.""" @abc.abstractmethod - def get_field_name_to_dtype_dict(self, data): + def get_field_name_to_dtype_dict(self, data) -> dict: """This method is supposed to return a dictionary with field name and numpy dtype instance for each field. """ @abc.abstractmethod - def get_length(self, data): + def get_length(self, data) -> int: """This method is supposed to return the length of the data table.""" -class NDArrayDataTableAccessor( - DataTableAccessor, -): +class NDArrayDataTableAccessor(DataTableAccessor): """This class provides an interface wrapper to access the data table stored as a structured numpy ndarray. """ def __init__(self, **kwargs): + """Creates a new instance of NDArrayDataTableAccessor.""" super().__init__(**kwargs) - def get_column(self, data, name): + def get_column(self, data: np.ndarray, name: str): """Gets the column data from the structured ndarray. Parameters ---------- - data : instance of numpy.ndarray + data The structured numpy ndarray holding the table data. - name : str + name The name of the column. """ return data[name] def get_field_names(self, data): + """Returns the list of field names of the data table.""" return data.dtype.names def get_field_name_to_dtype_dict(self, data): @@ -789,29 +812,29 @@ def get_length(self, data): return length -class DictDataTableAccessor( - DataTableAccessor, -): +class DictDataTableAccessor(DataTableAccessor): """This class provides an interface wrapper to access the data table stored as a Python dictionary. """ def __init__(self, **kwargs): + """Creates a new instance of DictDataTableAccessor.""" super().__init__(**kwargs) - def get_column(self, data, name): + def get_column(self, data: dict, name: str): """Gets the column data from the dictionary. Parameters ---------- - data : dict + data The dictionary holding the table data. - name : str + name The name of the column. """ return data[name] def get_field_names(self, data): + """Returns the list of field names of the data table.""" return list(data.keys()) def get_field_name_to_dtype_dict(self, data): @@ -829,29 +852,29 @@ def get_length(self, data): return length -class ParquetDataTableAccessor( - DataTableAccessor, -): +class ParquetDataTableAccessor(DataTableAccessor): """This class provides an interface wrapper to access the data table stored as a Parquet table. """ def __init__(self, **kwargs): + """Creates a new instance of ParquetDataTableAccessor.""" super().__init__(**kwargs) - def get_column(self, data, name): + def get_column(self, data, name: str): """Gets the column data from the Parquet table. Parameters ---------- - data : instance of pyarrow.Table + data The instance of pyarrow.Table holding the table data. - name : str + name The name of the column. """ return data[name].to_numpy() def get_field_names(self, data): + """Returns the list of field names of the data table.""" return data.column_names def get_field_name_to_dtype_dict(self, data): @@ -866,25 +889,29 @@ def get_length(self, data): return len(data) -class DataFieldRecordArrayDataTableAccessor( - DataTableAccessor, -): +class DataFieldRecordArrayDataTableAccessor(DataTableAccessor): + """This class provides an accessor for table data stored as an instance of + :class:`~skyllh.core.storage.DataFieldRecordArray`. + """ + def __init__(self, **kwargs): + """Creates a new instance of DataFieldRecordArrayDataTableAccessor.""" super().__init__(**kwargs) - def get_column(self, data, name): + def get_column(self, data, name: str): """Gets the column data from the Parquet table. Parameters ---------- - data : instance of pyarrow.Table + data The instance of pyarrow.Table holding the table data. - name : str + name The name of the column. """ return data[name] def get_field_names(self, data): + """Returns the list of field names of the data table.""" return data.field_name_list def get_field_name_to_dtype_dict(self, data): @@ -908,18 +935,18 @@ class DataFieldRecordArray: def __init__( self, - data, - data_table_accessor=None, - keep_fields=None, - dtype_conversions=None, - dtype_conversion_except_fields=None, - copy=True, + data: Any, + data_table_accessor: 'DataTableAccessor | None' = None, + keep_fields: str | Sequence[str] | None = None, + dtype_conversions: dict | None = None, + dtype_conversion_except_fields: str | Sequence[str] | None = None, + copy: bool = True, ): """Creates a DataFieldRecordArray from the given data. Parameters ---------- - data : any | None + data The tabulated data in any format. The only requirement is that there is a DataTableAccessor instance available for the given data format. Supported data types are: @@ -936,21 +963,21 @@ def __init__( If set to `None`, the DataFieldRecordArray instance is initialized with no data and the length of the array is set to 0. - data_table_accessor : instance of DataTableAccessor | None + data_table_accessor The instance of DataTableAccessor which provides column access to ``data``. If set to ``None``, an appropriate ``DataTableAccessor`` instance will be selected based on the type of ``data``. - keep_fields : str | sequence of str | None + keep_fields If not None (default), this specifies the data fields that should get kept from the given data. Otherwise all data fields get kept. - dtype_conversions : dict | None + dtype_conversions If not None, this dictionary defines how data fields of specific data types get converted into the specified data types. This can be used to use less memory. - dtype_conversion_except_fields : str | sequence of str | None + dtype_conversion_except_fields The sequence of field names whose data type should not get converted. - copy : bool + copy Flag if the input data should get copied. Default is True. If a DataFieldRecordArray instance is provided, this option is set to ``True`` automatically. @@ -995,9 +1022,10 @@ def __init__( else: raise TypeError(f'No TableDataAccessor instance has been specified for the data of type {type(data)}!') - field_names = data_table_accessor.get_field_names(data) - fname2dtype = data_table_accessor.get_field_name_to_dtype_dict(data) - length = data_table_accessor.get_length(data) + _dta = cast(DataTableAccessor, data_table_accessor) + field_names = _dta.get_field_names(data) + fname2dtype = _dta.get_field_name_to_dtype_dict(data) + length = _dta.get_length(data) for fname in field_names: # Ignore fields that should not get kept. @@ -1016,9 +1044,9 @@ def __init__( # Create a ndarray with the final data type and then assign the # values from the data, which technically is a copy. field_arr = np.empty((length,), dtype=dt) - np.copyto(field_arr, data_table_accessor.get_column(data, fname)) + np.copyto(field_arr, _dta.get_column(data, fname)) else: - field_arr = data_table_accessor.get_column(data, fname) + field_arr = _dta.get_column(data, fname) if self._len is None: self._len = len(field_arr) @@ -1039,29 +1067,33 @@ def __init__( self._field_name_list = list(self._data_fields.keys()) self._indices = None - def __contains__(self, name): + def __contains__(self, name: str) -> bool: """Checks if the given field exists in this DataFieldRecordArray instance. Parameters ---------- - name : str + name The name of the field. Returns ------- - check : bool + check True, if the given field exists in this DataFieldRecordArray instance, False otherwise. """ return name in self._data_fields - def __getitem__(self, name): + @overload + def __getitem__(self, name: str) -> np.ndarray: ... + @overload + def __getitem__(self, name: np.ndarray) -> 'DataFieldRecordArray': ... + def __getitem__(self, name: str | np.ndarray) -> 'np.ndarray | DataFieldRecordArray': """Implements data field value access. Parameters ---------- - name : str | numpy ndarray of int or bool + name The name of the data field. If a numpy ndarray is given, it must contain the indices for which to retrieve a data selection of the entire DataFieldRecordArray. A numpy ndarray of bools can be given @@ -1074,7 +1106,7 @@ def __getitem__(self, name): Returns ------- - data : numpy ndarray | instance of DataFieldRecordArray + data The requested field data or a DataFieldRecordArray holding the requested selection of the entire data. """ @@ -1086,17 +1118,21 @@ def __getitem__(self, name): return self._data_fields[name] - def __setitem__(self, name, arr): + @overload + def __setitem__(self, name: str, arr: np.ndarray) -> None: ... + @overload + def __setitem__(self, name: np.ndarray, arr: 'DataFieldRecordArray') -> None: ... + def __setitem__(self, name: str | np.ndarray, arr: 'np.ndarray | DataFieldRecordArray') -> None: """Implements data field value assignment. If values are assigned to a data field that does not exist yet, it will be added via the ``append_field`` method. Parameters ---------- - name : str | numpy ndarray of int or bool + name The name of the data field, or a numpy ndarray holding the indices or mask of a selection of this DataFieldRecordArray. - arr : numpy ndarray | instance of DataFieldRecordArray + arr The numpy ndarray holding the field values. It must be of the same length as this DataFieldRecordArray. If `name` is a numpy ndarray, `arr` must be a DataFieldRecordArray. @@ -1108,10 +1144,12 @@ def __setitem__(self, name, arr): DataFieldRecordArray instance. """ if isinstance(name, np.ndarray): + assert isinstance(arr, DataFieldRecordArray) self.set_selection(name, arr) return # Check if a new field is supposed to be added. + assert isinstance(arr, np.ndarray) if name not in self: self.append_field(name, arr) return @@ -1124,21 +1162,18 @@ def __setitem__(self, name, arr): 'instance!' ) - if not isinstance(arr, np.ndarray): - raise TypeError('When setting a field directly, the data must be provided as a numpy ndarray!') - self._data_fields[name] = arr - def __len__(self): - return self._len + def __len__(self) -> int: + return self._len or 0 - def __sizeof__(self): + def __sizeof__(self) -> int: """Calculates the size in bytes of this DataFieldRecordArray instance in memory. Returns ------- - memsize : int + memsize The memory size in bytes that this DataFieldRecordArray instance has. """ @@ -1155,6 +1190,9 @@ def __str__(self): # Generates a pretty string representation of the given field name. def _pretty_str_field(name): + """Creates a pretty string representation of the data field with the + given name. + """ field = self._data_fields[name] s = ( f'{name.ljust(max_field_name_len)}: ' @@ -1192,16 +1230,17 @@ def indices(self): DataFieldRecordArray. """ if self._indices is None: - self._indices = np.arange(self._len) + _len = int(self._len) if self._len is not None else 0 + self._indices = np.arange(_len) return self._indices - def append(self, arr): + def append(self, arr: 'DataFieldRecordArray'): """Appends the given DataFieldRecordArray to this DataFieldRecordArray instance. Parameters ---------- - arr : instance of DataFieldRecordArray + arr The instance of DataFieldRecordArray that should get appended to this DataFieldRecordArray. It must contain the same data fields. Additional data fields are ignored. @@ -1212,17 +1251,17 @@ def append(self, arr): for fname in self._field_name_list: self._data_fields[fname] = np.append(self._data_fields[fname], arr[fname]) - self._len += len(arr) + self._len = int(self._len or 0) + len(arr) self._indices = None - def append_field(self, name, data): + def append_field(self, name: str, data: np.ndarray): """Appends a field and its data to this DataFieldRecordArray instance. Parameters ---------- - name : str + name The name of the new data field. - data : numpy ndarray + data The numpy ndarray holding the data. The length of the ndarray must match the current length of this DataFieldRecordArray instance. @@ -1249,13 +1288,13 @@ def append_field(self, name, data): self._data_fields[name] = data self._field_name_list.append(name) - def as_numpy_record_array(self): + def as_numpy_record_array(self) -> np.ndarray: """Creates a numpy record ndarray instance holding the data of this DataFieldRecordArray instance. Returns ------- - arr : instance of numpy record ndarray + arr The numpy recarray ndarray holding the data of this DataFieldRecordArray instance. """ @@ -1267,25 +1306,25 @@ def as_numpy_record_array(self): return arr - def copy(self, keep_fields=None): + def copy(self, keep_fields: str | Sequence[str] | None = None): """Creates a new DataFieldRecordArray that is a copy of this DataFieldRecordArray instance. Parameters ---------- - keep_fields : str | sequence of str | None + keep_fields If not None (default), this specifies the data fields that should get kept from this DataFieldRecordArray. Otherwise all data fields get kept. """ return DataFieldRecordArray(self, keep_fields=keep_fields) - def remove_field(self, name): + def remove_field(self, name: str): """Removes the given field from this array. Parameters ---------- - name : str + name The name of the data field that is to be removed. """ self._data_fields.pop(name) @@ -1295,14 +1334,14 @@ def get_field_dtype(self, name): """Returns the numpy dtype object of the given data field.""" return self._data_fields[name].dtype - def set_field_dtype(self, name, dt): + def set_field_dtype(self, name: str, dt: np.ndarray): """Sets the data type of the given field. Parameters ---------- - name : str + name The name of the data field. - dt : numpy.dtype + dt The dtype instance defining the new data type. """ if name not in self: @@ -1312,15 +1351,15 @@ def set_field_dtype(self, name, dt): self._data_fields[name] = self._data_fields[name].astype(dt, copy=False) - def convert_dtypes(self, conversions, except_fields=None): + def convert_dtypes(self, conversions: dict, except_fields: Sequence[str] | None = None): """Converts the data type of the data fields of this DataFieldRecordArray. This method can be used to compress the data. Parameters ---------- - conversions : dict of `old_dtype` -> `new_dtype` + conversions The dictionary with the old dtype as key and the new dtype as value. - except_fields : sequence of str | None + except_fields The sequence of field names, which should not get converted. """ if not isinstance(conversions, dict): @@ -1340,18 +1379,18 @@ def convert_dtypes(self, conversions, except_fields=None): new_dtype = conversions[old_dtype] _data_fields[fname] = _data_fields[fname].astype(new_dtype) - def get_selection(self, indices): + def get_selection(self, indices: np.ndarray) -> 'DataFieldRecordArray': """Creates an DataFieldRecordArray that contains a selection of the data of this DataFieldRecordArray instance. Parameters ---------- - indices : (N,)-shaped numpy ndarray of int or bool + indices The numpy ndarray holding the indices for which to select the data. Returns ------- - data_field_array : instance of DataFieldRecordArray + data_field_array The DataFieldRecordArray that contains the selection of the original DataFieldRecordArray. The selection data is a copy of the original data. @@ -1363,16 +1402,16 @@ def get_selection(self, indices): data[fname] = self._data_fields[fname][indices] return DataFieldRecordArray(data, copy=False) - def set_selection(self, indices, arr): + def set_selection(self, indices: np.ndarray, arr: 'DataFieldRecordArray'): """Sets a selection of the data of this DataFieldRecordArray instance to the data given in arr. Parameters ---------- - indices : (N,)-shaped numpy ndarray of int or bool + indices The numpy ndarray holding the indices or mask for which to set the data. - arr : instance of DataFieldRecordArray + arr The instance of DataFieldRecordArray holding the selection data. It must have the same fields defined as this DataFieldRecordArray instance. @@ -1383,14 +1422,14 @@ def set_selection(self, indices, arr): for fname in self._field_name_list: self._data_fields[fname][indices] = arr[fname] - def rename_fields(self, conversions, must_exist=False): + def rename_fields(self, conversions: dict, must_exist: bool = False): """Renames the given fields of this array. Parameters ---------- - conversions : dict of `old_name` -> `new_name` + conversions The dictionary holding the old and new names of the data fields. - must_exist : bool + must_exist Flag if the given fields must exist. If set to ``True`` and a field does not exist, a KeyError is raised. @@ -1408,13 +1447,13 @@ def rename_fields(self, conversions, must_exist=False): self._field_name_list = list(self._data_fields.keys()) - def tidy_up(self, keep_fields): + def tidy_up(self, keep_fields: str | Sequence[str]): """Removes all fields that are not specified through the keep_fields argument. Parameters ---------- - keep_fields : str | sequence of str + keep_fields The field name(s), that should not be removed. Raises @@ -1435,17 +1474,17 @@ def tidy_up(self, keep_fields): if fname not in keep_fields: self.remove_field(fname) - def sort_by_field(self, name): + def sort_by_field(self, name: str) -> np.ndarray: """Sorts the data along the given field name in ascending order. Parameters ---------- - name : str + name The name of the field along the events should get sorted. Returns ------- - sorted_idxs : (n_events,)-shaped numpy ndarray + sorted_idxs The numpy ndarray holding the indices of the sorted array. Raises diff --git a/skyllh/core/test_statistic.py b/skyllh/core/test_statistic.py index beea5aa4ae..8ea894a1f4 100644 --- a/skyllh/core/test_statistic.py +++ b/skyllh/core/test_statistic.py @@ -6,6 +6,10 @@ import numpy as np +from skyllh.core.llhratio import TCLLHRatio +from skyllh.core.parameters import ParameterModelMapper +from skyllh.core.timing import TimeLord + class TestStatistic(metaclass=abc.ABCMeta): """This is the abstract base class for a test statistic class.""" @@ -15,25 +19,25 @@ def __init__(self, **kwargs): super().__init__(**kwargs) @abc.abstractmethod - def __call__(self, pmm, log_lambda, fitparam_values, **kwargs): + def __call__(self, pmm: ParameterModelMapper, log_lambda: float, fitparam_values: np.ndarray, **kwargs) -> float: """This method is supposed to evaluate the test-statistic function. Parameters ---------- - pmm : instance of ParameterModelMapper + pmm The ParameterModelMapper instance that defines the global parameter set. - log_lambda : float + log_lambda The value of the log-likelihood ratio function. Usually, this is its maximum. - fitparam_values : instance of numpy ndarray + fitparam_values The (N_fitparam,)-shaped 1D numpy ndarray holding the global fit parameter values of the log-likelihood ratio function for the given log_lambda value. Returns ------- - TS : float + TS The calculated test-statistic value. """ @@ -50,12 +54,12 @@ class WilksTestStatistic(TestStatistic): :math:`\hat{n}_{\text{s}} < 0`, and positive otherwise. """ - def __init__(self, ns_param_name='ns', **kwargs): + def __init__(self, ns_param_name: str = 'ns', **kwargs): """Constructs the test-statistic function instance. Parameters ---------- - ns_param_name : str + ns_param_name The name of the global fit parameter for the number of signal events in the detector, ns. """ @@ -70,25 +74,25 @@ def ns_param_name(self): """ return self._ns_param_name - def __call__(self, pmm, log_lambda, fitparam_values, **kwargs): + def __call__(self, pmm: ParameterModelMapper, log_lambda: float, fitparam_values: np.ndarray, **kwargs) -> float: """Evaluates the test-statistic function. Parameters ---------- - pmm : instance of ParameterModelMapper + pmm The ParameterModelMapper instance that defines the global parameter set. - log_lambda : float + log_lambda The value of the log-likelihood ratio function. Usually, this is its maximum. - fitparam_values : instance of numpy ndarray + fitparam_values The (N_fitparam,)-shaped 1D numpy ndarray holding the global fit parameter values of the log-likelihood ratio function for the given log_lambda value. Returns ------- - TS : float + TS The calculated test-statistic value. """ ns_pidx = pmm.get_gflp_idx(name=self._ns_param_name) @@ -101,7 +105,7 @@ def __call__(self, pmm, log_lambda, fitparam_values, **kwargs): TS = 2 * sgn_ns * log_lambda - return TS + return float(TS) class LLHRatioZeroNsTaylorWilksTestStatistic(TestStatistic): @@ -139,12 +143,12 @@ class LLHRatioZeroNsTaylorWilksTestStatistic(TestStatistic): being its second derivative w.r.t. ns. """ - def __init__(self, ns_param_name='ns', **kwargs): + def __init__(self, ns_param_name: str = 'ns', **kwargs): """Constructs the test-statistic function instance. Parameters ---------- - ns_param_name : str + ns_param_name The name of the global fit parameter for the number of signal events in the detector, ns. """ @@ -159,34 +163,43 @@ def ns_param_name(self): """ return self._ns_param_name - def __call__(self, pmm, log_lambda, fitparam_values, llhratio, grads, tl=None, **kwargs): + def __call__( # pyright: ignore[reportIncompatibleMethodOverride] + self, + pmm: ParameterModelMapper, + log_lambda: float, + fitparam_values: np.ndarray, + llhratio: TCLLHRatio, + grads: np.ndarray, + tl: TimeLord | None = None, + **kwargs, + ) -> float: """Evaluates the test-statistic function. Parameters ---------- - pmm : instance of ParameterModelMapper + pmm The ParameterModelMapper instance that defines the global parameter set. - log_lambda : float + log_lambda The value of the log-likelihood ratio function. Usually, this is its maximum. - fitparam_values : instance of numpy ndarray + fitparam_values The (N_fitparam,)-shaped 1D numpy ndarray holding the global fit parameter values of the log-likelihood ratio function for the given log_lambda value. - llhratio : instance of LLHRatio + llhratio The log-likelihood ratio function, which should be used for the test-statistic function. - grads : instance of numpy ndarray + grads The (N_fitparam,)-shaped 1D numpy ndarray holding the values of the first derivative of the log-likelihood ratio function w.r.t. each global fit parameter. - tl : instance of TimeLord | None + tl The optional instance of TimeLord to measure timing information. Returns ------- - TS : float + TS The calculated test-statistic value. """ ns_pidx = pmm.get_gflp_idx(name=self._ns_param_name) @@ -195,7 +208,7 @@ def __call__(self, pmm, log_lambda, fitparam_values, llhratio, grads, tl=None, * if ns == 0: nsgrad = grads[ns_pidx] - nsgrad2 = llhratio.calculate_ns_grad2(fitparam_values=fitparam_values, ns_pidx=ns_pidx, tl=tl) + nsgrad2 = llhratio.calculate_ns_grad2(ns=float(ns), src_params_recarray=np.empty(0), ns_pidx=ns_pidx, tl=tl) TS = -2 * nsgrad**2 / (4 * nsgrad2) diff --git a/skyllh/core/times.py b/skyllh/core/times.py index b9badec0c0..4137f77231 100644 --- a/skyllh/core/times.py +++ b/skyllh/core/times.py @@ -1,6 +1,9 @@ import abc +import numpy as np + from skyllh.core.livetime import Livetime +from skyllh.core.random import RandomStateService class TimeGenerationMethod( @@ -9,28 +12,30 @@ class TimeGenerationMethod( """Base class (type) for implementing a method to generate times.""" def __init__(self, **kwargs): + """Creates a new instance of TimeGenerationMethod.""" super().__init__(**kwargs) @abc.abstractmethod def generate_times( self, - rss, - size, - ): + rss: RandomStateService, + size: int, + **kwargs, + ) -> np.ndarray: """The ``generate_times`` method implements the actual generation of times, which is method dependent. Parameters ---------- - rss : instance of RandomStateService + rss The random state service providing the random number generator (RNG). - size : int + size The number of times that should get generated. Returns ------- - times : ndarray + times The 1d numpy ndarray holding the generated times. """ @@ -44,12 +49,12 @@ class LivetimeTimeGenerationMethod( method of the Livetime class. """ - def __init__(self, livetime, **kwargs): + def __init__(self, livetime: Livetime, **kwargs): """Creates a new LivetimeTimeGeneration instance. Parameters ---------- - livetime : Livetime + livetime The Livetime instance that should be used to generate times from. """ super().__init__(**kwargs) @@ -69,8 +74,8 @@ def livetime(self, livetime): def generate_times( self, - rss, - size, + rss: RandomStateService, + size: int, **kwargs, ): """Generates `size` MJD times according to the detector on-times @@ -78,15 +83,15 @@ def generate_times( Parameters ---------- - rss : instance of RandomStateService + rss The random state service providing the random number generator (RNG). - size : int + size The number of times that should get generated. Returns ------- - times : ndarray + times The 1d (`size`,)-shaped numpy ndarray holding the generated times. """ times = self._livetime.draw_ontimes(rss=rss, size=size, **kwargs) @@ -95,13 +100,17 @@ def generate_times( class TimeGenerator: - def __init__(self, method): + """This class provides a time generator that generates times according to a + defined time generation method. + """ + + def __init__(self, method: 'TimeGenerationMethod'): """Creates a time generator instance with a given defined time generation method. Parameters ---------- - method : instance of TimeGenerationMethod + method The instance of TimeGenerationMethod that defines the method of generating times. """ @@ -122,8 +131,8 @@ def method(self, method): def generate_times( self, - rss, - size, + rss: RandomStateService, + size: int, **kwargs, ): """Generates ``size`` amount of times by calling the ``generate_times`` @@ -131,10 +140,10 @@ def generate_times( Parameters ---------- - rss : instance of RandomStateService + rss The random state service providing the random number generator (RNG). - size : int + size The number of time that should get generated. **kwargs Additional keyword arguments are passed to the ``generate_times`` @@ -142,7 +151,7 @@ def generate_times( Returns ------- - times : ndarray + times The 1d (``size``,)-shaped ndarray holding the generated times. """ times = self._method.generate_times(rss=rss, size=size, **kwargs) diff --git a/skyllh/core/timing.py b/skyllh/core/timing.py index 82c2ad8fc8..f8fe6675ec 100644 --- a/skyllh/core/timing.py +++ b/skyllh/core/timing.py @@ -17,16 +17,20 @@ class TaskRecord: - def __init__(self, name, start_times, end_times): + """This class provides a record of a named task, holding the start and end + times of one or more executions of that task. + """ + + def __init__(self, name: str, start_times: list[float], end_times: list[float]): """Creates a new TaskRecord instance. Parameters ---------- - name : str + name The name of the task. - start_times : list of float + start_times The start times of the task in seconds. - end_times : list of float + end_times The end times of the task in seconds. """ self.name = name @@ -77,12 +81,12 @@ def niter(self): """(read-only) The number of times this task was executed.""" return len(self._start_times) - def join(self, tr): + def join(self, tr: 'TaskRecord'): """Joins this TaskRecord with the given TaskRecord instance. Parameters ---------- - tr : instance of TaskRecord + tr The instance of TaskRecord that should be joined with this TaskRecord instance. """ @@ -91,7 +95,14 @@ def join(self, tr): class TimeLord: + """This class provides a manager for keeping track of the execution times of + named tasks via :class:`TaskRecord` instances. + """ + def __init__(self): + """Creates a new instance of TimeLord with an empty list of task + records. + """ self._task_records = [] self._task_records_name_idx_map = {} @@ -113,44 +124,44 @@ def add_task_record(self, tr): self._task_records.append(tr) self._task_records_name_idx_map[tr.name] = len(self._task_records) - 1 - def get_task_record(self, name): + def get_task_record(self, name: str) -> 'TaskRecord': """Retrieves a task record of the given name. Parameters ---------- - name : str + name The name of the task record. Returns ------- - task_record : instance of TaskRecord + task_record The instance of TaskRecord with the requested name. """ return self._task_records[self._task_records_name_idx_map[name]] - def has_task_record(self, name): + def has_task_record(self, name: str) -> bool: """Checks if this TimeLord instance has a task record of the given name. Parameters ---------- - name : str + name The name of the task record. Returns ------- - check : bool + check ``True`` if this TimeLord instance has a task record of the given name, and ``False`` otherwise. """ return name in self._task_records_name_idx_map - def join(self, tl): + def join(self, tl: 'TimeLord'): """Joins a given TimeLord instance with this TimeLord instance. Tasks of the same name will be updated and new tasks will be added. Parameters ---------- - tl : instance of TimeLord + tl The instance of TimeLord whos tasks should be joined with the tasks of this TimeLord instance. """ @@ -198,20 +209,25 @@ def __str__(self): class TaskTimer: - def __init__(self, time_lord, name): - """ + """This class provides a context manager for timing the execution of a task + and recording it with a :class:`TimeLord` instance. + """ + + def __init__(self, time_lord: 'TimeLord | None', name: str): + """Creates a new TaskTimer instance. + Parameters ---------- - time_lord : instance of TimeLord + time_lord The TimeLord instance that keeps track of the recorded tasks. - name : str + name The name of the task. """ self.time_lord = time_lord self.name = name - self._start = None - self._end = None + self._start: float | None = None + self._end: float | None = None @property def time_lord(self): @@ -238,8 +254,11 @@ def name(self, name): self._name = name @property - def duration(self): + def duration(self) -> float: """The duration in seconds the task was executed.""" + assert self._end is not None and self._start is not None, ( + 'TaskTimer must be used as a context manager before accessing duration' + ) return self._end - self._start def __enter__(self): @@ -254,4 +273,5 @@ def __exit__(self, exc_type, exc_value, traceback): if self._time_lord is None: return + assert self._start is not None self._time_lord.add_task_record(TaskRecord(name=self._name, start_times=[self._start], end_times=[self._end])) diff --git a/skyllh/core/tool.py b/skyllh/core/tool.py index da7cc16c18..53053b7f1f 100644 --- a/skyllh/core/tool.py +++ b/skyllh/core/tool.py @@ -6,6 +6,7 @@ import importlib import importlib.util import sys +from types import ModuleType import numpy as np @@ -16,17 +17,17 @@ def assert_tool_version( - tool, - version, + tool: str, + version: str, ): """Asserts the required version of the tool. The tool module must have the attribute ``__version__``. Parameters ---------- - tool : str + tool The name of the tool. - version : str + version The required version of the tool in the format ``"X.Y.Z"``, where ```` is one of ``<=``, ``==``, and ``>=``. @@ -73,17 +74,17 @@ def assert_tool_version( raise ValueError(f'The version comparison operator "{comp_op}" for the tool "{tool}" is not supported!') -def is_available(name): +def is_available(name: str) -> bool: """Checks if the given Python package is available for import. Parameters ---------- - name : str + name The name of the Python package. Returns ------- - check : bool + check ``True`` if the given Python package is available, ``False`` otherwise. Raises @@ -99,18 +100,18 @@ def is_available(name): return spec is not None -def get(name): +def get(name: str) -> ModuleType: """Returns the module object of the given tool. This will import the Python package if it was not yet imported. Parameters ---------- - name : str + name The name of the Python package. Returns ------- - module : Python module + module The (imported) Python module object. """ if name in sys.modules: @@ -121,21 +122,21 @@ def get(name): def _get_tool_and_version( - tool, -): + tool: str | tuple[str, str], +) -> tuple[str, str | None]: """Returns the tool and and version based on the input value for the tool. Parameters ---------- - tool : str | (str, str) + tool Either the tool name or the tuple with the tool name and required version string. Returns ------- - tool : str + tool The name of the tool. - version : str | None + version The tool's version string, or ``None``, if no version was specified. """ if not (isinstance(tool, (str, tuple))): @@ -189,7 +190,12 @@ def requires(*tools): """ def decorator(f): + """Wraps the decorated function ``f`` with the tool availability check.""" + def wrapper(*args, **kwargs): + """Checks the availability of all required tools and then calls the + decorated function ``f``. + """ for tool in tools: (tool, version) = _get_tool_and_version(tool) if not is_available(tool): diff --git a/skyllh/core/trialdata.py b/skyllh/core/trialdata.py index 580d23168b..5f665110eb 100644 --- a/skyllh/core/trialdata.py +++ b/skyllh/core/trialdata.py @@ -6,22 +6,27 @@ """ from collections import OrderedDict +from collections.abc import Callable, Sequence import numpy as np from skyllh.core import display as dsp +from skyllh.core.event_selection import EventSelectionMethod from skyllh.core.logging import ( get_logger, ) +from skyllh.core.parameters import ParameterModelMapper from skyllh.core.py import ( classname, func_has_n_args, int_cast, issequenceof, ) +from skyllh.core.source_hypo_grouping import SourceHypoGroupManager from skyllh.core.storage import ( DataFieldRecordArray, ) +from skyllh.core.timing import TimeLord logger = get_logger(__name__) @@ -33,17 +38,24 @@ class DataField: """ def __init__( - self, name, func, global_fitparam_names=None, dt=None, is_src_field=False, is_srcevt_data=False, **kwargs + self, + name: str, + func: Callable, + global_fitparam_names: str | Sequence[str] | None = None, + dt: np.dtype | str | None = None, + is_src_field: bool = False, + is_srcevt_data: bool = False, + **kwargs, ): """Creates a new instance of DataField that might depend on fit parameters. Parameters ---------- - name : str + name The name of the data field. It serves as the identifier for the data field. - func : callable + func The function that calculates the values of this data field. The call signature must be @@ -54,20 +66,20 @@ def __init__( ``pmm`` is the instance of ParameterModelMapper, and ``global_fitparams_dict`` is the dictionary with the current global fit parameter names and values. - global_fitparam_names : str | sequence of str | None + global_fitparam_names The sequence of str instances specifying the names of the global fit parameters this data field depends on. If set to None, the data field does not depend on any fit parameters. - dt : numpy dtype | str | None + dt If specified it defines the data type this data field should have. If a str instance is given, it defines the name of the data field whose data type should be taken for this data field. - is_src_field : bool + is_src_field Flag if this data field is a source data field (``True``) and values should be stored within this DataField instance, instead of the events DataFieldRecordArray instance of the TrialDataManager (``False``). - is_srcevt_data : bool + is_srcevt_data Flag if the data field will hold source-event data, i.e. data of length N_values. In that case the data cannot be stored within the events attribute of the TrialDataManager, but must be stored in the @@ -100,7 +112,7 @@ def __init__( # Define the member variable that holds the numpy ndarray with the data # field values. - self._values = None + self._values: np.ndarray | None = None # Define the most efficient `calculate` method for this kind of data # field. @@ -204,48 +216,49 @@ def _convert_to_desired_dtype(self, tdm, values): values = values.astype(dt, copy=False) return values - def _calc_source_values(self, tdm, shg_mgr, pmm): + def _calc_source_values(self, tdm: 'TrialDataManager', shg_mgr: SourceHypoGroupManager, pmm: ParameterModelMapper): """Calculates the data field values utilizing the defined external function. The data field values solely depend on fixed source parameters. Parameters ---------- - tdm : instance of TrialDataManager + tdm The TrialDataManager instance this data field is part of and is holding the event data. - shg_mgr : instance of SourceHypoGroupManager + shg_mgr The instance of SourceHypoGroupManager, which defines the source hypothesis groups. - pmm : instance of ParameterModelMapper + pmm The instance of ParameterModelMapper, which defines the global parameters and their mapping to local source parameters. """ - self._values = self._func(tdm=tdm, shg_mgr=shg_mgr, pmm=pmm) + _func_result = self._func(tdm=tdm, shg_mgr=shg_mgr, pmm=pmm) - if not isinstance(self._values, np.ndarray): + if not isinstance(_func_result, np.ndarray): raise TypeError( f'The calculation function for the data field "{self._name}" ' 'must return an instance of numpy.ndarray! ' - f'Currently it is of type "{classname(self._values)}".' + f'Currently it is of type "{classname(_func_result)}".' ) + self._values = _func_result # Convert the data type. self._values = self._convert_to_desired_dtype(tdm, self._values) - def _calc_static_values(self, tdm, shg_mgr, pmm): + def _calc_static_values(self, tdm: 'TrialDataManager', shg_mgr: SourceHypoGroupManager, pmm: ParameterModelMapper): """Calculates the data field values utilizing the defined external function, that are static and only depend on source parameters. Parameters ---------- - tdm : instance of TrialDataManager + tdm The TrialDataManager instance this data field is part of and is holding the event data. - shg_mgr : instance of SourceHypoGroupManager + shg_mgr The instance of SourceHypoGroupManager, which defines the source hypothesis groups. - pmm : instance of ParameterModelMapper + pmm The instance of ParameterModelMapper, which defines the global parameters and their mapping to local source parameters. """ @@ -273,31 +286,39 @@ def _calc_static_values(self, tdm, shg_mgr, pmm): else: # Set the data values. This will add the data field to the # DataFieldRecordArray if it does not exist yet. + assert tdm.events is not None tdm.events[self._name] = values - def _calc_global_fitparam_dependent_values(self, tdm, shg_mgr, pmm, global_fitparams_dict): + def _calc_global_fitparam_dependent_values( + self, + tdm: 'TrialDataManager', + shg_mgr: SourceHypoGroupManager, + pmm: ParameterModelMapper, + global_fitparams_dict: dict, + ): """Calculate data field values utilizing the defined external function, that depend on fit parameter values. We check if the fit parameter values have changed. Parameters ---------- - tdm : instance of TrialDataManager + tdm The TrialDataManager instance this data field is part of and is holding the event data. - shg_mgr : instance of SourceHypoGroupManager + shg_mgr The instance of SourceHypoGroupManager, which defines the source hypothesis groups. - pmm : instance of ParameterModelMapper + pmm The instance of ParameterModelMapper defining the mapping of the global parameters to local source parameters. - global_fitparams_dict : dict + global_fitparams_dict The dictionary holding the current global fit parameter names and values. """ # Determine if we need to calculate the values. calc_values = False + assert tdm.events is not None if self._name not in tdm.events: calc_values = True else: @@ -350,12 +371,12 @@ class TrialDataManager: Hence, data fields are calculated only once. """ - def __init__(self, index_field_name=None, **kwargs): + def __init__(self, index_field_name: str | None = None, **kwargs): """Creates a new TrialDataManager instance. Parameters ---------- - index_field_name : str | None + index_field_name The name of the field that should be used as primary index field. If provided, the events will be sorted along this data field. This might be useful for run-time performance. @@ -461,6 +482,7 @@ def n_events(self, n): @property def n_selected_events(self): """(read-only) The number of selected events which should get evaluated.""" + assert self._events is not None return len(self._events) @property @@ -469,6 +491,8 @@ def n_pure_bkg_events(self): of the trial data, but must be considered for the test-statistic value. It is the difference of n_events and n_selected_events. """ + assert self._events is not None + assert self._n_events is not None return self._n_events - len(self._events) @property @@ -486,17 +510,17 @@ def trial_data_state_id(self): """ return self._trial_data_state_id - def __contains__(self, name): + def __contains__(self, name: str) -> bool: """Checks if the given data field is defined in this data field manager. Parameters ---------- - name : str + name The name of the data field. Returns ------- - check : bool + check True if the data field is defined in this data field manager, False otherwise. """ @@ -562,19 +586,19 @@ def __str__(self): return s - def broadcast_sources_array_to_values_array(self, arr): + def broadcast_sources_array_to_values_array(self, arr: np.ndarray) -> np.ndarray: """Broadcasts the given 1d numpy ndarray of length 1 or N_sources to a numpy ndarray of length N_values. Parameters ---------- - arr : instance of ndarray + arr The (N_sources,)- or (1,)-shaped numpy ndarray holding values for each source. Returns ------- - out_arr : instance of ndarray + out_arr The (N_values,)-shaped numpy ndarray holding the source values broadcasted to each event value. """ @@ -591,6 +615,7 @@ def broadcast_sources_array_to_values_array(self, arr): out_arr = np.empty((n_values,), dtype=arr.dtype) + assert self.src_evt_idxs is not None src_idxs = self.src_evt_idxs[0] v_start = 0 for src_idx, src_value in enumerate(arr): @@ -601,18 +626,18 @@ def broadcast_sources_array_to_values_array(self, arr): return out_arr - def broadcast_sources_arrays_to_values_arrays(self, arrays): + def broadcast_sources_arrays_to_values_arrays(self, arrays: Sequence[np.ndarray]) -> list[np.ndarray]: """Broadcasts the 1d numpy ndarrays to the values array. Parameters ---------- - arrays : sequence of numpy 1d ndarrays + arrays The sequence of (N_sources,)-shaped numpy ndarrays holding the parameter values. Returns ------- - out_arrays : list of numpy 1d ndarrays + out_arrays The list of (N_values,)-shaped numpy ndarrays holding the broadcasted array values. """ @@ -620,27 +645,28 @@ def broadcast_sources_arrays_to_values_arrays(self, arrays): return out_arrays - def broadcast_selected_events_arrays_to_values_arrays(self, arrays): + def broadcast_selected_events_arrays_to_values_arrays(self, arrays: Sequence[np.ndarray]) -> list[np.ndarray]: """Broadcasts the given arrays of length N_selected_events to arrays of length N_values. Parameters ---------- - arrays : sequence of instance of ndarray + arrays The sequence of instance of ndarray with the arrays to be broadcasted. Returns ------- - out_arrays : list of instance of ndarray + out_arrays The list of broadcasted numpy ndarray instances. """ + assert self._src_evt_idxs is not None evt_idxs = self._src_evt_idxs[1] out_arrays = [np.take(arr, evt_idxs) for arr in arrays] return out_arrays - def change_shg_mgr(self, shg_mgr, pmm): + def change_shg_mgr(self, shg_mgr: SourceHypoGroupManager, pmm: ParameterModelMapper): """This method is called when the source hypothesis group manager has changed. Hence, the source data fields need to get recalculated. @@ -649,16 +675,24 @@ def change_shg_mgr(self, shg_mgr, pmm): Parameters ---------- - shg_mgr : instance of SourceHypoGroupManager + shg_mgr The instance of SourceHypoGroupManager that defines the source hypothesis groups. - pmm : instance of ParameterModelMapper + pmm The instance of ParameterModelMapper that defines the global parameters and their mapping to local source parameter. """ self.calculate_source_data_fields(shg_mgr=shg_mgr, pmm=pmm) - def initialize_trial(self, shg_mgr, pmm, events, n_events=None, evt_sel_method=None, tl=None): + def initialize_trial( + self, + shg_mgr: SourceHypoGroupManager, + pmm: ParameterModelMapper, + events: DataFieldRecordArray, + n_events: int | None = None, + evt_sel_method: EventSelectionMethod | None = None, + tl: TimeLord | None = None, + ): """Initializes the trial data manager for a new trial. It sets the raw events, calculates pre-event-selection data fields, performs a possible event selection and calculates the static data fields for the left-over @@ -666,23 +700,23 @@ def initialize_trial(self, shg_mgr, pmm, events, n_events=None, evt_sel_method=N Parameters ---------- - shg_mgr : instance of SourceHypoGroupManager + shg_mgr The instance of SourceHypoGroupManager that defines the source hypothesis groups. - pmm : instance of ParameterModelMapper + pmm The instance of ParameterModelMapper, that defines the global parameters and their mapping to local source parameters. - events : DataFieldRecordArray instance + events The DataFieldRecordArray instance holding the entire raw events. - n_events : int | None + n_events The total number of events of the data set this trial data manager corresponds to. If None, the number of events is taken from the number of events present in the ``events`` array. - evt_sel_method : instance of EventSelectionMethod | None + evt_sel_method The optional event selection method that should be used to select potential signal events. - tl : instance of TimeLord | None + tl The optional TimeLord instance that should be used for timing measurements. """ @@ -694,6 +728,7 @@ def initialize_trial(self, shg_mgr, pmm, events, n_events=None, evt_sel_method=N # Save the number of sources. self._n_sources = shg_mgr.n_sources + assert self._events is not None if n_events is None: n_events = len(self._events) self.n_events = n_events @@ -704,7 +739,10 @@ def initialize_trial(self, shg_mgr, pmm, events, n_events=None, evt_sel_method=N if evt_sel_method is not None: logger.debug(f'Performing event selection method "{classname(evt_sel_method)}".') - (selected_events, src_evt_idxs) = evt_sel_method.select_events(events=self._events, tl=tl) + assert self._events is not None + _sel_result = evt_sel_method.select_events(events=self._events, tl=tl) + selected_events = _sel_result[0] + src_evt_idxs = _sel_result[1] logger.debug(f'Selected {len(selected_events)} out of {len(self._events)} events.') self.events = selected_events self._src_evt_idxs = src_evt_idxs @@ -712,76 +750,89 @@ def initialize_trial(self, shg_mgr, pmm, events, n_events=None, evt_sel_method=N # Sort the events by the index field, if a field was provided. if self._index_field_name is not None: logger.debug(f'Sorting events in index field "{self._index_field_name}"') + assert self._events is not None sorted_idxs = self._events.sort_by_field(self._index_field_name) # If event indices are stored, we need to re-assign also those event # indices according to the new order. if self._src_evt_idxs is not None: - self._src_evt_idxs[1] = np.take(sorted_idxs, self._src_evt_idxs[1]) + self._src_evt_idxs = ( + self._src_evt_idxs[0], + np.take(sorted_idxs, self._src_evt_idxs[1]), + ) # Create the src_evt_idxs property data in case it was not provided by # the event selection. In that case all events are selected for all # sources. This simplifies the implementations of the PDFs. if self._src_evt_idxs is None: + assert self._n_sources is not None + _n_sources = int(self._n_sources) + _n_sel_events = self.n_selected_events self._src_evt_idxs = ( - np.repeat(np.arange(self.n_sources), self.n_selected_events), - np.tile(np.arange(self.n_selected_events), self.n_sources), + np.repeat(np.arange(_n_sources), _n_sel_events), + np.tile(np.arange(_n_sel_events), _n_sources), ) # Now calculate all the static data fields. This will increment the # trial data state ID. self.calculate_static_data_fields(shg_mgr=shg_mgr, pmm=pmm) - def get_n_values(self): + def get_n_values(self) -> int: """Returns the expected size of the values array after a PDF evaluation, which will include PDF values for all trial data events and all sources. Returns ------- - n : int + n The length of the expected values array after a PDF evaluation. """ + assert self._src_evt_idxs is not None return len(self._src_evt_idxs[0]) - def get_values_mask_for_source_mask(self, src_mask): + def get_values_mask_for_source_mask(self, src_mask: np.ndarray) -> np.ndarray: """Creates a boolean mask for the values array where entries belonging to the sources given by the source mask are selected. Parameters ---------- - src_mask : instance of numpy ndarray + src_mask The (N_sources,)-shaped numpy ndarray holding the boolean selection of the sources. Returns ------- - values_mask : instance of numpy ndarray + values_mask The (N_values,)-shaped numpy ndarray holding the boolean selection of the values. """ + assert self.src_evt_idxs is not None + assert self._n_sources is not None tdm_src_idxs = self.src_evt_idxs[0] - src_idxs = np.arange(self.n_sources)[src_mask] + src_idxs = np.arange(self._n_sources)[src_mask] values_mask = np.zeros((self.get_n_values(),), dtype=np.bool_) def make_values_mask(src_idx): - global values_mask + """Updates the values mask in-place to also select the values + belonging to the source with the given index. + """ + nonlocal values_mask values_mask |= tdm_src_idxs == src_idx np.vectorize(make_values_mask)(src_idxs) return values_mask - def add_source_data_field(self, name, func, dt=None): + def add_source_data_field(self, name: str, func: Callable, dt: np.dtype | str | None = None): """Adds a new data field to the manager. The data field must depend solely on source parameters. Parameters ---------- - name : str + name The name of the data field. It serves as the identifier for the data field. - func : callable + func The function that calculates the data field values. The call signature must be @@ -790,7 +841,7 @@ def add_source_data_field(self, name, func, dt=None): where ``tdm`` is the TrialDataManager instance holding the event data, ``shg_mgr`` is the instance of SourceHypoGroupManager, and ``pmm`` is the instance of ParameterModelMapper. - dt : numpy dtype | str | None + dt If specified it defines the data type this data field should have. If a str instance is given, it defines the name of the data field whose data type should be taken for the data field. @@ -802,15 +853,23 @@ def add_source_data_field(self, name, func, dt=None): self._source_data_fields_dict[name] = data_field - def add_data_field(self, name, func, global_fitparam_names=None, dt=None, pre_evt_sel=False, is_srcevt_data=False): + def add_data_field( + self, + name: str, + func: Callable, + global_fitparam_names: str | Sequence[str] | None = None, + dt: np.dtype | str | None = None, + pre_evt_sel: bool = False, + is_srcevt_data: bool = False, + ): """Adds a new data field to the manager. Parameters ---------- - name : str + name The name of the data field. It serves as the identifier for the data field. - func : callable + func The function that calculates the data field values. The call signature must be @@ -822,19 +881,19 @@ def add_data_field(self, name, func, global_fitparam_names=None, dt=None, pre_ev ``global_fitparams_dict`` is the dictionary with the current global fit parameter names and values. The shape of the returned array must be (N_selected_events,). - global_fitparam_names : str | sequence of str | None + global_fitparam_names The sequence of str instances specifying the names of the global fit parameters this data field depends on. If set to ``None``, it means that the data field does not depend on any fit parameters. - dt : numpy dtype | str | None + dt If specified it defines the data type this data field should have. If a str instance is given, it defines the name of the data field whose data type should be taken for the data field. - pre_evt_sel : bool + pre_evt_sel Flag if this data field should get calculated before potential signal events get selected (True), or afterwards (False). Default is False. - is_srcevt_data : bool + is_srcevt_data Flag if this data field contains source-event data, hence the length of the data array will be N_values. Default is False. @@ -871,16 +930,16 @@ def add_data_field(self, name, func, global_fitparam_names=None, dt=None, pre_ev else: self._global_fitparam_data_fields_dict[name] = data_field - def calculate_source_data_fields(self, shg_mgr, pmm): + def calculate_source_data_fields(self, shg_mgr: SourceHypoGroupManager, pmm: ParameterModelMapper): """Calculates the data values of the data fields that solely depend on source parameters. Parameters ---------- - shg_mgr : instance of SourceHypoGroupManager + shg_mgr The instance of SourceHypoGroupManager, which defines the groups of source hypotheses. - pmm : instance of ParameterModelMapper + pmm The instance of ParameterModelMapper, that defines the global parameters and their mapping to local source parameters. """ @@ -892,17 +951,17 @@ def calculate_source_data_fields(self, shg_mgr, pmm): self._trial_data_state_id += 1 - def calculate_pre_evt_sel_static_data_fields(self, shg_mgr, pmm): + def calculate_pre_evt_sel_static_data_fields(self, shg_mgr: SourceHypoGroupManager, pmm: ParameterModelMapper): """Calculates the data values of the data fields that should be available for the event selection method and do not depend on any fit parameters. Parameters ---------- - shg_mgr : instance of SourceHypoGroupManager + shg_mgr The instance of SourceHypoGroupManager, which defines the groups of source hypotheses. - pmm : instance of ParameterModelMapper + pmm The instance of ParameterModelMapper, that defines the global parameters and their mapping to local source parameters. """ @@ -914,16 +973,16 @@ def calculate_pre_evt_sel_static_data_fields(self, shg_mgr, pmm): self._trial_data_state_id += 1 - def calculate_static_data_fields(self, shg_mgr, pmm): + def calculate_static_data_fields(self, shg_mgr: SourceHypoGroupManager, pmm: ParameterModelMapper): """Calculates the data values of the data fields that do not depend on any source or fit parameters. Parameters ---------- - shg_mgr : instance of SourceHypoGroupManager + shg_mgr The instance of SourceHypoGroupManager, which defines the groups of source hypotheses. - pmm : instance of ParameterModelMapper + pmm The instance of ParameterModelMapper, that defines the global parameters and their mapping to local source parameters. """ @@ -935,19 +994,21 @@ def calculate_static_data_fields(self, shg_mgr, pmm): self._trial_data_state_id += 1 - def calculate_global_fitparam_data_fields(self, shg_mgr, pmm, global_fitparams_dict): + def calculate_global_fitparam_data_fields( + self, shg_mgr: SourceHypoGroupManager, pmm: ParameterModelMapper, global_fitparams_dict: dict + ): """Calculates the data values of the data fields that depend on global fit parameter values. Parameters ---------- - shg_mgr : instance of SourceHypoGroupManager + shg_mgr The instance of SourceHypoGroupManager, which defines the groups of source hypotheses. - pmm : instance of ParameterModelMapper + pmm The instance of ParameterModelMapper, that defines the global parameters and their mapping to local source parameters. - global_fitparams_dict : dict + global_fitparams_dict The dictionary holding the current global fit parameter names and values. """ @@ -959,7 +1020,7 @@ def calculate_global_fitparam_data_fields(self, shg_mgr, pmm, global_fitparams_d self._trial_data_state_id += 1 - def get_data(self, name): + def get_data(self, name: str) -> np.ndarray: """Gets the data for the given data field name. The data is stored either in the raw events DataFieldRecordArray or in one of the additional defined data fields. Data from the raw events @@ -967,12 +1028,12 @@ def get_data(self, name): Parameters ---------- - name : str + name The name of the data field for which to retrieve the data. Returns ------- - data : instance of numpy ndarray + data The numpy ndarray holding the data of the requested data field. The length of the array is either N_sources, N_selected_events, or N_values. @@ -1000,17 +1061,17 @@ def get_data(self, name): raise KeyError(f'The data field "{name}" is not defined!') - def get_dtype(self, name): + def get_dtype(self, name: str) -> np.dtype: """Gets the data type of the given data field. Parameters ---------- - name : str + name The name of the data field whose data type should get retrieved. Returns ------- - dt : numpy dtype + dt The numpy dtype object of the given data field. Raises @@ -1022,52 +1083,52 @@ def get_dtype(self, name): return dt - def is_event_data_field(self, name): + def is_event_data_field(self, name: str) -> bool: """Checks if the given data field is an events data field, i.e. its length is N_selected_events. Parameters ---------- - name : str + name The name of the data field. Returns ------- - check : bool + check ``True`` if the given data field contains event data, ``False`` otherwise. """ return self._events is not None and name in self._events.field_name_list - def is_source_data_field(self, name): + def is_source_data_field(self, name: str) -> bool: """Checks if the given data field is a source data field, i.e. its length is N_sources. Parameters ---------- - name : str + name The name of the data field. Returns ------- - check : bool + check ``True`` if the given data field contains source data, ``False`` otherwise. """ return name in self._source_data_fields_dict - def is_srcevt_data_field(self, name): + def is_srcevt_data_field(self, name: str) -> bool: """Checks if the given data field is a source-event data field, i.e. its length is N_values. Parameters ---------- - name : str + name The name of the data field. Returns ------- - check : bool + check ``True`` if the given data field contains source-event data, ``False`` otherwise. """ diff --git a/skyllh/core/types.py b/skyllh/core/types.py index 63d886e41b..8717a48782 100644 --- a/skyllh/core/types.py +++ b/skyllh/core/types.py @@ -4,5 +4,12 @@ class SourceHypoGroup_t: + """This is the base type for the + :class:`~skyllh.core.source_hypo_grouping.SourceHypoGroup` class. It exists + to allow type checks without importing the actual class, avoiding circular + imports. + """ + def __init__(self, *args, **kwargs) -> None: + """Creates a new instance of SourceHypoGroup_t.""" super().__init__(*args, **kwargs) diff --git a/skyllh/core/utils/analysis.py b/skyllh/core/utils/analysis.py index df4afabef5..19752ae7fe 100644 --- a/skyllh/core/utils/analysis.py +++ b/skyllh/core/utils/analysis.py @@ -18,13 +18,18 @@ try: from iminuit import minimize except ImportError: + minimize = None IMINUIT_LOADED = False else: IMINUIT_LOADED = True +from collections.abc import Callable, Sequence + +from skyllh.core.analysis import Analysis from skyllh.core.logging import ( get_logger, ) +from skyllh.core.parameters import ParameterModelMapper from skyllh.core.progressbar import ( ProgressBar, ) @@ -34,15 +39,19 @@ issequence, issequenceof, ) +from skyllh.core.random import RandomStateService from skyllh.core.session import ( is_interactive_session, ) +from skyllh.core.source_hypo_grouping import SourceHypoGroupManager from skyllh.core.source_model import ( PointLikeSource, ) from skyllh.core.storage import ( NPYFileLoader, ) +from skyllh.core.timing import TimeLord +from skyllh.core.trialdata import TrialDataManager from skyllh.core.utils.spline import ( make_spline_1d, ) @@ -51,7 +60,9 @@ """ -def pointlikesource_to_data_field_array(tdm, shg_mgr, pmm): +def pointlikesource_to_data_field_array( + tdm: TrialDataManager, shg_mgr: SourceHypoGroupManager, pmm: ParameterModelMapper +) -> np.ndarray: """Function to transform a list of PointLikeSource sources into a numpy record ndarray. The resulting numpy record ndarray contains the following fields: @@ -65,17 +76,17 @@ def pointlikesource_to_data_field_array(tdm, shg_mgr, pmm): Parameters ---------- - tdm : instance of TrialDataManager + tdm The TrialDataManager instance. - shg_mgr : instance of SourceHypoGroupManager + shg_mgr The instance of SourceHypoGroupManager that defines the sources. - pmm : instance of ParameterModelMapper + pmm The instance of ParameterModelMapper that defines the mapping of global parameters to local model parameters. Returns ------- - arr : (N_sources,)-shaped numpy record ndarray + arr The numpy record ndarray holding the source parameters. """ sources = shg_mgr.source_list @@ -101,7 +112,9 @@ def pointlikesource_to_data_field_array(tdm, shg_mgr, pmm): return arr -def calculate_pval_from_trials(ts_vals, ts_threshold, comp_operator='greater_equal'): +def calculate_pval_from_trials( + ts_vals: np.ndarray, ts_threshold: float, comp_operator: str = 'greater_equal' +) -> tuple[float, float]: """Calculates the percentage (p-value) of test-statistic trials that are above the given test-statistic critical value. In addition it calculates the standard deviation of the p-value assuming @@ -109,11 +122,11 @@ def calculate_pval_from_trials(ts_vals, ts_threshold, comp_operator='greater_equ Parameters ---------- - ts_vals : (n_trials,)-shaped 1D ndarray of float + ts_vals The ndarray holding the test-statistic values of the trials. - ts_threshold : float + ts_threshold The critical test-statistic value. - comp_operator: string, optional + comp_operator The comparison operator for p-value calculation. It can be set to one of the following options: 'greater' or 'greater_equal'. @@ -133,21 +146,23 @@ def calculate_pval_from_trials(ts_vals, ts_threshold, comp_operator='greater_equ return (p, p_sigma) -def calculate_pval_from_gammafit_to_trials(ts_vals, ts_threshold, eta=3.0, n_max=500000): +def calculate_pval_from_gammafit_to_trials( + ts_vals: np.ndarray, ts_threshold: float, eta: float = 3.0, n_max: int = 500000 +): """Calculates the probability (p-value) of test-statistic exceeding the given test-statistic threshold. This calculation relies on fitting a gamma distribution to a list of ts values. Parameters ---------- - ts_vals : (n_trials,)-shaped 1D ndarray of float + ts_vals The ndarray holding the test-statistic values of the trials. - ts_threshold : float + ts_threshold The critical test-statistic value. - eta : float, optional + eta Test-statistic value at which the gamma function is truncated from below. Default = 3.0. - n_max : int, optional + n_max The maximum number of trials that should be used during fitting. Default = 500,000 @@ -179,7 +194,12 @@ def calculate_pval_from_gammafit_to_trials(ts_vals, ts_threshold, eta=3.0, n_max def calculate_pval_from_trials_mixed( - ts_vals, ts_threshold, switch_at_ts=3.0, eta=None, n_max=500000, comp_operator='greater_equal' + ts_vals: np.ndarray, + ts_threshold: float, + switch_at_ts: float = 3.0, + eta: float | None = None, + n_max: int = 500000, + comp_operator: str = 'greater_equal', ): """Calculates the probability (p-value) of test-statistic exceeding the given test-statistic threshold. This calculation relies on fitting @@ -189,21 +209,21 @@ def calculate_pval_from_trials_mixed( Parameters ---------- - ts_vals : (n_trials,)-shaped 1D ndarray of float + ts_vals The ndarray holding the test-statistic values of the trials. - ts_threshold : float + ts_threshold The critical test-statistic value. - switch_at_ts : float, optional + switch_at_ts Test-statistic value below which p-value is computed from trials directly. For thresholds greater than switch_at_ts the pvalue is calculated using a gamma fit. - eta : float, optional + eta Test-statistic value at which the gamma function is truncated from below. Default is None. - n_max : int, optional + n_max The maximum number of trials that should be used during fitting. Default = 500,000 - comp_operator: string, optional + comp_operator The comparison operator for p-value calculation. It can be set to one of the following options: 'greater' or 'greater_equal'. @@ -222,23 +242,23 @@ def calculate_pval_from_trials_mixed( return calculate_pval_from_gammafit_to_trials(ts_vals, ts_threshold, eta=eta, n_max=n_max) -def truncated_gamma_logpdf(a, scale, eta, ts_above_eta, N_above_eta): +def truncated_gamma_logpdf(a: float, scale: float, eta: float, ts_above_eta: np.ndarray, N_above_eta: int): """Calculates the -log(likelihood) of a sample of random numbers generated from a gamma pdf truncated from below at x=eta. Parameters ---------- - a : float + a Shape parameter. - scale : float + scale Scale parameter. - eta : float + eta Test-statistic value at which the gamma function is truncated from below. - ts_above_eta : (n_trials,)-shaped 1D ndarray + ts_above_eta The ndarray holding the test-statistic values falling in the truncated gamma pdf. - N_above_eta : int + N_above_eta Number of test-statistic values falling in the truncated gamma pdf. @@ -254,7 +274,7 @@ def truncated_gamma_logpdf(a, scale, eta, ts_above_eta, N_above_eta): return -logl -def fit_truncated_gamma(vals, eta): +def fit_truncated_gamma(vals: np.ndarray, eta: float) -> tuple[np.ndarray, float]: """ Fits a truncated gamma function to a set of values. Returns the best-fit parameters and the normalization constant @@ -262,15 +282,15 @@ def fit_truncated_gamma(vals, eta): Parameters ---------- - vals : (n_trials,)-shaped 1D ndarray of float - eta : float + vals + eta Value at which the gamma function is truncated from below. Returns ------- - pars : (2,)-shaped 1D array of float + pars `a` and `scale` parameters of the truncated gamma function. - norm : float + norm Normalization constant of the truncated gamma function. """ @@ -280,6 +300,7 @@ def fit_truncated_gamma(vals, eta): 'This module is a requirement of the function ' '"calculate_critical_ts_from_gamma"!' ) + assert minimize is not None Ntot = len(vals) vals_eta = vals[vals > eta] @@ -287,6 +308,9 @@ def fit_truncated_gamma(vals, eta): alpha = N_prime / Ntot def obj(x): + """The objective function passed to the minimizer, evaluating the + truncated gamma log-pdf for the parameter vector ``x`` = (a, scale). + """ return truncated_gamma_logpdf(x[0], x[1], eta=eta, ts_above_eta=vals_eta, N_above_eta=N_prime) x0 = [0.75, 1.8] # Initial values of function parameters. @@ -294,29 +318,29 @@ def obj(x): r = minimize(obj, x0, bounds=bounds) pars = r.x - norm = alpha / gamma.sf(eta, a=pars[0], scale=pars[1]) + norm = float(alpha / gamma.sf(eta, a=pars[0], scale=pars[1])) return pars, norm -def calculate_critical_ts_from_gamma(ts, h0_ts_quantile, eta=3.0): +def calculate_critical_ts_from_gamma(ts: np.ndarray, h0_ts_quantile: float, eta: float = 3.0) -> float: """Calculates the critical test-statistic value corresponding to h0_ts_quantile by fitting the ts distribution with a truncated gamma function. Parameters ---------- - ts : (n_trials,)-shaped 1D ndarray + ts The ndarray holding the test-statistic values of the trials. - h0_ts_quantile : float + h0_ts_quantile Null-hypothesis test statistic quantile. - eta : float, optional + eta Test-statistic value at which the gamma function is truncated from below. Returns ------- - critical_ts : float + critical_ts """ (pars, norm) = fit_truncated_gamma(vals=ts, eta=eta) critical_ts = gamma.ppf(1 - 1.0 / norm * h0_ts_quantile, a=pars[0], scale=pars[1]) @@ -330,32 +354,33 @@ def calculate_critical_ts_from_gamma(ts, h0_ts_quantile, eta=3.0): eta, ) - return critical_ts + return float(critical_ts) -def polynomial_fit(ns, p, p_weight, deg, p_thr): +def polynomial_fit(ns: Sequence, p: Sequence, p_weight: Sequence, deg: int, p_thr: float) -> float: """Performs a polynomial fit on the p-values of test-statistic trials - associated to each ns.. + associated to each ns. Using the fitted parameters it computes the number of signal events correponding to the given p-value critical value. Parameters ---------- - ns : 1D array_like object + ns x-coordinates of the sample. - p : 1D array_like object + p y-coordinates of the sample. - p_weight : 1D array_like object + p_weight Weights to apply to the y-coordinates of the sample points. For gaussian uncertainties, use 1/sigma. - deg : int + deg Degree of the fitting polynomial function. - p_thr : float within [0,1] + p_thr The critical p-value. Returns ------- - ns : float + x + The number of signal events corresponding to the given p-value critical value. """ (params, _) = np.polyfit(ns, p, deg, w=p_weight, cov=True) @@ -368,92 +393,90 @@ def polynomial_fit(ns, p, p_weight, deg, p_thr): if deg == 1: (a, b) = (params[0], params[1]) - ns = (p_thr - b) / a - return ns + return float((p_thr - b) / a) elif deg == 2: (a, b, c) = (params[0], params[1], params[2]) - ns = (-b + np.sqrt((b**2) - 4 * a * (c - p_thr))) / (2 * a) - return ns + return float((-b + np.sqrt((b**2) - 4 * a * (c - p_thr))) / (2 * a)) else: raise ValueError('deg = %g is not valid. The order of the polynomial function must be 1 or 2.', deg) def estimate_mean_nsignal_for_ts_quantile( - ana, - rss, - p, - eps_p, - mu_range, - critical_ts=None, - h0_trials=None, - h0_ts_quantile=None, - min_dmu=0.5, - bkg_kwargs=None, - sig_kwargs=None, - ppbar=None, - tl=None, - pathfilename=None, -): + ana: Analysis, + rss: RandomStateService, + p: float, + eps_p: float, + mu_range: Sequence, + critical_ts: float | None = None, + h0_trials: np.ndarray | None = None, + h0_ts_quantile: float | None = None, + min_dmu: float = 0.5, + bkg_kwargs: dict | None = None, + sig_kwargs: dict | None = None, + ppbar: ProgressBar | None = None, + tl: TimeLord | None = None, + pathfilename: str | None = None, +) -> tuple[float, float | None]: """Calculates the mean number of signal events needed to be injected to reach a test statistic distribution with defined properties for the given analysis. Parameters ---------- - ana : Analysis instance + ana The Analysis instance to use for the calculation. - rss : instance of RandomStateService + rss The RandomStateService instance to use for generating random numbers. - p : float + p Desired probability of signal test statistic for exceeding `h0_ts_quantile` part of null-hypothesis test statistic threshold. - eps_p : float + eps_p Precision in `p` as stopping condition for the calculation. - mu_range : 2-element sequence + mu_range The range of mu (lower,upper) to search for mean number of signal events. - critical_ts : float | None + critical_ts The critical test-statistic value that should be overcome by the signal distribution. If set to None, the null-hypothesis test-statistic distribution will be used to compute the critical TS value. - h0_trials : (n_h0_trials,)-shaped ndarray | None + h0_trials The structured ndarray holding the trials for the null-hypothesis. If set to `None`, the number of trials is calculated from binomial statistics via `h0_ts_quantile*(1-h0_ts_quantile)/eps**2`, where `eps` is `min(5e-3, h0_ts_quantile/10)`. - h0_ts_quantile : float | None + h0_ts_quantile Null-hypothesis test statistic quantile. If set to None, the critical test-statistic value that should be overcome by the signal distribution MUST be given. - min_dmu : float + min_dmu The minimum delta mu to use for calculating the derivative dmu/dp. The default is ``0.5``. - bkg_kwargs : dict | None + bkg_kwargs Additional keyword arguments for the `generate_events` method of the background generation method class. An usual keyword argument is `poisson`. - sig_kwargs : dict | None + sig_kwargs Additional keyword arguments for the `generate_signal_events` method of the `SignalGenerator` class. An usual keyword argument is `poisson`. If `poisson` is set to True, the actual number of generated signal events will be drawn from a Poisson distribution with the mean number of signal events, mu. - ppbar : instance of ProgressBar | None + ppbar The possible parent ProgressBar instance. - tl: instance of TimeLord | None + tl The optional TimeLord instance that should be used to collect timing information about this function. - pathfilename: string | None + pathfilename Trial data file path including the filename. If set to None, generatedtrials won't be saved. Returns ------- - mu : float + mu Estimated mean number of signal events. - mu_err : None + mu_err Error estimate needs to be implemented. """ logger = get_logger(__name__) @@ -468,6 +491,7 @@ def estimate_mean_nsignal_for_ts_quantile( 'the type of test to run.' ) elif critical_ts is None: + assert h0_ts_quantile is not None n_trials_max = int(5.0e5) # Via binomial statistics, calcuate the minimum number of trials # needed to get the required precision on the critial TS value. @@ -501,7 +525,7 @@ def estimate_mean_nsignal_for_ts_quantile( np.save(pathfilename, h0_ts_vals) else: if h0_trials.size < n_trials_total: - if 'seed' not in h0_trials.dtype.names: + if h0_trials.dtype.names is None or 'seed' not in h0_trials.dtype.names: logger.debug( 'Uploaded trials miss the rss_seed field. ' 'Will not be possible to extend the trial file ' @@ -586,7 +610,8 @@ def estimate_mean_nsignal_for_ts_quantile( # Initially generate trials for a 5-times larger uncertainty ``eps_p`` # to catch ns0 points far away from the desired propability quicker. dn_trials = max(100, int(n_trials / 5**2 + 0.5)) - (ts_vals0, p0_sigma, delta_p) = ([], 2 * eps_p, 0) + ts_vals0: np.ndarray = np.array([]) + (p0, p0_sigma, delta_p) = (0.0, 2 * eps_p, 0.0) while (delta_p < p0_sigma * 5) and (p0_sigma > eps_p): ts_vals0 = np.concatenate( ( @@ -860,20 +885,20 @@ def estimate_mean_nsignal_for_ts_quantile( def estimate_sensitivity( - ana, - rss, - h0_trials=None, - h0_ts_quantile=0.5, - p=0.9, - eps_p=0.005, - mu_range=None, - min_dmu=0.5, - bkg_kwargs=None, - sig_kwargs=None, - ppbar=None, - tl=None, - pathfilename=None, -): + ana: Analysis, + rss: RandomStateService, + h0_trials: np.ndarray | None = None, + h0_ts_quantile: float = 0.5, + p: float = 0.9, + eps_p: float = 0.005, + mu_range: Sequence | None = None, + min_dmu: float = 0.5, + bkg_kwargs: dict | None = None, + sig_kwargs: dict | None = None, + ppbar: ProgressBar | None = None, + tl: TimeLord | None = None, + pathfilename: str | None = None, +) -> tuple[float, float | None]: """Estimates the mean number of signal events that whould have to be injected into the data such that the test-statistic value of p*100% of all trials are larger than the critical test-statistic value c, which @@ -885,54 +910,54 @@ def estimate_sensitivity( Parameters ---------- - ana : Analysis + ana The Analysis instance to use for sensitivity estimation. - rss : RandomStateService + rss The RandomStateService instance to use for generating random numbers. - h0_trials : (n_h0_ts_vals,)-shaped ndarray | None + h0_trials The strutured ndarray holding the trials for the null-hypothesis. If set to `None`, the number of trials is calculated from binomial statistics via `h0_ts_quantile*(1-h0_ts_quantile)/eps**2`, where `eps` is `min(5e-3, h0_ts_quantile/10)`. - h0_ts_quantile : float, optional + h0_ts_quantile Null-hypothesis test statistic quantile that defines the critical value. - p : float, optional + p Desired probability of the signal test statistic value to exceed the null-hypothesis test statistic value threshold, which is defined through the `h0_ts_quantile` value. - eps_p : float, optional + eps_p Precision in `p` for execution to break. - mu_range : 2-element sequence | None + mu_range Range to search for the mean number of signal events. If set to None, the range (0, 10) will be used. - min_dmu : float + min_dmu The minimum delta mu to use for calculating the derivative dmu/dp. The default is ``0.5``. - bkg_kwargs : dict | None + bkg_kwargs Additional keyword arguments for the `generate_events` method of the background generation method class. An usual keyword argument is `poisson`. - sig_kwargs : dict | None + sig_kwargs Additional keyword arguments for the `generate_signal_events` method of the `SignalGenerator` class. An usual keyword argument is `poisson`. If `poisson` is set to True, the actual number of generated signal events will be drawn from a Poisson distribution with the mean number of signal events, mu. - ppbar : instance of ProgressBar | None + ppbar The possible parent ProgressBar instance. - tl: instance of TimeLord | None + tl The optional TimeLord instance that should be used to collect timing information about this function. - pathfilename : string | None + pathfilename Trial data file path including the filename. If set to None, generated trials won't be saved. Returns ------- - mu : float + mu Estimated median number of signal events to reach desired sensitivity. - mu_err : float + mu_err The uncertainty of the estimated mean number of signal events. """ if mu_range is None: @@ -958,20 +983,20 @@ def estimate_sensitivity( def estimate_discovery_potential( - ana, - rss, - h0_trials=None, - h0_ts_quantile=2.8665e-7, - p=0.5, - eps_p=0.005, - mu_range=None, - min_dmu=0.5, - bkg_kwargs=None, - sig_kwargs=None, - ppbar=None, - tl=None, - pathfilename=None, -): + ana: Analysis, + rss: RandomStateService, + h0_trials: np.ndarray | None = None, + h0_ts_quantile: float = 2.8665e-7, + p: float = 0.5, + eps_p: float = 0.005, + mu_range: Sequence | None = None, + min_dmu: float = 0.5, + bkg_kwargs: dict | None = None, + sig_kwargs: dict | None = None, + ppbar: ProgressBar | None = None, + tl: TimeLord | None = None, + pathfilename: str | None = None, +) -> tuple[float, float | None]: """Estimates the mean number of signal events that whould have to be injected into the data such that the test-statistic value of p*100% of all trials are larger than the critical test-statistic value c, which @@ -983,54 +1008,54 @@ def estimate_discovery_potential( Parameters ---------- - ana : Analysis + ana The Analysis instance to use for discovery potential estimation. - rss : RandomStateService + rss The RandomStateService instance to use for generating random numbers. - h0_trials : (n_h0_ts_vals,)-shaped ndarray | None + h0_trials The structured ndarray holding the trials for the null-hypothesis. If set to `None`, the number of trials is calculated from binomial statistics via `h0_ts_quantile*(1-h0_ts_quantile)/eps**2`, where `eps` is `min(5e-3, h0_ts_quantile/10)`. - h0_ts_quantile : float, optional + h0_ts_quantile Null-hypothesis test statistic quantile that defines the critical value. - p : float, optional + p Desired probability of the signal test statistic value to exceed the critical value. - eps_p : float, optional + eps_p Precision in `p` for execution to break. - mu_range : 2-element sequence | None + mu_range Range to search for the mean number of signal events. If set to None, the range (0, 10) will be used. - min_dmu : float + min_dmu The minimum delta mu to use for calculating the derivative dmu/dp. The default is ``0.5``. - bkg_kwargs : dict | None + bkg_kwargs Additional keyword arguments for the `generate_events` method of the background generation method class. An usual keyword argument is `poisson`. - sig_kwargs : dict | None + sig_kwargs Additional keyword arguments for the `generate_signal_events` method of the `SignalGenerator` class. An usual keyword argument is `poisson`. If `poisson` is set to True, the actual number of generated signal events will be drawn from a Poisson distribution with the mean number of signal events, mu. - ppbar : instance of ProgressBar | None + ppbar The possible parent ProgressBar instance. - tl: instance of TimeLord | None + tl The optional TimeLord instance that should be used to collect timing information about this function. - pathfilename : string | None + pathfilename Trial data file path including the filename. If set to None, generated trials won't be saved. Returns ------- - mu : float + mu Estimated mean number of injected signal events to reach the desired discovery potential. - mu_err : float + mu_err Estimated error of `mu`. """ if mu_range is None: @@ -1056,19 +1081,19 @@ def estimate_discovery_potential( def generate_mu_of_p_spline_interpolation( - ana, - rss, - h0_ts_vals, - h0_ts_quantile, - eps_p, - mu_range, - mu_step, - kind='cubic', - bkg_kwargs=None, - sig_kwargs=None, - ppbar=None, - tl=None, -): + ana: Analysis, + rss: RandomStateService, + h0_ts_vals: np.ndarray | None, + h0_ts_quantile: float, + eps_p: float, + mu_range: Sequence, + mu_step: float, + kind: str = 'cubic', + bkg_kwargs: dict | None = None, + sig_kwargs: dict | None = None, + ppbar: ProgressBar | None = None, + tl: TimeLord | None = None, +) -> Callable: """Generates a spline interpolation for mu(p) function for a pre-defined range of mu, where mu is the mean number of injected signal events and p the probability for the ts value larger than the ts value corresponding to the @@ -1077,47 +1102,47 @@ def generate_mu_of_p_spline_interpolation( Parameters ---------- - ana : instance of Analysis + ana The Analysis instance to use for the calculation. - rss : instance of RandomStateService + rss The RandomStateService instance to use for generating random numbers. - h0_ts_vals : (n_h0_ts_vals,)-shaped 1D ndarray | None + h0_ts_vals The 1D ndarray holding the test-statistic values for the null-hypothesis. If set to `None`, 100/(1-h0_ts_quantile) null-hypothesis trials will be generated. - h0_ts_quantile : float + h0_ts_quantile Null-hypothesis test statistic quantile, which should be exceeded by the alternative hypothesis ts value. - eps_p : float + eps_p The one sigma precision in `p` as stopping condition for the calculation for a single mu value. - mu_range : 2-element sequence + mu_range The range (lower,upper) of mean number of injected signal events to create the interpolation spline for. - mu_step : float + mu_step The step size of the mean number of signal events. - kind : str + kind The kind of spline to generate. Possble values are 'linear' and 'cubic' (default). - bkg_kwargs : dict | None + bkg_kwargs Additional keyword arguments for the `generate_events` method of the background generation method class. An usual keyword argument is `poisson`. - sig_kwargs : dict | None + sig_kwargs Additional keyword arguments for the `generate_signal_events` method of the `SignalGenerator` class. An usual keyword argument is `poisson`. If `poisson` is set to True, the actual number of generated signal events will be drawn from a Poisson distribution with the mean number of signal events, mu. - ppbar : instance of ProgressBar | None + ppbar The possible parent ProgressBar instance. - tl: instance of TimeLord | None + tl The optional TimeLord instance that should be used to collect timing information about this function. Returns ------- - spline : callable + spline The spline function mu(p). """ logger = get_logger(__name__) @@ -1154,7 +1179,8 @@ def generate_mu_of_p_spline_interpolation( for idx, mu in enumerate(mu_vals): p = None - (ts_vals, p_sigma) = ([], 2 * eps_p) + ts_vals: np.ndarray = np.array([]) + p_sigma = 2 * eps_p while p_sigma > eps_p: ts_vals = np.concatenate( ( @@ -1182,34 +1208,34 @@ def generate_mu_of_p_spline_interpolation( def create_trial_data_file( - ana, - rss, - n_trials, + ana: Analysis, + rss: RandomStateService, + n_trials: int, mean_n_sig=0, mean_n_sig_null=0, - mean_n_bkg_list=None, - minimizer_rss=None, - bkg_kwargs=None, - sig_kwargs=None, - pathfilename=None, - ncpu=None, - ppbar=None, - tl=None, -): + mean_n_bkg_list: list[float] | None = None, + minimizer_rss: RandomStateService | None = None, + bkg_kwargs: dict | None = None, + sig_kwargs: dict | None = None, + pathfilename: str | None = None, + ncpu: int | None = None, + ppbar: ProgressBar | None = None, + tl: TimeLord | None = None, +) -> tuple[int, np.ndarray, np.ndarray, np.ndarray]: """Creates and fills a trial data file with `n_trials` generated trials for each mean number of injected signal events specified by `mean_n_sig` for a given analysis. Parameters ---------- - ana : instance of Analysis + ana The Analysis instance to use for the trial generation. - rss : instance of RandomStateService + rss The RandomStateService instance to use for generating random numbers. - n_trials : int + n_trials The number of trials to perform for each hypothesis test. - mean_n_sig : ndarray of float | float | 2- or 3-element sequence of float + mean_n_sig The array of mean number of injected signal events (MNOISEs) for which to generate trials. If this argument is not a ndarray, an array of MNOISEs is generated based on this argument. @@ -1218,7 +1244,7 @@ def create_trial_data_file( MNOISEs with a step size of one. If a 3-element sequence of floats is given, it specifies the range plus the step size of the MNOISEs. - mean_n_sig_null : ndarray of float | float | 2- or 3-element sequence of float + mean_n_sig_null The array of the fixed mean number of signal events (FMNOSEs) for the null-hypothesis for which to generate trials. If this argument is not a ndarray, an array of FMNOSEs is generated based on this argument. @@ -1227,50 +1253,53 @@ def create_trial_data_file( FMNOSEs with a step size of one. If a 3-element sequence of floats is given, it specifies the range plus the step size of the FMNOSEs. - mean_n_bkg_list : list of float | None + mean_n_bkg_list The mean number of background events that should be generated for each dataset. This parameter is passed to the ``do_trials`` method of the ``Analysis`` class. If set to None (the default), the background generation method needs to obtain this number itself. - minimizer_rss : instance of RandomStateService | None + minimizer_rss The instance of RandomStateService to use for generating random numbers for the minimizer, e.g. new initial fit parameter values. If set to ``None``, a rss with the same seed as ``rss`` will be initialized. - bkg_kwargs : dict | None + bkg_kwargs Additional keyword arguments for the `generate_events` method of the background generation method class. An usual keyword argument is `poisson`. - sig_kwargs : dict | None + sig_kwargs Additional keyword arguments for the `generate_signal_events` method of the `SignalGenerator` class. An usual keyword argument is `poisson`. - pathfilename : string | None + pathfilename Trial data file path including the filename. If set to None generated trials won't be saved. - ncpu : int | None + ncpu The number of CPUs to use. - ppbar : instance of ProgressBar | None + ppbar The optional instance of the parent progress bar. - tl: instance of TimeLord | None + tl The instance of TimeLord that should be used to measure individual tasks. Returns ------- - seed : int + seed The seed used to generate the trials. - mean_n_sig : 1d ndarray + mean_n_sig The array holding the mean number of signal events used to generate the trials. - mean_n_sig_null : 1d ndarray + mean_n_sig_null The array holding the fixed mean number of signal events for the null-hypothesis used to generate the trials. - trial_data : structured numpy ndarray + trial_data The generated trial data. """ n_trials = int_cast(n_trials, 'The n_trials argument must be castable to type int!') + mean_n_sig_min: float = 0 + mean_n_sig_max: float = 0 + mean_n_sig_step: float = 1 if not isinstance(mean_n_sig, np.ndarray): if not issequence(mean_n_sig): mean_n_sig = float_cast(mean_n_sig, 'The mean_n_sig argument must be castable to type float!') @@ -1278,17 +1307,20 @@ def create_trial_data_file( mean_n_sig_max = mean_n_sig mean_n_sig_step = 1 else: - mean_n_sig = float_cast( + _mean_n_sig_list: list[float] = float_cast( # pyright: ignore[reportAssignmentType] mean_n_sig, 'The sequence elements of the mean_n_sig argument must be castable to float values!' ) - if len(mean_n_sig) == 2: - (mean_n_sig_min, mean_n_sig_max) = mean_n_sig + if len(_mean_n_sig_list) == 2: + (mean_n_sig_min, mean_n_sig_max) = _mean_n_sig_list mean_n_sig_step = 1 - elif len(mean_n_sig) == 3: - (mean_n_sig_min, mean_n_sig_max, mean_n_sig_step) = mean_n_sig + elif len(_mean_n_sig_list) == 3: + (mean_n_sig_min, mean_n_sig_max, mean_n_sig_step) = _mean_n_sig_list mean_n_sig = np.arange(mean_n_sig_min, mean_n_sig_max + 1, mean_n_sig_step, dtype=np.float64) + mean_n_sig_null_min: float = 0 + mean_n_sig_null_max: float = 0 + mean_n_sig_null_step: float = 1 if not isinstance(mean_n_sig_null, np.ndarray): if not issequence(mean_n_sig_null): mean_n_sig_null = float_cast( @@ -1298,15 +1330,15 @@ def create_trial_data_file( mean_n_sig_null_max = mean_n_sig_null mean_n_sig_null_step = 1 else: - mean_n_sig_null = float_cast( + _mean_n_sig_null_list: list[float] = float_cast( # pyright: ignore[reportAssignmentType] mean_n_sig_null, 'The sequence elements of the mean_n_sig_null argument must be castable to float values!', ) - if len(mean_n_sig_null) == 2: - (mean_n_sig_null_min, mean_n_sig_null_max) = mean_n_sig_null + if len(_mean_n_sig_null_list) == 2: + (mean_n_sig_null_min, mean_n_sig_null_max) = _mean_n_sig_null_list mean_n_sig_null_step = 1 - elif len(mean_n_sig_null) == 3: - (mean_n_sig_null_min, mean_n_sig_null_max, mean_n_sig_null_step) = mean_n_sig_null + elif len(_mean_n_sig_null_list) == 3: + (mean_n_sig_null_min, mean_n_sig_null_max, mean_n_sig_null_step) = _mean_n_sig_null_list mean_n_sig_null = np.arange( mean_n_sig_null_min, mean_n_sig_null_max + 1, mean_n_sig_null_step, dtype=np.float64 @@ -1348,20 +1380,21 @@ def create_trial_data_file( # Save the trial data to file. np.save(pathfilename, trial_data) + assert rss.seed is not None return (rss.seed, mean_n_sig, mean_n_sig_null, trial_data) def extend_trial_data_file( - ana, - rss, - n_trials, - trial_data, + ana: Analysis, + rss: RandomStateService, + n_trials: int, + trial_data: np.ndarray, mean_n_sig=0, mean_n_sig_null=0, mean_n_bkg_list=None, - bkg_kwargs=None, - sig_kwargs=None, - pathfilename=None, + bkg_kwargs: dict | None = None, + sig_kwargs: dict | None = None, + pathfilename: str | None = None, **kwargs, ): """Appends to the trial data file `n_trials` generated trials for each @@ -1370,16 +1403,16 @@ def extend_trial_data_file( Parameters ---------- - ana : instance of Analysis + ana The Analysis instance to use for sensitivity estimation. - rss : instance of RandomStateService + rss The RandomStateService instance to use for generating random numbers. - n_trials : int + n_trials The number of trials the trial data file needs to be extended by. - trial_data : structured numpy ndarray + trial_data The structured numpy ndarray holding the trials. - mean_n_sig : ndarray of float | float | 2- or 3-element sequence of float + mean_n_sig The array of mean number of injected signal events (MNOISEs) for which to generate trials. If this argument is not a ndarray, an array of MNOISEs is generated based on this argument. @@ -1388,7 +1421,7 @@ def extend_trial_data_file( MNOISEs with a step size of one. If a 3-element sequence of floats is given, it specifies the range plus the step size of the MNOISEs. - mean_n_sig_null : ndarray of float | float | 2- or 3-element sequence of float + mean_n_sig_null The array of the fixed mean number of signal events (FMNOSEs) for the null-hypothesis for which to generate trials. If this argument is not a ndarray, an array of FMNOSEs is generated based on this argument. @@ -1397,15 +1430,15 @@ def extend_trial_data_file( FMNOSEs with a step size of one. If a 3-element sequence of floats is given, it specifies the range plus the step size of the FMNOSEs. - bkg_kwargs : dict | None + bkg_kwargs Additional keyword arguments for the `generate_events` method of the background generation method class. An usual keyword argument is `poisson`. - sig_kwargs : dict | None + sig_kwargs Additional keyword arguments for the `generate_signal_events` method of the `SignalGenerator` class. An usual keyword argument is `poisson`. - pathfilename : string | None + pathfilename Trial data file path including the filename. Additional keyword arguments @@ -1448,7 +1481,9 @@ def extend_trial_data_file( return trial_data -def calculate_upper_limit_distribution(ana, rss, pathfilename, n_bkg=5000, n_bins=100): +def calculate_upper_limit_distribution( + ana: Analysis, rss: RandomStateService, pathfilename: str, n_bkg: int = 5000, n_bins: int = 100 +) -> dict: """Function to calculate upper limit distribution. It loads the trial data file containing test statistic distribution and calculates 10 percentile value for each mean number of injected signal event. Then it finds upper @@ -1457,36 +1492,36 @@ def calculate_upper_limit_distribution(ana, rss, pathfilename, n_bkg=5000, n_bin Parameters ---------- - ana : instance of Analysis + ana The Analysis instance to use for sensitivity estimation. - rss : instance of RandomStateService + rss The RandomStateService instance to use for generating random numbers. - pathfilename : string + pathfilename Trial data file path including the filename. - n_bkg : int, optional + n_bkg Number of times to perform background analysis trial. - n_bins : int, optional + n_bins Number of returned test statistic histograms bins. Returns ------- - result : dict + result Result dictionary which contains the following fields: - ul : list of float + ul List of upper limit values. - mean : float + mean Mean of upper limit values. - median : float + median Median of upper limit values. - var : float + var Variance of upper limit values. - ts_hist : numpy ndarray + ts_hist 2D array of test statistic histograms calculated by axis 1. - extent : list of float + extent Test statistic histogram boundaries. - q_values : list of float + q_values `q` percentile values of test statistic for different injected events means. """ @@ -1511,7 +1546,7 @@ def calculate_upper_limit_distribution(ana, rss, pathfilename, n_bkg=5000, n_bin # `ts_inv_f` interpolation boundary. ts_bkg = ts_bkg[ts_bkg >= min(trial_data_q_values)] - ul_list = map(ts_inv_f, ts_bkg) + ul_list = list(map(ts_inv_f, ts_bkg)) ul_mean = np.mean(ul_list) ul_median = np.median(ul_list) ul_var = np.var(ul_list) diff --git a/skyllh/core/utils/coords.py b/skyllh/core/utils/coords.py index a11d13d84d..0febb1126d 100644 --- a/skyllh/core/utils/coords.py +++ b/skyllh/core/utils/coords.py @@ -1,3 +1,5 @@ +from typing import Any + import numpy as np from astropy.coordinates import ( SkyCoord, @@ -80,44 +82,44 @@ def rotate_spherical_vector(ra1, dec1, ra2, dec2, ra3, dec3): def rotate_signal_events_on_sphere( - src_ra, - src_dec, - evt_true_ra, - evt_true_dec, - evt_reco_ra, - evt_reco_dec, -): + src_ra: np.ndarray, + src_dec: np.ndarray, + evt_true_ra: np.ndarray, + evt_true_dec: np.ndarray, + evt_reco_ra: np.ndarray, + evt_reco_dec: np.ndarray, +) -> tuple[np.ndarray, np.ndarray]: """Rotate signal events on a sphere to a given source position preserving position angle and separation (great circle distance) between the event's true and reco directions. Parameters ---------- - src_ra : instance of numpy.ndarray + src_ra The (N_events,)-shaped 1D numpy.ndarray holding the true right-ascension of the source. - src_dec : instance of numpy.ndarray + src_dec The (N_events,)-shaped 1D numpy.ndarray holding the true declination of the source. - evt_true_ra : instance of numpy.ndarray + evt_true_ra The (N_events,)-shaped 1D numpy.ndarray holding the true right-ascension of the MC event. - evt_true_dec : instance of numpy.ndarray + evt_true_dec The (N_events,)-shaped 1D numpy.ndarray holding the true declination of the MC event. - evt_reco_ra : instance of numpy.ndarray + evt_reco_ra The (N_events,)-shaped 1D numpy.ndarray holding the reconstructed right-ascension of the MC event. - evt_reco_dec : instance of numpy.ndarray + evt_reco_dec The (N_events,)-shaped 1D numpy.ndarray holding the reconstructed declination of the MC event. Returns ------- - rot_evt_reco_ra : instance of numpy.ndarray + rot_evt_reco_ra The (N_events,)-shaped 1D numpy.ndarray holding the rotated reconstructed event right-ascension. - rot_evt_reco_dec : instance of numpy.ndarray + rot_evt_reco_dec The (N_events,)-shaped 1D numpy.ndarray holding the rotated reconstructed event declination. """ @@ -132,36 +134,39 @@ def rotate_signal_events_on_sphere( position_angle = v_evt_true.position_angle(v_evt_reco) separation = v_evt_true.separation(v_evt_reco) - v_rotated = v_source.directional_offset_by(position_angle, separation) - (rot_evt_reco_ra, rot_evt_reco_dec) = (v_rotated.ra.rad, v_rotated.dec.rad) + v_rotated: Any = v_source.directional_offset_by(position_angle, separation) + rot_evt_reco_ra: np.ndarray = v_rotated.ra.rad + rot_evt_reco_dec: np.ndarray = v_rotated.dec.rad return (rot_evt_reco_ra, rot_evt_reco_dec) -def angular_separation(ra1, dec1, ra2, dec2, psi_floor=None): +def angular_separation( + ra1: np.ndarray, dec1: np.ndarray, ra2: np.ndarray, dec2: np.ndarray, psi_floor: float | None = None +) -> np.ndarray: """Calculates the angular separation on the sphere between two vectors on the sphere. Parameters ---------- - ra1 : instance of numpy.ndarray + ra1 The (N_events,)-shaped numpy.ndarray holding the right-ascension or longitude coordinate of the first vector in radians. - dec1 : instance of numpy.ndarray + dec1 The (N_events,)-shaped numpy.ndarray holding declination or latitude coordinate of the first vector in radians. - ra2 : instance of numpy.ndarray + ra2 The (N_events,)-shaped numpy.ndarray holding the right-ascension or longitude coordinate of the second vector in radians. - dec2 : instance of numpy.ndarray + dec2 The (N_events,)-shaped numpy.ndarray holding declination coordinate of the second vector in radians. - psi_floor : float | None + psi_floor If not ``None``, specifies the floor value of psi. Returns ------- - psi : instance of numpy.ndarray + psi The (N_events,)-shaped numpy.ndarray holding the calculated angular separation value of each event. """ diff --git a/skyllh/core/utils/flux_model.py b/skyllh/core/utils/flux_model.py index c4de16a6ec..3fdbea3907 100644 --- a/skyllh/core/utils/flux_model.py +++ b/skyllh/core/utils/flux_model.py @@ -11,7 +11,7 @@ def create_scipy_stats_rv_continuous_from_TimeFluxProfile( - profile, + profile: TimeFluxProfile, ): """This function builds a scipy.stats.rv_continuous instance for a given :class:`~skyllh.core.flux_model.TimeFluxProfile` instance. @@ -21,13 +21,13 @@ def create_scipy_stats_rv_continuous_from_TimeFluxProfile( Parameters ---------- - profile : instance of TimeFluxProfile + profile The instance of TimeFluxProfile providing the function of the time flux profile. Returns ------- - rv : instance of rv_continuous_frozen + rv The instance of rv_continuous_frozen representing the time flux profile as a continuous random variate instance. """ @@ -42,6 +42,10 @@ def create_scipy_stats_rv_continuous_from_TimeFluxProfile( norm = 1 / tot_integral class rv_continuous_from_TimeFluxProfile(rv_continuous): + """This class provides a scipy ``rv_continuous`` random variable whose + probability density is given by a time flux profile instance. + """ + def __init__(self, *args, **kwargs): """Creates a new instance of the subclass of rv_continuous using the time flux profile. @@ -51,24 +55,25 @@ def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) - def _pdf(self, t): + def _pdf(self, x, *args): """Calculates the probability density of the time flux profile function for given time values. """ - pd = self._profile(t=t) * self._norm + pd = self._profile(t=x) * self._norm return pd - def _cdf(self, t): + def _cdf(self, x, *args): """Calculates the cumulative distribution function values for rhe given time values. If the time flux profile instance provides a ``cdf`` method, it will be used. Otherwise the generic ``_cdf`` method of the ``rv_continuous`` class will be used. """ - if hasattr(self._profile, 'cdf') and callable(self._profile.cdf): - return self._profile.cdf(t=t) + profile_cdf = getattr(self._profile, 'cdf', None) + if callable(profile_cdf): + return profile_cdf(t=x) - return super()._cdf(t) + return super()._cdf(x) rv = rv_continuous_from_TimeFluxProfile( a=profile.t_start, diff --git a/skyllh/core/utils/multidimgridpdf.py b/skyllh/core/utils/multidimgridpdf.py index 4a56ccebe7..aa125c7d7a 100644 --- a/skyllh/core/utils/multidimgridpdf.py +++ b/skyllh/core/utils/multidimgridpdf.py @@ -2,20 +2,25 @@ MultiDimGridPDF instances. """ +from collections.abc import Callable + import numpy as np from skyllh.core.binning import ( BinningDefinition, ) +from skyllh.core.dataset import Dataset, DatasetData +from skyllh.core.parameters import ParameterModelMapper from skyllh.core.pdf import ( MultiDimGridPDF, ) from skyllh.core.py import ( classname, ) +from skyllh.core.timing import TimeLord -def get_kde_pdf_sig_spatial_norm_factor_func(log10_psi_name='log10_psi'): +def get_kde_pdf_sig_spatial_norm_factor_func(log10_psi_name: str = 'log10_psi'): """Returns the standard normalization factor function for the spatial signal MultiDimGridPDF, which is created from KDE PDF values. It can be used for the ``norm_factor_func`` argument of the @@ -24,12 +29,14 @@ def get_kde_pdf_sig_spatial_norm_factor_func(log10_psi_name='log10_psi'): Parameters ---------- - log10_psi_name : str + log10_psi_name The name of the event data field for the log10(psi) values. """ def kde_pdf_sig_spatial_norm_factor_func(pdf, tdm, params_recarray, eventdata, evt_mask=None): - + """Calculates the normalization factor for the signal spatial KDE PDF for + each event from its log10(psi) value. + """ log10_psi_idx = pdf._axes.get_index_by_name(log10_psi_name) if evt_mask is None: # noqa: SIM108 @@ -53,23 +60,25 @@ def get_kde_pdf_bkg_norm_factor_func(): """ def kde_pdf_bkg_norm_factor_func(pdf, tdm, params_recarray, eventdata, evt_mask=None): - + """Returns the constant normalization factor for the background KDE + PDF. + """ return 1.0 / (2 * np.pi) return kde_pdf_bkg_norm_factor_func def create_MultiDimGridPDF_from_photosplinetable( - multidimgridpdf_cls, - pmm, - ds, - data, - info_key, - splinetable_key, - kde_pdf_axis_name_map_key='KDE_PDF_axis_name_map', - norm_factor_func=None, - cache_pd_values=False, - tl=None, + multidimgridpdf_cls: type[MultiDimGridPDF], + pmm: ParameterModelMapper, + ds: Dataset, + data: DatasetData, + info_key: str, + splinetable_key: str, + kde_pdf_axis_name_map_key: str = 'KDE_PDF_axis_name_map', + norm_factor_func: Callable | None = None, + cache_pd_values: bool = False, + tl: TimeLord | None = None, **kwargs, ): """ @@ -79,38 +88,38 @@ def create_MultiDimGridPDF_from_photosplinetable( Parameters ---------- - multidimgridpdf_cls : subclass of MultiDimGridPDF + multidimgridpdf_cls The MultiDimGridPDF class, which should be used. - pmm : instance of ParameterModelMapper + pmm The instance of ParameterModelMapper, which defines the mapping of global parameters to local model parameters. - ds : instance of Dataset + ds The instance of Dataset the PDF applies to. - data : instance of DatasetData + data The instance of DatasetData that holds the experimental and monte-carlo data of the dataset. - info_key : str + info_key The auxiliary data name for the file containing PDF information. - splinetable_key : str + splinetable_key The auxiliary data name for the name of the file containing the photospline spline table. - kde_pdf_axis_name_map_key : str + kde_pdf_axis_name_map_key The auxiliary data name for the KDE PDF axis name map. - norm_factor_func : callable | None + norm_factor_func The function that calculates a possible required normalization factor for the PDF value based on the event properties. For more information about this argument see the documentation of the :meth:`skyllh.core.pdf.MultiDimGridPDF.__init__` method. - cache_pd_values : bool + cache_pd_values Flag if the probability density values should get cached by the MultiDimGridPDF class. - tl : instance of TimeLord | None + tl The optional instance of TimeLord to use for measuring timing information. Returns ------- - pdf : instance of ``multidimgridpdf_cls`` + pdf The created PDF instance of MultiDimGridPDF. """ if not issubclass(multidimgridpdf_cls, MultiDimGridPDF): @@ -157,16 +166,16 @@ def create_MultiDimGridPDF_from_photosplinetable( def create_MultiDimGridPDF_from_kde_pdf( - multidimgridpdf_cls, - pmm, - ds, - data, - numerator_key, - denumerator_key=None, - kde_pdf_axis_name_map_key='KDE_PDF_axis_name_map', - norm_factor_func=None, - cache_pd_values=False, - tl=None, + multidimgridpdf_cls: type[MultiDimGridPDF], + pmm: ParameterModelMapper, + ds: Dataset, + data: DatasetData, + numerator_key: str, + denumerator_key: str | None = None, + kde_pdf_axis_name_map_key: str = 'KDE_PDF_axis_name_map', + norm_factor_func: Callable | None = None, + cache_pd_values: bool = False, + tl: TimeLord | None = None, **kwargs, ): """Creates a MultiDimGridPDF instance with pdf values taken from KDE PDF @@ -174,38 +183,38 @@ def create_MultiDimGridPDF_from_kde_pdf( Parameters ---------- - multidimgridpdf_cls : subclass of MultiDimGridPDF + multidimgridpdf_cls The MultiDimGridPDF class, which should be used. - pmm : instance of ParameterModelMapper + pmm The instance of ParameterModelMapper, which defines the mapping of global parameters to local model parameters. - ds : instance of Dataset + ds The instance of Dataset the PDF applies to. - data : instance of DatasetData + data The instance of DatasetData that holds the auxiliary data of the dataset. - numerator_key : str + numerator_key The auxiliary data name for the PDF numerator array. - denumerator_key : str | None + denumerator_key The auxiliary data name for the PDF denumerator array. This can be ``None``, if no denumerator array is required. - kde_pdf_axis_name_map_key : str + kde_pdf_axis_name_map_key The auxiliary data name for the KDE PDF axis name map. - norm_factor_func : callable | None + norm_factor_func The function that calculates a possible required normalization factor for the PDF value based on the event properties. For more information about this argument see the documentation of the :meth:`skyllh.core.pdf.MultiDimGridPDF.__init__` method. - cache_pd_values : bool + cache_pd_values Flag if the probability density values should get cached by the MultiDimGridPDF class. - tl : instance of TimeLord | None + tl The optional instance of TimeLord to use for measuring timing information. Returns ------- - pdf : instance of ``multidimgridpdf_cls`` + pdf The created PDF instance of MultiDimGridPDF. """ if not issubclass(multidimgridpdf_cls, MultiDimGridPDF): diff --git a/skyllh/core/utils/spline.py b/skyllh/core/utils/spline.py index 901cf138a6..bc76dd891c 100644 --- a/skyllh/core/utils/spline.py +++ b/skyllh/core/utils/spline.py @@ -2,17 +2,17 @@ from scipy.interpolate import interp1d -def make_spline_1d(x, y, kind='linear', **kwargs): +def make_spline_1d(x: np.ndarray, y: np.ndarray, kind: str = 'linear', **kwargs): """Creates a 1D spline for the function y(x) using :class:`scipy.interpolate.interp1d`. Parameters ---------- - x : array_like + x The x values. - y : array_like + y The y values. - kind : str + kind The kind of the spline. See the :class:`scipy.interpolate.interp1d` documentation for possible values. Default is ``'linear'``. **kwargs @@ -52,17 +52,17 @@ class CatmullRomRegular1DSpline: def __init__( self, - x, - y, + x: np.ndarray, + y: np.ndarray, **kwargs, ): """Creates a new CatmullRom1DSpline instance. Parameters ---------- - x : instance of ndarray + x The x values of the data points. - y : instance of ndarray + y The y values of the data points. """ super().__init__(**kwargs) @@ -93,19 +93,19 @@ def __init__( def _eval_for_valid_x( self, - x, - ): + x: np.ndarray, + ) -> np.ndarray: """Evaluates the spline given valid x-values in data coordinates. Parameters ---------- - x : instance of ndarray + x The instance of ndarray holding the valid values for which the spline should get evaluated. Returns ------- - y : instance of ndarray + y The instance of ndarray with the spline values at the given x values. """ @@ -159,20 +159,20 @@ def _eval_for_valid_x( def __call__( self, - x, + x: np.ndarray, oor_value=np.nan, - ): + ) -> np.ndarray: """Evaluates the spline given x-values in data coordinates. Parameters ---------- - x : instance of ndarray + x The instance of ndarray holding the values for which the spline should get evaluated. Returns ------- - y : instance of ndarray + y The instance of ndarray with the spline values at the given x values. """ @@ -187,32 +187,32 @@ def __call__( def _calc_tj( self, - ti, - Pi_x, - Pi_y, - Pj_x, - Pj_y, - ): + ti: float, + Pi_x: float, + Pi_y: float, + Pj_x: float, + Pj_y: float, + ) -> float: """Calculates the next segment coefficient ``tj`` given the previous segment coefficient ``ti`` and the previous and next data point ``(Pi_x, Pi_y)`` and ``(Pj_x, Pj_y)``, respectively. Parameters ---------- - ti : float + ti The previous segment coefficient. - Pi_x : float + Pi_x The x-value of the previous data point. - Pi_y : float + Pi_y The y-value of the previous data point. - Pj_x : float + Pj_x The x-value of the next data point. - Pj_y : float + Pj_y The y-value of the next data point. Returns ------- - tj : float + tj The next segment coefficient. """ dx = Pj_x - Pi_x @@ -223,8 +223,8 @@ def _calc_tj( def _calc_segment_coefficients( self, - Px, - Py, + Px: np.ndarray, + Py: np.ndarray, ): """Calculates the segment coefficients t1, t2, and t3 given the 4 data (control) points of the segment. The coefficient t0 is 0 by @@ -232,10 +232,10 @@ def _calc_segment_coefficients( Parameters ---------- - Px : instance of ndarray + Px The (4,)-shaped numpy ndarray holding the 4 x-values of the segment's data points. - Py : instance of ndarray + Py The (4,)-shaped numpy ndarray holding the 4 y-values of the segment's data points. """ diff --git a/skyllh/core/utils/tdm.py b/skyllh/core/utils/tdm.py index 8db01264fd..4507723332 100644 --- a/skyllh/core/utils/tdm.py +++ b/skyllh/core/utils/tdm.py @@ -1,3 +1,5 @@ +from collections.abc import Callable + import numpy as np from skyllh.core.utils.coords import ( @@ -5,20 +7,20 @@ ) -def get_tdm_field_func_psi(psi_floor=None): +def get_tdm_field_func_psi(psi_floor: float | None = None) -> Callable: """Returns the TrialDataManager (TDM) field function for psi with an optional psi value floor. Parameters ---------- - psi_floor : float | None + psi_floor The optional floor value for psi. This should be ``None`` for a standard point-source analysis that uses an analytic function for the detector's point-spread-function (PSF). Returns ------- - tdm_field_func_psi : function + tdm_field_func_psi TrialDataManager (TDM) field function for psi. """ diff --git a/skyllh/core/utils/trials.py b/skyllh/core/utils/trials.py index bb9d8d0bf1..acc1953292 100644 --- a/skyllh/core/utils/trials.py +++ b/skyllh/core/utils/trials.py @@ -2,48 +2,58 @@ import pickle -from skyllh.core.timing import ( - TaskTimer, -) +from skyllh.core.analysis import Analysis +from skyllh.core.random import RandomStateService +from skyllh.core.storage import DataFieldRecordArray +from skyllh.core.timing import TaskTimer, TimeLord def create_pseudo_data_file( - ana, rss, filename, mean_n_bkg_list=None, mean_n_sig=0, bkg_kwargs=None, sig_kwargs=None, tl=None + ana: Analysis, + rss: RandomStateService, + filename: str, + mean_n_bkg_list: list[float] | None = None, + mean_n_sig: float = 0, + bkg_kwargs: dict | None = None, + sig_kwargs: dict | None = None, + tl: TimeLord | None = None, ): """Creates a pickle file that contains the pseudo data for a single trial by generating background and signal events. Parameters ---------- - ana : Analysis + ana The Analysis instance that should be used to generate the pseudo data. - rss : RandomStateService + rss The RandomStateService instance to use for generating random numbers. - filename : str + filename The data file name into which the generated pseudo data should get written to. - mean_n_bkg_list : list of float | None + mean_n_bkg_list The mean number of background events that should be generated for each dataset. If set to None (the default), the background generation method needs to obtain this number itself. - mean_n_sig : float + mean_n_sig The mean number of signal events that should be generated for the trial. The actual number of generated events will be drawn from a Poisson distribution with this given signal mean as mean. - bkg_kwargs : dict | None + bkg_kwargs Additional keyword arguments for the `generate_events` method of the background generation method class. An usual keyword argument is `poisson`. - sig_kwargs : dict | None + sig_kwargs Additional keyword arguments for the `generate_signal_events` method of the `SignalGenerator` class. An usual keyword argument is `poisson`. - tl : TimeLord | None + tl The instance of TimeLord that should be used to time individual tasks. """ + from typing import cast as _cast + (n_bkg_events_list, bkg_events_list) = ana.generate_background_events( - rss=rss, mean_n_bkg_list=mean_n_bkg_list, bkg_kwargs=bkg_kwargs, tl=tl + rss=rss, mean_n_bkg_list=_cast('list[float | None] | None', mean_n_bkg_list), bkg_kwargs=bkg_kwargs, tl=tl ) (n_sig, n_sig_events_list, sig_events_list) = ana.generate_signal_events( @@ -66,32 +76,34 @@ def create_pseudo_data_file( pickle.dump(trial_data, fp) -def load_pseudo_data(filename, tl=None): +def load_pseudo_data( + filename: str, tl: TimeLord | None = None +) -> tuple[float, int, list[int], list[int], list[DataFieldRecordArray], list[DataFieldRecordArray]]: """Loads the pseudo data for a single trial from the given file name. Parameters ---------- - filename : str + filename The name of the file that contains the pseudo data. - tl : TimeLord | None + tl The instance of TimeLord that should be used to time individual tasks. Returns ------- - mean_n_sig : float + mean_n_sig The mean number of signal events that was used to generate the pseudo data. - n_sig : int + n_sig The actual total number of signal events in the pseudo data. - n_bkg_events_list : list of int + n_bkg_events_list The total number of background events for each data set of the pseudo data. - n_sig_events_list : list of int + n_sig_events_list The total number of signal events for each data set of the pseudo data. - bkg_events_list : list of DataFieldRecordArray instances + bkg_events_list The list of DataFieldRecordArray instances containing the background pseudo data events for each data set. - sig_events_list : list of DataFieldRecordArray instances | None + sig_events_list The list of DataFieldRecordArray instances containing the signal pseudo data events for each data set. If a particular dataset has no signal events, the entry for that dataset can be None. diff --git a/skyllh/datasets/i3/PublicData_10y_ps.py b/skyllh/datasets/i3/PublicData_10y_ps.py index bd062814e4..d63e9ed10b 100644 --- a/skyllh/datasets/i3/PublicData_10y_ps.py +++ b/skyllh/datasets/i3/PublicData_10y_ps.py @@ -1,5 +1,8 @@ +from typing import TypedDict + import numpy as np +from skyllh.core.config import Config from skyllh.core.dataset import ( DatasetCollection, DatasetOrigin, @@ -13,11 +16,25 @@ DATASET_NAMES = ('IC40', 'IC59', 'IC79', 'IC86_I', 'IC86_II-VII') +class _DsKwargs(TypedDict): + """Typed dictionary of the keyword arguments passed to each dataset + definition in this collection. + """ + + cfg: Config + version: int + verqualifiers: dict[str, int] | None + base_path: str | None + default_sub_path_fmt: str + sub_path_fmt: str | None + origin: DatasetOrigin + + def create_dataset_collection( - cfg, - base_path=None, - sub_path_fmt=None, -): + cfg: Config, + base_path: str | None = None, + sub_path_fmt: str | None = None, +) -> DatasetCollection: """Defines the dataset collection for IceCube's 10-year point-source public data, which is available at https://doi.org/10.7910/DVN/VKL316. @@ -25,21 +42,21 @@ def create_dataset_collection( Parameters ---------- - cfg : instance of Config + cfg The instance of Config holding the local configuration. - base_path : str | None + base_path The base path of the data files. The actual path of a data file is assumed to be of the structure //. If ``None``, ``cfg['repository']['base_path']`` is used, which defaults to ``~/.cache/skyllh``. - sub_path_fmt : str | None + sub_path_fmt The sub path format of the data files of the public data sample. If None, use the default sub path format 'icecube_10year_ps'. Returns ------- - dsc : DatasetCollection + dsc The dataset collection containing all the seasons as individual I3Dataset objects. """ @@ -270,9 +287,8 @@ def create_dataset_collection( ) # Define the common keyword arguments for all data sets. - ds_kwargs = { + ds_kwargs: _DsKwargs = { 'cfg': cfg, - 'livetime': None, 'version': version, 'verqualifiers': verqualifiers, 'base_path': base_path, @@ -580,6 +596,9 @@ def create_dataset_collection( ) def convert_deg2rad(data): + """Converts the angular experimental data fields from degrees to + radians. + """ exp = data.exp exp['ang_err'] = np.deg2rad(exp['ang_err']) exp['ra'] = np.deg2rad(exp['ra']) diff --git a/skyllh/datasets/i3/PublicData_10y_ps_wMC.py b/skyllh/datasets/i3/PublicData_10y_ps_wMC.py index fb45c5712f..b59bb4fdd6 100644 --- a/skyllh/datasets/i3/PublicData_10y_ps_wMC.py +++ b/skyllh/datasets/i3/PublicData_10y_ps_wMC.py @@ -1,5 +1,7 @@ import numpy as np +from skyllh.core.config import Config +from skyllh.core.dataset import DatasetCollection from skyllh.datasets.i3 import ( PublicData_10y_ps, ) @@ -8,10 +10,10 @@ def create_dataset_collection( - cfg, - base_path=None, - sub_path_fmt=None, -): + cfg: Config, + base_path: str | None = None, + sub_path_fmt: str | None = None, +) -> DatasetCollection: """Defines the dataset collection for IceCube's 10-year point-source public data, which is available at http://icecube.wisc.edu/data-releases/20210126_PS-IC40-IC86_VII.zip, and @@ -19,21 +21,21 @@ def create_dataset_collection( Parameters ---------- - cfg : instance of Config + cfg The instance of Config holding the local configuration. - base_path : str | None + base_path The base path of the data files. The actual path of a data file is assumed to be of the structure //. If ``None``, ``cfg['repository']['base_path']`` is used, which defaults to ``~/.cache/skyllh``. - sub_path_fmt : str | None + sub_path_fmt The sub path format of the data files of the public data sample. If None, use the default sub path format 'icecube_10year_ps'. Returns ------- - dsc : DatasetCollection + dsc The dataset collection containing all the seasons as individual I3Dataset objects. """ @@ -83,10 +85,14 @@ def create_dataset_collection( IC86_II_VII.mc_pathfilename_list = IC86_II.mc_pathfilename_list def add_time(data): + """Adds a zero-valued ``time`` data field to the monte-carlo data.""" mc = data.mc mc.append_field('time', np.repeat(0, len(mc))) def add_azimuth_and_zenith(data): + """Adds zero-valued ``azi`` and ``zen`` data fields to the monte-carlo + data. + """ mc = data.mc mc.append_field('azi', np.repeat(0, len(mc))) mc.append_field('zen', np.repeat(0, len(mc))) diff --git a/skyllh/datasets/i3/PublicData_14y_ps.py b/skyllh/datasets/i3/PublicData_14y_ps.py index 6456033d9a..ab50a82500 100644 --- a/skyllh/datasets/i3/PublicData_14y_ps.py +++ b/skyllh/datasets/i3/PublicData_14y_ps.py @@ -1,5 +1,8 @@ +from typing import TypedDict + import numpy as np +from skyllh.core.config import Config from skyllh.core.dataset import ( DatasetCollection, DatasetOrigin, @@ -13,11 +16,25 @@ DATASET_NAMES = ('IC40', 'IC59', 'IC79', 'IC86_I-XI') +class _DsKwargs(TypedDict): + """Typed dictionary of the keyword arguments passed to each dataset + definition in this collection. + """ + + cfg: Config + version: int + verqualifiers: dict[str, int] | None + base_path: str | None + default_sub_path_fmt: str + sub_path_fmt: str | None + origin: DatasetOrigin + + def create_dataset_collection( - cfg, - base_path=None, - sub_path_fmt=None, -): + cfg: Config, + base_path: str | None = None, + sub_path_fmt: str | None = None, +) -> DatasetCollection: """Defines the dataset collection for IceCube's 14-year point-source public data, which is available at https://doi.org/10.7910/DVN/MMIIZA. @@ -31,20 +48,20 @@ def create_dataset_collection( Parameters ---------- - cfg : instance of Config + cfg The instance of Config holding the local configuration. - base_path : str | None + base_path The base path of the data files. The actual path of a data file is assumed to be of the structure //. If ``None``, ``cfg['repository']['base_path']`` is used, which defaults to ``~/.cache/skyllh``. - sub_path_fmt : str | None + sub_path_fmt The sub path format of the data files of the public data sample. If None, use the default sub path format 'icecube_14year_ps'. Returns ------- - dsc : DatasetCollection + dsc The dataset collection containing all the seasons as individual I3Dataset objects. """ @@ -238,9 +255,8 @@ def create_dataset_collection( ) # Define the common keyword arguments for all data sets. - ds_kwargs = { + ds_kwargs: _DsKwargs = { 'cfg': cfg, - 'livetime': None, 'version': version, 'verqualifiers': verqualifiers, 'base_path': base_path, @@ -616,6 +632,9 @@ def create_dataset_collection( ) def convert_deg2rad(data): + """Converts the angular experimental data fields from degrees to + radians. + """ exp = data.exp exp['ang_err'] = np.deg2rad(exp['ang_err']) exp['ra'] = np.deg2rad(exp['ra']) diff --git a/skyllh/datasets/i3/TestData.py b/skyllh/datasets/i3/TestData.py index b9bd3118f3..d1d2972cc7 100644 --- a/skyllh/datasets/i3/TestData.py +++ b/skyllh/datasets/i3/TestData.py @@ -1,3 +1,6 @@ +from typing import TypedDict + +from skyllh.core.config import Config from skyllh.core.dataset import ( DatasetCollection, ) @@ -8,30 +11,43 @@ DATASET_NAMES = ('TestData',) +class _DsKwargs(TypedDict): + """Typed dictionary of the keyword arguments passed to each dataset + definition in this collection. + """ + + cfg: Config + version: int + verqualifiers: dict[str, int] | None + base_path: str | None + default_sub_path_fmt: str + sub_path_fmt: str | None + + def create_dataset_collection( - cfg, - base_path=None, - sub_path_fmt=None, -): + cfg: Config, + base_path: str | None = None, + sub_path_fmt: str | None = None, +) -> DatasetCollection: """Defines a dataset collection with a test dataset. Parameters ---------- - cfg : instance of Config + cfg The instance of Config holding the local configuration. - base_path : str | None + base_path The base path of the data files. The actual path of a data file is assumed to be of the structure //. If ``None``, ``cfg['repository']['base_path']`` is used, which defaults to ``~/.cache/skyllh``. - sub_path_fmt : str | None + sub_path_fmt The sub path format of the data files of the public data sample. If None, use the default sub path format 'testdata'. Returns ------- - dsc : DatasetCollection + dsc The dataset collection containing all the seasons as individual I3Dataset objects. """ @@ -47,9 +63,8 @@ def create_dataset_collection( """ # Define the common keyword arguments for all data sets. - ds_kwargs = { + ds_kwargs: _DsKwargs = { 'cfg': cfg, - 'livetime': None, 'version': version, 'verqualifiers': verqualifiers, 'base_path': base_path, diff --git a/skyllh/i3/background_generation.py b/skyllh/i3/background_generation.py index 5c113499e9..5b49ab8024 100644 --- a/skyllh/i3/background_generation.py +++ b/skyllh/i3/background_generation.py @@ -1,12 +1,15 @@ from skyllh.core.background_generation import ( BackgroundGenerationMethod, ) +from skyllh.core.dataset import Dataset, DatasetData from skyllh.core.py import ( classname, ) +from skyllh.core.random import RandomStateService from skyllh.core.scrambling import ( DataScrambler, ) +from skyllh.core.storage import DataFieldRecordArray class FixedScrambledExpDataI3BkgGenMethod( @@ -20,7 +23,7 @@ class FixedScrambledExpDataI3BkgGenMethod( def __init__( self, - data_scrambler, + data_scrambler: DataScrambler, **kwargs, ): """Creates a new background generation method instance to generate @@ -29,7 +32,7 @@ def __init__( Parameters ---------- - data_scrambler : instance of DataScrambler + data_scrambler The DataScrambler instance to use to generate scrambled experimental data. """ @@ -52,37 +55,38 @@ def data_scrambler(self, scrambler): ) self._data_scrambler = scrambler - def generate_events( + def generate_events( # pyright: ignore[reportIncompatibleMethodOverride] self, - rss, - dataset, - data, + rss: RandomStateService, + dataset: Dataset, + data: DatasetData, **kwargs, - ): + ) -> tuple[int, DataFieldRecordArray]: """Generates background events from the given data, by scrambling the experimental data. The number of events is equal to the size of the given dataset. Parameters ---------- - rss : instance of RandomStateService + rss The instance of RandomStateService that should be used to generate random numbers from. It is used to scramble the experimental data. - dataset : instance of Dataset + dataset The Dataset instance describing the dataset for which background events should get generated. - data : instance of DatasetData + data The DatasetData instance holding the data of the dataset for which background events should get generated. Returns ------- - n_bkg : int + n_bkg The number of generated background events. - bkg_events : instance of DataFieldRecordArray + bkg_events The instance of DataFieldRecordArray holding the generated background events. """ + assert data.exp is not None bkg_events = self._data_scrambler.scramble_data(rss=rss, dataset=dataset, data=data.exp, copy=True) return (len(bkg_events), bkg_events) diff --git a/skyllh/i3/backgroundpdf.py b/skyllh/i3/backgroundpdf.py index d74945271b..d0ef210bd9 100644 --- a/skyllh/i3/backgroundpdf.py +++ b/skyllh/i3/backgroundpdf.py @@ -1,9 +1,7 @@ import numpy as np import scipy.interpolate -from skyllh.core.binning import ( - UsesBinning, -) +from skyllh.core.binning import BinningDefinition, UsesBinning from skyllh.core.pdf import ( IsBackgroundPDF, SpatialPDF, @@ -14,12 +12,12 @@ issequence, issequenceof, ) +from skyllh.core.smoothing import SmoothingFilter from skyllh.core.storage import ( DataFieldRecordArray, ) -from skyllh.core.timing import ( - TaskTimer, -) +from skyllh.core.timing import TaskTimer, TimeLord +from skyllh.core.trialdata import TrialDataManager from skyllh.i3.pdf import ( I3EnergyPDF, ) @@ -40,23 +38,23 @@ class BackgroundI3SpatialPDF( def __init__( self, - data_sin_dec, - data_weights, - sin_dec_binning, - spline_order_sin_dec, + data_sin_dec: np.ndarray, + data_weights: np.ndarray, + sin_dec_binning: BinningDefinition, + spline_order_sin_dec: int, **kwargs, ): """Creates a new IceCube spatial background PDF object. Parameters ---------- - data_sin_dec : 1d ndarray + data_sin_dec The array holding the sin(dec) values of the events. - data_weights : 1d ndarray + data_weights The array holding the weight of each event used for histogramming. - sin_dec_binning : BinningDefinition + sin_dec_binning The binning definition for the sin(declination) axis. - spline_order_sin_dec : int + spline_order_sin_dec The order of the spline function for the logarithmic values of the spatial background PDF along the sin(dec) axis. """ @@ -116,17 +114,17 @@ def spline_order_sin_dec(self): def spline_order_sin_dec(self, order): self._spline_order_sin_dec = int_cast(order, 'The spline_order_sin_dec property must be castable to type int!') - def add_events(self, events): + def add_events(self, events: np.ndarray): """Add events to spatial background PDF object and recalculate logarithmic spline function. Parameters ---------- - events : numpy record ndarray + events The array holding the event data. The following data fields must exist: - sin_dec : float + sin_dec The sin(declination) value of the event. """ @@ -162,31 +160,33 @@ def initialize_for_new_trial(self, tdm, tl=None, **kwargs): self._pd = 0.5 / np.pi * np.exp(log_spline_val) - def get_pd(self, tdm, params_recarray=None, tl=None): + def get_pd( # pyright: ignore[reportIncompatibleMethodOverride] + self, tdm: TrialDataManager, params_recarray: None = None, tl: TimeLord | None = None + ) -> tuple[np.ndarray, dict]: """Calculates the spatial background probability on the sphere of each event. Parameters ---------- - tdm : instance of TrialDataManager + tdm The TrialDataManager instance holding the trial event data for which to calculate the PDF values. The following data fields must exist: - sin_dec : float + sin_dec The sin(declination) value of the event. - params_recarray : None + params_recarray Unused interface parameter. - tl : instance of TimeLord | None + tl The optional TimeLord instance that should be used to measure timing information. Returns ------- - pd : instance of numpy ndarray + pd The (N_events,)-shaped numpy ndarray holding the background probability density value for each event. - grads : dict + grads The dictionary holding the gradients of the probability density w.r.t. each global fit parameter. The background PDF does not depend on any global fit parameter, @@ -204,9 +204,9 @@ class DataBackgroundI3SpatialPDF( def __init__( self, - data_exp, - sin_dec_binning, - spline_order_sin_dec=2, + data_exp: DataFieldRecordArray, + sin_dec_binning: BinningDefinition, + spline_order_sin_dec: int = 2, **kwargs, ): """Constructs a new IceCube spatial background PDF from experimental @@ -214,16 +214,16 @@ def __init__( Parameters ---------- - data_exp : instance of DataFieldRecordArray + data_exp The instance of DataFieldRecordArray holding the experimental data. The following data fields must exist: - sin_dec : float + sin_dec The sin(declination) of the data event. - sin_dec_binning : BinningDefinition + sin_dec_binning The binning definition for the sin(declination). - spline_order_sin_dec : int + spline_order_sin_dec The order of the spline function for the logarithmic values of the spatial background PDF along the sin(dec) axis. The default is 2. @@ -257,10 +257,10 @@ class MCBackgroundI3SpatialPDF( def __init__( self, - data_mc, - physics_weight_field_names, - sin_dec_binning, - spline_order_sin_dec=2, + data_mc: DataFieldRecordArray, + physics_weight_field_names: str | list[str], + sin_dec_binning: BinningDefinition, + spline_order_sin_dec: int = 2, **kwargs, ): """Constructs a new IceCube spatial background PDF from monte-carlo @@ -268,21 +268,21 @@ def __init__( Parameters ---------- - data_mc : instance of DataFieldRecordArray + data_mc The array holding the monte-carlo data. The following data fields must exist: - sin_dec : float + sin_dec The sine of the reconstructed declination of the data event. - physics_weight_field_names : str | list of str + physics_weight_field_names The name or the list of names of the monte-carlo data fields, which should be used as event weights. If a list is given, the weight values of all the fields will be summed to construct the final event weight. - sin_dec_binning : BinningDefinition + sin_dec_binning The binning definition for the sin(declination). - spline_order_sin_dec : int + spline_order_sin_dec The order of the spline function for the logarithmic values of the spatial background PDF along the sin(dec) axis. The default is 2. @@ -293,7 +293,7 @@ def __init__( ) if not issequence(physics_weight_field_names): - physics_weight_field_names = [physics_weight_field_names] + physics_weight_field_names = [physics_weight_field_names] # pyright: ignore[reportAssignmentType] if not issequenceof(physics_weight_field_names, str): raise TypeError( 'The physics_weight_field_names argument must be of type str ' @@ -331,10 +331,10 @@ class DataBackgroundI3EnergyPDF( def __init__( self, - data_exp, - log10_energy_binning, - sin_dec_binning, - smoothing_filter=None, + data_exp: DataFieldRecordArray, + log10_energy_binning: BinningDefinition, + sin_dec_binning: BinningDefinition, + smoothing_filter: SmoothingFilter | None = None, **kwargs, ): """Constructs a new IceCube energy background PDF from experimental @@ -342,21 +342,21 @@ def __init__( Parameters ---------- - data_exp : instance of DataFieldRecordArray + data_exp The array holding the experimental data. The following data fields must exist: - log_energy : float + log_energy The logarithm of the reconstructed energy value of the data event. - sin_dec : float + sin_dec The sine of the reconstructed declination of the data event. - log10_energy_binning : instance of BinningDefinition + log10_energy_binning The binning definition for the binning in log10(E). - sin_dec_binning : instance of BinningDefinition + sin_dec_binning The binning definition for the sin(declination). - smoothing_filter : instance of SmoothingFilter | None + smoothing_filter The smoothing filter to use for smoothing the energy histogram. If None, no smoothing will be applied. """ @@ -397,11 +397,11 @@ class MCBackgroundI3EnergyPDF( def __init__( self, - data_mc, - physics_weight_field_names, - log10_energy_binning, - sin_dec_binning, - smoothing_filter=None, + data_mc: DataFieldRecordArray, + physics_weight_field_names: str | list[str], + log10_energy_binning: BinningDefinition, + sin_dec_binning: BinningDefinition, + smoothing_filter: SmoothingFilter | None = None, **kwargs, ): """Constructs a new IceCube energy background PDF from monte-carlo @@ -409,28 +409,28 @@ def __init__( Parameters ---------- - data_mc : instance of DataFieldRecordArray + data_mc The array holding the monte-carlo data. The following data fields must exist: - log_energy : float + log_energy The logarithm of the reconstructed energy value of the data event. - sin_dec : float + sin_dec The sine of the reconstructed declination of the data event. - mcweight: float + mcweight The monte-carlo weight of the event. - physics_weight_field_names : str | list of str + physics_weight_field_names The name or the list of names of the monte-carlo data fields, which should be used as physics event weights. If a list is given, the weight values of all the fields will be summed to construct the final event physics weight. - log10_energy_binning : BinningDefinition + log10_energy_binning The binning definition for the binning in log10(E). - sin_dec_binning : BinningDefinition + sin_dec_binning The binning definition for the sin(declination). - smoothing_filter : SmoothingFilter instance | None + smoothing_filter The smoothing filter to use for smoothing the energy histogram. If None, no smoothing will be applied. """ @@ -442,7 +442,7 @@ def __init__( ) if not issequence(physics_weight_field_names): - physics_weight_field_names = [physics_weight_field_names] + physics_weight_field_names = [physics_weight_field_names] # pyright: ignore[reportAssignmentType] if not issequenceof(physics_weight_field_names, str): raise TypeError( 'The physics_weight_field_names argument must be ' diff --git a/skyllh/i3/config.py b/skyllh/i3/config.py index 68c1a830c1..bb5a2553dd 100644 --- a/skyllh/i3/config.py +++ b/skyllh/i3/config.py @@ -1,17 +1,18 @@ """This file defines IceCube specific global configuration.""" +from skyllh.core.config import Config from skyllh.core.datafields import ( DataFieldStages as DFS, ) -def add_icecube_specific_analysis_required_data_fields(cfg): +def add_icecube_specific_analysis_required_data_fields(cfg: Config): """Adds IceCube specific data fields required by an IceCube analysis to the given local configuration. Parameters ---------- - cfg : instance of Config + cfg The instance of Config holding the local configuration. """ cfg['datafields']['azi'] = DFS.ANALYSIS_EXP diff --git a/skyllh/i3/dataset.py b/skyllh/i3/dataset.py index 600463b860..b479e209c4 100644 --- a/skyllh/i3/dataset.py +++ b/skyllh/i3/dataset.py @@ -1,4 +1,5 @@ import os.path +from collections.abc import Sequence import numpy as np @@ -9,6 +10,7 @@ Dataset, DatasetData, ) +from skyllh.core.livetime import Livetime from skyllh.core.logging import ( get_logger, ) @@ -20,9 +22,7 @@ DataFieldRecordArray, create_FileLoader, ) -from skyllh.core.timing import ( - TaskTimer, -) +from skyllh.core.timing import TaskTimer, TimeLord class I3Dataset( @@ -37,18 +37,18 @@ class I3Dataset( """ @staticmethod - def get_combined_grl_pathfilenames(datasets): + def get_combined_grl_pathfilenames(datasets: 'Sequence[I3Dataset]') -> list: """Creates the combined list of grl pathfilenames of all the given datasets. Parameters ---------- - datasets : sequence of I3Dataset + datasets The sequence of I3Dataset instances. Returns ------- - grl_pathfilenames : list + grl_pathfilenames The combined list of grl pathfilenames. """ if not issequenceof(datasets, I3Dataset): @@ -62,8 +62,8 @@ def get_combined_grl_pathfilenames(datasets): def __init__( self, - livetime=None, - grl_pathfilenames=None, + livetime: float | None = None, + grl_pathfilenames: str | Sequence[str] | None = None, **kwargs, ): """Creates a new IceCube specific dataset, that also can hold a list @@ -71,10 +71,10 @@ def __init__( Parameters ---------- - livetime : float | None + livetime The live-time of the dataset in days. It can be ``None``, if good-run-list data files are provided. - grl_pathfilenames : str | sequence of str + grl_pathfilenames The sequence of pathfilenames pointing to the good-run-list (GRL) data files. """ @@ -160,13 +160,13 @@ def __str__(self): def create_file_list( self, - ): + ) -> list[str]: """Creates the list of files of this dataset. The file paths are relative to the dataset's root directory. Returns ------- - file_list : list of str + file_list The list of files of this dataset. """ file_list = super().create_file_list() + self._grl_pathfilename_list @@ -175,9 +175,9 @@ def create_file_list( def load_grl( self, - efficiency_mode=None, - tl=None, - ): + efficiency_mode: str | None = None, + tl: TimeLord | None = None, + ) -> DataFieldRecordArray: """Loads the good-run-list and returns a DataFieldRecordArray instance which should contain the following data fields: @@ -194,7 +194,7 @@ def load_grl( Parameters ---------- - efficiency_mode : str | None + efficiency_mode The efficiency mode the data should get loaded with. Possible values are: @@ -209,12 +209,12 @@ def load_grl( The default value is ``'time'``. If set to ``None``, the default value will be used. - tl : TimeLord instance | None + tl The TimeLord instance to use to time the data loading procedure. Returns ------- - grl_data : instance of DataFieldRecordArray + grl_data The DataFieldRecordArray instance holding the good-run-list information of the dataset. """ @@ -230,36 +230,36 @@ def load_grl( def load_data( self, - livetime=None, - keep_fields=None, - dtc_dict=None, - dtc_except_fields=None, - efficiency_mode=None, - tl=None, - ): + livetime: Livetime | float | None = None, + keep_fields: list[str] | None = None, + dtc_dict: dict | None = None, + dtc_except_fields: str | Sequence[str] | None = None, + efficiency_mode: str | None = None, + tl: TimeLord | None = None, + ) -> DatasetData: """Loads the data, which is described by the dataset. If a good-run-list (GRL) is provided for this dataset, only experimental data will be selected which matches the GRL. Parameters ---------- - livetime : instance of Livetime | float | None + livetime If not None, uses this livetime (if float livetime in days) as livetime for the DatasetData instance, otherwise uses the live time from the Dataset instance or, if available, the livetime from the good-run-list (GRL). - keep_fields : list of str | None + keep_fields The list of user-defined data fields that should get loaded and kept in addition to the analysis required data fields. - dtc_dict : dict | None + dtc_dict This dictionary defines how data fields of specific data types (key) should get converted into other data types (value). This can be used to use less memory. If set to None, no data convertion is performed. - dtc_except_fields : str | sequence of str | None + dtc_except_fields The sequence of field names whose data type should not get converted. - efficiency_mode : str | None + efficiency_mode The efficiency mode the data should get loaded with. Possible values are: @@ -274,13 +274,13 @@ def load_data( The default value is ``'time'``. If set to ``None``, the default value will be used. - tl : TimeLord instance | None + tl The TimeLord instance that should be used to time the data load operation. Returns ------- - data : instance of DatasetData + data A DatasetData instance holding the experimental and monte-carlo data of this data set. """ @@ -322,10 +322,10 @@ def load_data( return data - def prepare_data( + def prepare_data( # pyright: ignore[reportIncompatibleMethodOverride] self, - data, - tl=None, + data: 'I3DatasetData', + tl: TimeLord | None = None, ): """Prepares the data for IceCube by pre-calculating the following experimental data fields: @@ -340,9 +340,9 @@ def prepare_data( Parameters ---------- - data : DatasetData instance + data The DatasetData instance holding the data as numpy record ndarray. - tl : TimeLord instance | None + tl The TimeLord instance that should be used to time the data preparation. """ @@ -424,17 +424,17 @@ class I3DatasetData( def __init__( self, - data, - data_grl, + data: DatasetData, + data_grl: DataFieldRecordArray | None, ): """Constructs a new I3DatasetData instance. Parameters ---------- - data : instance of DatasetData + data The instance of DatasetData holding the experimental and monte-carlo data. - data_grl : instance of DataFieldRecordArray | None + data_grl The instance of DataFieldRecordArray holding the good-run-list data of the dataset. This can be None, if no GRL data is available. """ @@ -443,7 +443,7 @@ def __init__( self.grl = data_grl @property - def grl(self): + def grl(self) -> DataFieldRecordArray | None: """The DataFieldRecordArray instance holding the good-run-list (GRL) data of the IceCube data set. It is None, if there is no GRL data available for this IceCube data set. @@ -451,7 +451,7 @@ def grl(self): return self._grl @grl.setter - def grl(self, data): + def grl(self, data: DataFieldRecordArray | None): if data is not None and not isinstance(data, DataFieldRecordArray): raise TypeError('The grl property must be an instance of DataFieldRecordArray!') self._grl = data diff --git a/skyllh/i3/detsigyield.py b/skyllh/i3/detsigyield.py index b5c8afe4e8..11416cf020 100644 --- a/skyllh/i3/detsigyield.py +++ b/skyllh/i3/detsigyield.py @@ -3,6 +3,8 @@ """ import abc +from collections.abc import Callable, Sequence +from typing import cast import numpy as np import scipy.interpolate @@ -14,23 +16,25 @@ from skyllh.core.binning import ( BinningDefinition, ) +from skyllh.core.dataset import Dataset, DatasetData from skyllh.core.detsigyield import ( DetSigYield, DetSigYieldBuilder, ) +from skyllh.core.flux_model import FluxModel from skyllh.core.livetime import ( Livetime, ) from skyllh.core.parameters import ( ParameterGrid, ) +from skyllh.core.progressbar import ProgressBar from skyllh.core.py import ( classname, issequenceof, ) -from skyllh.core.source_model import ( - PointLikeSource, -) +from skyllh.core.source_hypo_grouping import SourceHypoGroup +from skyllh.core.source_model import PointLikeSource, SourceModel class I3DetSigYield(DetSigYield, metaclass=abc.ABCMeta): @@ -39,22 +43,30 @@ class I3DetSigYield(DetSigYield, metaclass=abc.ABCMeta): detector effective area and hence the detector signal yield. """ - def __init__(self, param_names, dataset, fluxmodel, livetime, sin_dec_binning, **kwargs): + def __init__( + self, + param_names: Sequence[str], + dataset: Dataset, + fluxmodel: FluxModel, + livetime: float | Livetime, + sin_dec_binning: BinningDefinition, + **kwargs, + ): """Constructor of the IceCube specific detector signal yield base class. Parameters ---------- - param_names : sequence of str + param_names The sequence of parameter names this detector signal yield depends on. These are either fixed or floating parameters. - dataset : Dataset instance + dataset The Dataset instance holding the monte-carlo event data. - fluxmodel : FluxModel + fluxmodel The flux model instance. Must be an instance of FluxModel. - livetime : float | Livetime instance + livetime The live-time. - sin_dec_binning : BinningDefinition instance + sin_dec_binning The BinningDefinition instance defining the sin(dec) binning. """ super().__init__(param_names=param_names, dataset=dataset, fluxmodel=fluxmodel, livetime=livetime, **kwargs) @@ -83,7 +95,7 @@ class I3DetSigYieldBuilder(DetSigYieldBuilder, metaclass=abc.ABCMeta): def __init__( self, - sin_dec_binning=None, + sin_dec_binning: BinningDefinition | None = None, **kwargs, ): """Constructor of the IceCube specific detector signal yield @@ -91,7 +103,7 @@ def __init__( Parameters ---------- - sin_dec_binning : BinningDefinition instance + sin_dec_binning The instance of BinningDefinition defining the sin(dec) binning. """ super().__init__(**kwargs) @@ -140,22 +152,30 @@ class PointLikeSourceI3DetSigYield(I3DetSigYield): classes for point-like sources. """ - def __init__(self, param_names, dataset, fluxmodel, livetime, sin_dec_binning, **kwargs): + def __init__( + self, + param_names: Sequence[str], + dataset: Dataset, + fluxmodel: FluxModel, + livetime: float | Livetime, + sin_dec_binning: BinningDefinition, + **kwargs, + ): """Constructor of the IceCube specific detector signal yield base class for point-like sources. Parameters ---------- - param_names : sequence of str + param_names The sequence of parameter names this detector signal yield depends on. These are either fixed or floating parameters. - dataset : Dataset instance + dataset The Dataset instance holding the monte-carlo event data. - fluxmodel : FluxModel + fluxmodel The flux model instance. Must be an instance of FluxModel. - livetime : float | Livetime instance + livetime The livetime in days or an instance of Livetime. - sin_dec_binning : BinningDefinition instance + sin_dec_binning The BinningDefinition instance defining the sin(dec) binning. """ super().__init__( @@ -167,19 +187,19 @@ class for point-like sources. **kwargs, ) - def sources_to_recarray(self, sources): + def sources_to_recarray(self, sources: SourceModel | Sequence[SourceModel]) -> np.ndarray: # pyright: ignore[reportIncompatibleMethodOverride] """Converts the sequence of PointLikeSource sources into a numpy record array holding the information of the sources needed for the detector signal yield calculation. Parameters ---------- - sources : SourceModel | sequence of SourceModel + sources The source model(s) containing the information of the source(s). Returns ------- - recarr : numpy record ndarray + recarr The generated (N_sources,)-shaped 1D numpy record ndarray holding the information for each source. """ @@ -188,8 +208,9 @@ def sources_to_recarray(self, sources): if not issequenceof(sources, PointLikeSource): raise TypeError('The sources argument must be an instance or a sequence of instances of PointLikeSource!') - recarr = np.empty((len(sources),), dtype=[('dec', np.float64)]) - for i, src in enumerate(sources): + _sources = cast(list[PointLikeSource], sources) + recarr = np.empty((len(_sources),), dtype=[('dec', np.float64)]) + for i, src in enumerate(_sources): recarr['dec'][i] = src.dec return recarr @@ -209,14 +230,14 @@ class PointLikeSourceI3DetSigYieldBuilder( def __init__( self, - sin_dec_binning=None, + sin_dec_binning: BinningDefinition | None = None, **kwargs, ): """Initializes a new detector signal yield builder object. Parameters ---------- - sin_dec_binning : BinningDefinition | None + sin_dec_binning The BinningDefinition instance defining the sin(dec) binning that should be used to compute the sin(dec) dependency of the detector effective area. If set to None, the binning will be taken from the @@ -228,26 +249,35 @@ def __init__( class FixedFluxPointLikeSourceI3DetSigYield(PointLikeSourceI3DetSigYield): """The detector signal yield class for a point-source with a fixed flux.""" - def __init__(self, param_names, dataset, fluxmodel, livetime, sin_dec_binning, log_spl_sinDec, **kwargs): + def __init__( + self, + param_names: Sequence[str], + dataset: Dataset, + fluxmodel: FluxModel, + livetime: float | Livetime, + sin_dec_binning: BinningDefinition, + log_spl_sinDec: scipy.interpolate.InterpolatedUnivariateSpline, + **kwargs, + ): """Constructs an IceCube detector signal yield instance for a point-like source with a fixed flux. Parameters ---------- - param_names : sequence of str + param_names The sequence of parameter names this detector signal yield depends on. These are either fixed or floating parameters. - dataset : Dataset instance + dataset The instance of Dataset holding the monte-carlo data this detector signal yield is made for. - fluxmodel : FluxModel instance + fluxmodel The instance of FluxModel with fixed parameters this detector signal yield is made for. - livetime : float | Livetime instance + livetime The livetime in days or an instance of Livetime. - sin_dec_binning : BinningDefinition instance + sin_dec_binning The binning definition for sin(dec). - log_spl_sinDec : scipy.interpolate.InterpolatedUnivariateSpline + log_spl_sinDec The spline instance representing the log value of the detector signal yield as a function of sin(dec). """ @@ -278,23 +308,25 @@ def log_spl_sinDec(self, spl): ) self._log_spl_sinDec = spl - def __call__(self, src_recarray, src_params_recarray=None): + def __call__( # pyright: ignore[reportIncompatibleMethodOverride] + self, src_recarray: np.ndarray, src_params_recarray: None = None + ) -> tuple[np.ndarray, dict]: """Retrieves the detector signal yield for the list of given sources. Parameters ---------- - src_recarray : numpy record ndarray + src_recarray The numpy record ndarray with the field ``dec`` holding the declination of the source. - src_params_recarray : None + src_params_recarray Unused interface argument, because this detector signal yield does not depend on any source parameters. Returns ------- - values : numpy 1d ndarray + values The array with the detector signal yield for each source. - grads : dict + grads This detector signal yield does not depend on any parameters. So there are no gradients and the dictionary is empty. """ @@ -335,8 +367,8 @@ class FixedFluxPointLikeSourceI3DetSigYieldBuilder( def __init__( self, - sin_dec_binning=None, - spline_order_sinDec=2, + sin_dec_binning: BinningDefinition | None = None, + spline_order_sinDec: int = 2, **kwargs, ): """Creates a new IceCube detector signal yield builder object for a @@ -348,11 +380,11 @@ def __init__( Parameters ---------- - sin_dec_binning : BinningDefinition | None + sin_dec_binning The BinningDefinition instance which defines the sin(dec) binning. If set to None, the binning will be taken from the Dataset binning definitions. - spline_order_sinDec : int + spline_order_sinDec The order of the spline function for the logarithmic values of the detector signal yield along the sin(dec) axis. The default is 2. @@ -378,39 +410,39 @@ def spline_order_sinDec(self, order): def _create_hist( self, - data_sin_true_dec, - data_true_energy, - sin_dec_binning, - weights, - fluxmodel, - to_internal_flux_unit_factor, - ): + data_sin_true_dec: np.ndarray, + data_true_energy: np.ndarray, + sin_dec_binning: BinningDefinition, + weights: np.ndarray, + fluxmodel: FluxModel, + to_internal_flux_unit_factor: float, + ) -> np.ndarray: """Creates a histogram of the detector signal yield with the given sin(dec) binning for the given flux model. Parameters ---------- - data_sin_true_dec : instance of numpy.ndarray + data_sin_true_dec The (N_data,)-shaped numpy.ndarray holding the sin(true_dec) values of the monte-carlo events. - data_true_energy : instance of numpy.ndarray + data_true_energy The (N_data,)-shaped numpy.ndarray holding the true energy of the monte-carlo events. - sin_dec_binning : instance of BinningDefinition + sin_dec_binning The sin(dec) binning definition to use for the histogram. - weights : 1d ndarray + weights The (N_data,)-shaped numpy.ndarray holding the weight factor of each monte-carlo event where only the flux value needs to be multiplied with in order to get the detector signal yield. - fluxmodel : instance of FluxModel + fluxmodel The flux model to get the flux values from. - to_internal_flux_unit_factor : float + to_internal_flux_unit_factor The conversion factor to convert the flux unit into the internal flux unit. Returns ------- - hist : instance of numpy.ndarray + hist The (N_sin_dec_bins,)-shaped numpy.ndarray containing the histogram values. """ @@ -426,19 +458,19 @@ def _create_hist( def _create_detsigyield_from_hist( self, - hist, - sin_dec_binning, + hist: np.ndarray, + sin_dec_binning: BinningDefinition, **kwargs, - ): + ) -> 'FixedFluxPointLikeSourceI3DetSigYield': """Create a single instance of FixedFluxPointLikeSourceI3DetSigYield from the given histogram. Parameters ---------- - hist : instance of numpy.ndarray + hist The (N_sin_dec_bins,)-shaped numpy.ndarray holding the normalized histogram of the detector signal yield. - sin_dec_binning : instance of BinningDefinition + sin_dec_binning The sin(dec) binning definition to use for the histogram. **kwargs Additional keyword arguments are passed to the constructor of the @@ -446,7 +478,7 @@ def _create_detsigyield_from_hist( Returns ------- - detsigyield : instance of FixedFluxPointLikeSourceI3DetSigYield + detsigyield The instance of FixedFluxPointLikeSourceI3DetSigYield for the given flux model. """ @@ -463,19 +495,19 @@ def _create_detsigyield_from_hist( def construct_detsigyields( self, - dataset, - data, - shgs, - ppbar=None, - ): + dataset: Dataset, + data: DatasetData, + shgs: Sequence[SourceHypoGroup], + ppbar: ProgressBar | None = None, + ) -> 'list[FixedFluxPointLikeSourceI3DetSigYield]': """Constructs a set of FixedFluxPointLikeSourceI3DetSigYield instances, one for each provided fluxmodel. Parameters ---------- - dataset : instance of Dataset + dataset The instance of Dataset holding meta information about the data. - data : instance of DatasetData + data The instance of DatasetData holding the monte-carlo event data. The numpy record ndarray holding the monte-carlo event data must contain the following data fields: @@ -488,16 +520,16 @@ def construct_detsigyields( The monte-carlo weight of the data event in the unit GeV cm^2 sr. - shgs : sequence of instance of SourceHypoGroup + shgs The sequence of instance of SourceHypoGroup specifying the source hypothesis groups (i.e. flux model) for which the detector signal yields should get constructed. - ppbar : instance of ProgressBar | None + ppbar The optional instance of ProgressBar of the parent progress bar. Returns ------- - detsigyields : list of instance of FixedFluxPointLikeSourceI3DetSigYield + detsigyields The list of instance of FixedFluxPointLikeSourceI3DetSigYield providing the detector signal yield function for a point-like source with each of the given fixed flux models. @@ -511,6 +543,8 @@ def construct_detsigyields( to_internal_time_unit_factor = self._cfg.to_internal_time_unit(time_unit=units.day) # Get integrated live-time in days. + assert data.livetime is not None + assert data.mc is not None livetime_days = Livetime.get_integrated_livetime(data.livetime) # Get the sin(dec) binning definition either as setting from this @@ -566,13 +600,13 @@ def construct_detsigyields( return detsigyields - def construct_detsigyield( + def construct_detsigyield( # pyright: ignore[reportIncompatibleMethodOverride] self, - dataset, - data, - shg, - ppbar=None, - ): + dataset: Dataset, + data: DatasetData, + shg: SourceHypoGroup, + ppbar: ProgressBar | None = None, + ) -> 'FixedFluxPointLikeSourceI3DetSigYield': """Constructs a detector signal yield log spline function for the given fixed flux model. @@ -581,9 +615,9 @@ def construct_detsigyield( Parameters ---------- - dataset : instance of Dataset + dataset The instance of Dataset holding meta information about the data. - data : instance of DatasetData + data The instance of DatasetData holding the monte-carlo event data. The numpy record ndarray holding the monte-carlo event data must contain the following data fields: @@ -596,15 +630,15 @@ def construct_detsigyield( The monte-carlo weight of the data event in the unit GeV cm^2 sr. - shg : instance of SourceHypoGroup + shg The instance of SourceHypoGroup (i.e. sources and flux model) for which the detector signal yield should get constructed. - ppbar : instance of ProgressBar | None + ppbar The optional instance of ProgressBar of the parent progress bar. Returns ------- - detsigyield : instance of FixedFluxPointLikeSourceI3DetSigYield + detsigyield The instance of FixedFluxPointLikeSourceI3DetSigYield providing the detector signal yield function for a point-like source with a fixed flux. @@ -618,13 +652,13 @@ def construct_detsigyield( return detsigyield - def get_detsigyield_construction_factory(self): + def get_detsigyield_construction_factory(self) -> Callable: """Returns the factory callable for constructing a set of instance of FixedFluxPointLikeSourceI3DetSigYield. Returns ------- - factory : callable + factory The factory callable for constructing a set of instance of FixedFluxPointLikeSourceI3DetSigYield. """ @@ -637,23 +671,32 @@ class SingleParamFluxPointLikeSourceI3DetSigYield(PointLikeSourceI3DetSigYield): source parameter. """ - def __init__(self, param_name, dataset, fluxmodel, livetime, sin_dec_binning, log_spl_sinDec_param, **kwargs): + def __init__( + self, + param_name: str, + dataset: Dataset, + fluxmodel: FluxModel, + livetime: float | Livetime, + sin_dec_binning: BinningDefinition, + log_spl_sinDec_param: scipy.interpolate.RectBivariateSpline, + **kwargs, + ): """Constructs the detector signal yield instance. Parameters ---------- - param_name : str + param_name The parameter name this detector signal yield depends on. These are either fixed or floating parameter. - dataset : Dataset instance + dataset The Dataset instance holding the monte-carlo event data. - fluxmodel : FluxModel + fluxmodel The flux model instance. Must be an instance of FluxModel. - livetime : float | Livetime instance + livetime The live-time. - sin_dec_binning : BinningDefinition instance + sin_dec_binning The BinningDefinition instance defining the sin(dec) binning. - log_spl_sinDec_param : scipy.interpolate.RectBivariateSpline instance + log_spl_sinDec_param The 2D spline in sin(dec) and the parameter this detector signal yield depends on. """ @@ -686,16 +729,18 @@ def log_spl_sinDec_param(self, spl): ) self._log_spl_sinDec_param = spl - def __call__(self, src_recarray, src_params_recarray): + def __call__( # pyright: ignore[reportIncompatibleMethodOverride] + self, src_recarray: np.ndarray, src_params_recarray: np.ndarray + ) -> tuple[np.ndarray, dict]: """Retrieves the detector signal yield for the given list of sources and their flux parameters. Parameters ---------- - src_recarray : numpy record ndarray + src_recarray The numpy record ndarray with the field ``dec`` holding the declination of the source. - src_params_recarray : (N_sources,)-shaped numpy record ndarray + src_params_recarray The numpy record ndarray containing the parameter values of the sources. The parameter values can be different for the different sources. @@ -708,9 +753,9 @@ def __call__(self, src_recarray, src_params_recarray): Returns ------- - values : numpy (N_sources,)-shaped 1D ndarray + values The array with the detector signal yield for each source. - grads : dict + grads The dictionary holding the gradient values for each global floating parameter. The key is the global floating parameter index and the value is the (N_sources,)-shaped numpy ndarray holding the gradient @@ -792,11 +837,11 @@ class SingleParamFluxPointLikeSourceI3DetSigYieldBuilder( def __init__( self, - param_grid, - sin_dec_binning=None, - spline_order_sinDec=2, + param_grid: ParameterGrid, + sin_dec_binning: BinningDefinition | None = None, + spline_order_sinDec: int = 2, spline_order_param=2, - ncpu=None, + ncpu: int | None = None, **kwargs, ): """Creates a new IceCube detector signal yield builder instance for a @@ -807,23 +852,23 @@ def __init__( Parameters ---------- - param_grid : instance of ParameterGrid + param_grid The instance of ParameterGrid which defines the grid of the parameter values. The name of the parameter is defined via the name property of the ParameterGrid instance. - sin_dec_binning : instance of BinningDefinition | None + sin_dec_binning The instance of BinningDefinition which defines the sin(dec) binning. If set to None, the sin(dec) binning will be taken from the dataset's binning definitions. - spline_order_sinDec : int + spline_order_sinDec The order of the spline function for the logarithmic values of the detector signal yield along the sin(dec) axis. The default is 2. - spline_order_gamma : int + spline_order_param The order of the spline function for the logarithmic values of the - detector signal yield along the gamma axis. + detector signal yield along the parameter axis. The default is 2. - ncpu : int | None + ncpu The number of CPUs to utilize. If set to ``None``, global setting will take place. """ @@ -878,21 +923,21 @@ def spline_order_param(self, order): ) self._spline_order_param = order - def construct_detsigyield( + def construct_detsigyield( # pyright: ignore[reportIncompatibleMethodOverride] self, - dataset, - data, - shg, - ppbar=None, - ): + dataset: Dataset, + data: DatasetData, + shg: SourceHypoGroup, + ppbar: ProgressBar | None = None, + ) -> 'SingleParamFluxPointLikeSourceI3DetSigYield': """Constructs a detector signal yield 2-dimensional log spline function for the given flux model with varying parameter values. Parameters ---------- - dataset : instance of Dataset + dataset The instance of Dataset holding the sin(dec) binning definition. - data : instance of DatasetData + data The instance of DatasetData holding the monte-carlo event data. The numpy record array for the monte-carlo data of the dataset must contain the following data fields: @@ -905,21 +950,23 @@ def construct_detsigyield( ``'true_energy'`` : float The true energy value of the data event. - shg : instance of SourceHypoGroup + shg The instance of SourceHypoGroup for which the detector signal yield should get constructed. - ppbar : instance of ProgressBar | None + ppbar The instance of ProgressBar of the optional parent progress bar. Returns ------- - detsigyield : instance of SingleParamFluxPointLikeSourceI3DetSigYield + detsigyield The I3DetSigYield instance for a point-like source with a flux model of a single parameter. """ self.assert_types_of_construct_detsigyield_arguments(dataset=dataset, data=data, shgs=shg, ppbar=ppbar) # Get integrated live-time in days. + assert data.livetime is not None + assert data.mc is not None livetime_days = Livetime.get_integrated_livetime(data.livetime) # Get the sin(dec) binning definition either as setting from this @@ -936,37 +983,37 @@ def construct_detsigyield( # along sin(dec) for a given flux model, i.e. for given spectral index, # gamma. def _create_hist( - data_sin_true_dec, - data_true_energy, - sin_dec_binning, - weights, - fluxmodel, - to_internal_flux_unit_factor, - ): + data_sin_true_dec: np.ndarray, + data_true_energy: np.ndarray, + sin_dec_binning: BinningDefinition, + weights: np.ndarray, + fluxmodel: FluxModel, + to_internal_flux_unit_factor: float, + ) -> np.ndarray: """Creates a histogram of the detector signal yield with the given sin(dec) binning. Parameters ---------- - data_sin_true_dec : 1d ndarray + data_sin_true_dec The sin(true_dec) values of the monte-carlo events. - data_true_energy : 1d ndarray + data_true_energy The true energy of the monte-carlo events. - sin_dec_binning : BinningDefinition + sin_dec_binning The sin(dec) binning definition to use for the histogram. - weights : 1d ndarray + weights The weight factors of each monte-carlo event where only the flux value needs to be multiplied with in order to get the detector signal yield. - fluxmodel : FluxModel + fluxmodel The flux model to get the flux values from. - to_internal_flux_unit_factor : float + to_internal_flux_unit_factor The conversion factor to convert the flux unit into the internal flux unit. Returns ------- - h : 1d ndarray + h The numpy array containing the histogram values. """ weights = weights * fluxmodel(E=data_true_energy).squeeze() * to_internal_flux_unit_factor diff --git a/skyllh/i3/livetime.py b/skyllh/i3/livetime.py index 4bf49f5f4d..5f485310ab 100644 --- a/skyllh/i3/livetime.py +++ b/skyllh/i3/livetime.py @@ -17,24 +17,24 @@ class I3Livetime(Livetime): """ @classmethod - def from_grl_data(cls, grl_data): + def from_grl_data(cls, grl_data: np.ndarray) -> 'I3Livetime': """Creates an I3LiveTime instance from the given good-run-list (GRL) data. Parameters ---------- - grl_data : instance of numpy structured ndarray. + grl_data The numpy structured ndarray of length N_runs holding the start end end times of the good runs. The following fields need to exist: - start : float + start The MJD of the run start. - end : float + end The MJD of the run stop. Returns ------- - livetime : instance of I3Livetime + livetime The created instance of I3Livetime for the provided GRL data. """ uptime_mjd_intervals_arr = np.hstack( @@ -46,7 +46,7 @@ def from_grl_data(cls, grl_data): return livetime @staticmethod - def from_grl_files(pathfilenames): + def from_grl_files(pathfilenames: str | list[str]) -> 'I3Livetime': """Loads an I3Livetime instance from the given good-run-list (GRL) data file. The data file needs to contain the following data fields: @@ -57,12 +57,12 @@ def from_grl_files(pathfilenames): Parameters ---------- - pathfilenames : str | list of str - The list of fully qualified file names of the GRL data files. + pathfilenames + The list of fully qualified file name(s) of the GRL data file(s). Returns ------- - livetime : instance of I3Livetime + livetime The created instance of I3Livetime for the provided GRL data. """ grl_data = create_FileLoader(pathfilenames).load_data() @@ -76,19 +76,19 @@ def from_grl_files(pathfilenames): return livetime @staticmethod - def from_I3Dataset(ds): + def from_I3Dataset(ds: I3Dataset) -> 'I3Livetime': """Loads an I3Livetime instance from a given I3Dataset instance, which must have a good-run-list (GRL) files defined. Parameters ---------- - ds : I3Dataset instance + ds The instance of I3Dataset which defined the good-run-list (GRL) files for the dataset. Returns ------- - livetime : instance of I3Livetime + livetime The created instance of I3Livetime for the GRL data from the provided dataset. """ diff --git a/skyllh/i3/pdf.py b/skyllh/i3/pdf.py index a91782086a..d0d312ae6a 100644 --- a/skyllh/i3/pdf.py +++ b/skyllh/i3/pdf.py @@ -1,11 +1,10 @@ import numpy as np -from skyllh.core.binning import ( - UsesBinning, -) +from skyllh.core.binning import BinningDefinition, UsesBinning from skyllh.core.logging import ( get_logger, ) +from skyllh.core.parameters import ParameterModelMapper from skyllh.core.pdf import ( EnergyPDF, PDFAxis, @@ -20,9 +19,8 @@ NoHistSmoothingMethod, SmoothingFilter, ) -from skyllh.core.timing import ( - TaskTimer, -) +from skyllh.core.timing import TaskTimer, TimeLord +from skyllh.core.trialdata import TrialDataManager logger = get_logger(__name__) @@ -41,41 +39,41 @@ class I3EnergyPDF( def __init__( self, - pmm, - data_log10_energy, - data_sin_dec, - data_mcweight, - data_physicsweight, - log10_energy_binning, - sin_dec_binning, - smoothing_filter, + pmm: ParameterModelMapper | None, + data_log10_energy: np.ndarray, + data_sin_dec: np.ndarray, + data_mcweight: np.ndarray, + data_physicsweight: np.ndarray, + log10_energy_binning: BinningDefinition, + sin_dec_binning: BinningDefinition, + smoothing_filter: SmoothingFilter | None, **kwargs, ): """Creates a new IceCube energy PDF object. Parameters ---------- - pmm : instance of ParameterModelMapper | None + pmm The instance of ParameterModelMapper defining the global parameters and their mapping to local model/source parameters. It can be ``None``, if the PDF does not depend on any parameters. - data_log10_energy : 1d ndarray + data_log10_energy The array holding the log10(E) values of the events. - data_sin_dec : 1d ndarray + data_sin_dec The array holding the sin(dec) values of the events. - data_mcweight : 1d ndarray + data_mcweight The array holding the monte-carlo weights of the events. The final data weight will be the product of data_mcweight and data_physicsweight. - data_physicsweight : 1d ndarray + data_physicsweight The array holding the physics weights of the events. The final data weight will be the product of data_mcweight and data_physicsweight. - log10_energy_binning : instance of BinningDefinition + log10_energy_binning The binning definition for the log10(E) axis. - sin_dec_binning : instance of BinningDefinition + sin_dec_binning The binning definition for the sin(declination) axis. - smoothing_filter : instance of SmoothingFilter | None + smoothing_filter The smoothing filter to use for smoothing the energy histogram. If ``None``, no smoothing will be applied. """ @@ -205,23 +203,23 @@ def hist_mask_mc_covered_with_physics(self): mask = self._hist_mask_mc_covered & ~self._hist_mask_mc_covered_zero_physics return mask - def assert_is_valid_for_trial_data(self, tdm, tl=None, **kwargs): + def assert_is_valid_for_trial_data(self, tdm: TrialDataManager, tl: TimeLord | None = None, **kwargs): """Checks if this energy PDF is valid for all the given trial events. It checks if all the data is within the log10(E) and sin(dec) binning range. Parameters ---------- - tdm : instance of TrialDataManager + tdm The instance of TrialDataManager holding the trial data events. The following data fields must exist: - log_energy : float + log_energy The base-10 logarithm of the energy value of the data event. - dec : float + dec The declination of the data event. - tl : instance of TimeLord | None + tl The optional instance of TimeLord for measuring timing information. Raises @@ -254,33 +252,35 @@ def assert_is_valid_for_trial_data(self, tdm, tl=None, **kwargs): f'The following data values are out of range: {oor_data}' ) - def get_pd(self, tdm, params_recarray=None, tl=None): + def get_pd( # pyright: ignore[reportIncompatibleMethodOverride] + self, tdm: TrialDataManager, params_recarray: None = None, tl: TimeLord | None = None + ) -> tuple[np.ndarray, dict]: """Calculates the energy probability density of each event. Parameters ---------- - tdm : instance of TrialDataManager + tdm The TrialDataManager instance holding the data events for which the probability density should be calculated. The following data fields must exist: - log_energy : float + log_energy The base-10 logarithm of the energy value of the event. - sin_dec : float + sin_dec The sin(declination) value of the event. - params_recarray : None + params_recarray Unused interface parameter. - tl : TimeLord instance | None + tl The optional TimeLord instance that should be used to measure timing information. Returns ------- - pd : instance of ndarray + pd The 1D (N_events,)-shaped numpy ndarray with the energy probability density for each event. - grads : dict + grads The dictionary holding the gradients of the probability density w.r.t. each fit parameter. The key of the dictionary is the id of the global fit parameter. Because this energy PDF does not depend diff --git a/skyllh/i3/pdfratio.py b/skyllh/i3/pdfratio.py index 1df6fa1fd9..8f5c1bc11f 100644 --- a/skyllh/i3/pdfratio.py +++ b/skyllh/i3/pdfratio.py @@ -4,6 +4,7 @@ repack_fields, ) +from skyllh.core.interpolate import GridManifoldInterpolationMethod from skyllh.core.multiproc import ( IsParallelizable, parallelize, @@ -15,9 +16,12 @@ MostSignalLikePDFRatioFillMethod, PDFRatioFillMethod, ) +from skyllh.core.progressbar import ProgressBar from skyllh.core.py import ( make_dict_hash, ) +from skyllh.core.timing import TimeLord +from skyllh.core.trialdata import TrialDataManager class SplinedI3EnergySigSetOverBkgPDFRatio(SigSetOverBkgPDFRatio, IsParallelizable): @@ -35,10 +39,10 @@ def __init__( self, sig_pdf_set, bkg_pdf, - fillmethod=None, - interpolmethod_cls=None, - ncpu=None, - ppbar=None, + fillmethod: PDFRatioFillMethod | None = None, + interpolmethod_cls: type[GridManifoldInterpolationMethod] | None = None, + ncpu: int | None = None, + ppbar: ProgressBar | None = None, **kwargs, ): """Creates a new IceCube signal-over-background energy PDF ratio spline @@ -46,25 +50,25 @@ def __init__( Parameters ---------- - sig_pdf_set : class instance derived from PDFSet (for PDF type + sig_pdf_set I3EnergyPDF), IsSignalPDF, and UsesBinning The PDF set, which provides signal energy PDFs for a set of discrete signal parameters. - bkg_pdf : class instance derived from I3EnergyPDF, and + bkg_pdf IsBackgroundPDF The background energy PDF object. - fillmethod : instance of PDFRatioFillMethod | None + fillmethod An instance of class derived from PDFRatioFillMethod that implements the desired ratio fill method. If set to None (default), the default ratio fill method MostSignalLikePDFRatioFillMethod will be used. - interpolmethod_cls : class of GridManifoldInterpolationMethod + interpolmethod_cls The class implementing the parameter interpolation method for the PDF ratio manifold grid. - ncpu : int | None + ncpu The number of CPUs to use to create the ratio splines for the different sets of signal parameters. - ppbar : ProgressBar instance | None + ppbar The instance of ProgressBar of the optional parent progress bar. Raises @@ -92,7 +96,7 @@ def create_log_ratio_spline(sig_pdf_set, bkg_pdf, fillmethod, gridparams): Returns ------- - log_ratio_spline : instance of RegularGridInterpolator + log_ratio_spline The spline of the logarithmic PDF ratio values. """ # Get the signal PDF for the given signal parameters. @@ -149,7 +153,7 @@ def create_log_ratio_spline(sig_pdf_set, bkg_pdf, fillmethod, gridparams): self._gridparams_hash_log_ratio_spline_dict[gridparams_hash] = log_ratio_spline # Save the list of data field names. - self._data_field_names = [binning.name for binning in self._bkg_pdf.binnings] + self._data_field_names = [binning.name for binning in self._bkg_pdf.binnings] # pyright: ignore[reportAttributeAccessIssue] # Construct the instance for the parameter interpolation method. self._interpolmethod = self._interpolmethod_cls( @@ -180,20 +184,26 @@ def fillmethod(self, obj): raise TypeError('The fillmethod property must be an instance of PDFRatioFillMethod!') self._fillmethod = obj - def _create_cache(self, trial_data_state_id, interpol_params_recarray, ratio, grads): + def _create_cache( + self, + trial_data_state_id: int | None, + interpol_params_recarray: np.ndarray | None, + ratio: np.ndarray | None, + grads: np.ndarray | None, + ): """Creates a cache dictionary holding cache data. Parameters ---------- - trial_data_state_id : int | None + trial_data_state_id The trial data state ID of the TrialDataManager. - interpol_params_recarray : instance of numpy record ndarray | None + interpol_params_recarray The numpy record ndarray of length N_sources holding the parameter names and values necessary for the interpolation for all sources. - ratio : instance of numpy ndarray + ratio The (N_values,)-shaped numpy ndarray holding the PDF ratio values for all sources and trial events. - grads : instance of numpy ndarray + grads The (D,N_values)-shaped numpy ndarray holding the gradients for each PDF ratio value w.r.t. each interpolation parameter. """ @@ -218,18 +228,20 @@ def _is_cached(self, trial_data_state_id, interpol_params_recarray): return np.all(self._cache['interpol_params_recarray'] == interpol_params_recarray) - def _get_spline_for_param_values(self, interpol_param_values): + def _get_spline_for_param_values( + self, interpol_param_values: np.ndarray + ) -> scipy.interpolate.RegularGridInterpolator: """Retrieves the spline for a given set of parameter values. Parameters ---------- - interpol_param_values : instance of numpy ndarray + interpol_param_values The (N_interpol_params,)-shaped numpy ndarray holding the values of the interpolation parameters. Returns ------- - spline : instance of scipy.interpolate.RegularGridInterpolator + spline The requested spline instance. """ gridparams = dict(zip(self._interpol_param_names, interpol_param_values, strict=True)) @@ -239,37 +251,40 @@ def _get_spline_for_param_values(self, interpol_param_values): return spline - def _evaluate_splines(self, tdm, eventdata, gridparams_recarray, n_values): + def _evaluate_splines( + self, tdm: TrialDataManager, eventdata: np.ndarray, gridparams_recarray: np.ndarray, n_values: int + ) -> np.ndarray: """For each set of parameter values given by ``gridparams_recarray``, the spline is retrieved and evaluated for the events suitable for that source model. Parameters ---------- - tdm : instance of TrialDataManager + tdm The TrialDataManager instance holding the trial data and the event mapping to the sources via the ``src_evt_idx`` property. - eventdata : instance of numpy ndarray + eventdata The (V,N_events)-shaped numpy ndarray holding the event data, where N_events is the number of events, and V the dimensionality of the event data. - gridparams_recarray : instance of numpy structured ndarray + gridparams_recarray The numpy structured ndarray of length N_sources with the parameter names and values needed for the interpolation on the grid for all sources. If the length of this record array is 1, the set of parameters will be used for all sources. - n_values : int + n_values The size of the output array. Returns ------- - values : instance of ndarray + values The (N_values,)-shaped numpy ndarray holding the values for each set of parameter values of the ``gridparams_recarray``. The length of the array depends on the ``src_evt_idx`` property of the TrialDataManager. In the worst case it is ``N_sources * N_selected_events``. """ + assert tdm.src_evt_idxs is not None (src_idxs, evt_idxs) = tdm.src_evt_idxs # Check for special case when a single set of parameters are provided. @@ -300,7 +315,7 @@ def _evaluate_splines(self, tdm, eventdata, gridparams_recarray, n_values): return values - def _create_interpol_params_recarray(self, src_params_recarray): + def _create_interpol_params_recarray(self, src_params_recarray: np.ndarray) -> np.ndarray: """Creates the params_recarray needed for the interpolation. It selects The interpolation parameters from the ``params_recarray`` argument. If all parameters have the same value for all sources, the length will @@ -308,13 +323,13 @@ def _create_interpol_params_recarray(self, src_params_recarray): Parameters ---------- - src_params_recarray : instance of numpy record ndarray + src_params_recarray The numpy record ndarray of length N_sources holding all local parameter names and values. Returns ------- - interpol_params_recarray : instance of numpy record ndarray + interpol_params_recarray The numpy record ndarray of length N_sources or 1 holding only the parameters needed for the interpolation. """ @@ -330,16 +345,16 @@ def _create_interpol_params_recarray(self, src_params_recarray): return interpol_params_recarray - def _calculate_ratio_and_grads(self, tdm, interpol_params_recarray): + def _calculate_ratio_and_grads(self, tdm: TrialDataManager, interpol_params_recarray: np.ndarray): """Calculates the ratio values and ratio gradients for all the sources and trial events given the source parameter values. The result is stored in the class member variable ``_cache``. Parameters ---------- - tdm : instance of TrialDataManager + tdm The TrialDataManager instance holding the trial data. - interpol_params_recarray : instance of numpy record ndarray + interpol_params_recarray The numpy record ndarray of length N_sources holding the parameter names and values for all sources. It must contain only the parameters necessary for the interpolation. @@ -364,7 +379,9 @@ def _calculate_ratio_and_grads(self, tdm, interpol_params_recarray): grads=grads, ) - def get_ratio(self, tdm, src_params_recarray, tl=None): + def get_ratio( + self, tdm: TrialDataManager, src_params_recarray: np.ndarray, tl: TimeLord | None = None + ) -> np.ndarray: """Retrieves the PDF ratio values for each given trial event data, given the given set of fit parameters. This method is called during the likelihood maximization process. @@ -373,21 +390,21 @@ def get_ratio(self, tdm, src_params_recarray, tl=None): Parameters ---------- - tdm : instance of TrialDataManager + tdm The TrialDataManager instance holding the trial event data for which the PDF ratio values should get calculated. - src_params_recarray : instance of numpy record ndarray | None + src_params_recarray The (N_sources,)-shaped numpy record ndarray holding the parameter names and values of the sources. See the :meth:`skyllh.core.parameters.ParameterModelMapper.create_src_params_recarray` for more information. - tl : instance of TimeLord | None + tl The optional TimeLord instance that should be used to measure timing information. Returns ------- - ratio : instance of numpy ndarray + ratio The (N_values,)-shaped numpy ndarray of float holding the PDF ratio value for each source and trial event. """ @@ -404,35 +421,42 @@ def get_ratio(self, tdm, src_params_recarray, tl=None): return self._cache['ratio'] - def get_gradient(self, tdm, src_params_recarray, fitparam_id, tl=None): + def get_gradient( + self, + tdm: TrialDataManager, + src_params_recarray: np.ndarray | None, + fitparam_id: int, + tl: TimeLord | None = None, + ) -> np.ndarray: """Retrieves the PDF ratio gradient for the given fit parameter ``fitparam_id``. Parameters ---------- - tdm : instance of TrialDataManager + tdm The TrialDataManager instance holding the trial event data for which the PDF ratio gradient values should get calculated. - src_params_recarray : instance of numpy record ndarray | None + src_params_recarray The (N_sources,)-shaped numpy record ndarray holding the local parameter names and values of all sources. See the :meth:`skyllh.core.parameters.ParameterModelMapper.create_src_params_recarray` method for more information. - fitparam_id : int + fitparam_id The ID of the global fit parameter for which the gradient should get calculated. - tl : instance of TimeLord | None + tl The optional TimeLord instance that should be used to measure timing information. Returns ------- - grad : instance of ndarray + grad The (N_values,)-shaped numpy ndarray holding the gradient values for all sources and trial events w.r.t. the given global fit parameter. """ # Select only the parameters necessary for the interpolation. + assert src_params_recarray is not None interpol_params_recarray = self._create_interpol_params_recarray(src_params_recarray) # Calculate the gradients if necessary. diff --git a/skyllh/i3/scrambling.py b/skyllh/i3/scrambling.py index 6567fafa22..8d43c47420 100644 --- a/skyllh/i3/scrambling.py +++ b/skyllh/i3/scrambling.py @@ -1,9 +1,14 @@ import numpy as np +from skyllh.core.dataset import Dataset +from skyllh.core.random import RandomStateService from skyllh.core.scrambling import ( DataScramblingMethod, TimeScramblingMethod, ) +from skyllh.core.storage import DataFieldRecordArray +from skyllh.core.times import TimeGenerator +from skyllh.i3.dataset import I3DatasetData from skyllh.i3.utils.coords import ( azi_to_ra_transform, hor_to_equ_transform, @@ -20,14 +25,14 @@ class I3TimeScramblingMethod( def __init__( self, - timegen, + timegen: TimeGenerator, **kwargs, ): """Initializes a new I3 time scrambling instance. Parameters ---------- - timegen : TimeGenerator + timegen The time generator that should be used to generate random MJD times. """ super().__init__(timegen=timegen, hor_to_equ_transform=hor_to_equ_transform, **kwargs) @@ -36,28 +41,28 @@ def __init__( # the ``ra`` field. def scramble( self, - rss, - dataset, - data, - ): + rss: RandomStateService, + dataset: Dataset, + data: DataFieldRecordArray, + ) -> DataFieldRecordArray: """Draws a time from the time generator and calculates the right ascention coordinate from the azimuth angle according to the time. Sets the values of the ``time`` and ``ra`` keys of data. Parameters ---------- - rss : RandomStateService + rss The random state service providing the random number generator (RNG). - dataset : instance of Dataset + dataset The instance of Dataset for which the data should get scrambled. - data : DataFieldRecordArray instance + data The DataFieldRecordArray instance containing the to be scrambled data. Returns ------- - data : numpy record ndarray + data The given numpy record ndarray holding the scrambled data. """ mjds = self._timegen.generate_times(rss, len(data)) @@ -78,14 +83,14 @@ class I3SeasonalVariationTimeScramblingMethod( def __init__( self, - data, + data: I3DatasetData, **kwargs, ): """Initializes a new seasonal time scrambling instance. Parameters ---------- - data : instance of I3DatasetData + data The instance of I3DatasetData holding the experimental data and good-run-list information. """ @@ -93,42 +98,46 @@ def __init__( # The run weights are the number of events in each run relative to all # the events to account for possible seasonal variations. - self.run_weights = np.zeros((len(data.grl),), dtype=np.float64) + _grl = data.grl + assert _grl is not None + assert data.exp is not None + self.run_weights = np.zeros((len(_grl),), dtype=np.float64) n_events = len(data.exp['time']) - for i, (start, stop) in enumerate(zip(data.grl['start'], data.grl['stop'], strict=True)): + for i, (start, stop) in enumerate(zip(_grl['start'], _grl['stop'], strict=True)): mask = (data.exp['time'] >= start) & (data.exp['time'] < stop) self.run_weights[i] = len(data.exp[mask]) / n_events self.run_weights /= np.sum(self.run_weights) - self.grl = data.grl + self.grl = _grl def scramble( self, - rss, - dataset, - data, - ): + rss: RandomStateService, + dataset: Dataset, + data: DataFieldRecordArray, + ) -> DataFieldRecordArray: """Scrambles the given data based on random MJD times, which are generated uniformely within the data runs, where the data runs are weighted based on their amount of events compared to the total events. Parameters ---------- - rss : instance of RandomStateService + rss The random state service providing the random number generator (RNG). - dataset : instance of Dataset + dataset The instance of Dataset for which the data should get scrambled. - data : instance of DataFieldRecordArray + data The DataFieldRecordArray instance containing the to be scrambled data. Returns ------- - data : instance of DataFieldRecordArray + data The given DataFieldRecordArray holding the scrambled data. """ # Get run indices based on their seasonal weights. + assert self.grl is not None run_idxs = rss.random.choice(self.grl['start'].size, size=len(data['time']), p=self.run_weights) # Draw random times uniformely within the runs. diff --git a/skyllh/i3/signal_generation.py b/skyllh/i3/signal_generation.py index ed52338f50..2ddc3d4879 100644 --- a/skyllh/i3/signal_generation.py +++ b/skyllh/i3/signal_generation.py @@ -1,3 +1,5 @@ +from collections.abc import Callable + import numpy as np from skyllh.core.py import ( @@ -8,6 +10,7 @@ from skyllh.core.signal_generation import ( SignalGenerationMethod, ) +from skyllh.core.source_hypo_grouping import SourceHypoGroup from skyllh.core.source_model import ( PointLikeSource, ) @@ -16,7 +19,7 @@ ) -def source_sin_dec_shift_linear(x, w, L, U): +def source_sin_dec_shift_linear(x: np.ndarray, w: float, L: float, U: float) -> np.ndarray: """Calculates the shift of the sine of the source declination, in order to allow the construction of the source sine declination band with sin(dec_src) +/- w. This shift function, S(x), is implemented as a line @@ -28,18 +31,18 @@ def source_sin_dec_shift_linear(x, w, L, U): Parameters ---------- - x : 1D numpy ndarray + x The sine of the source declination for each source. - w : float + w The half size of the sin(dec)-window. - L : float + L The lower value of the allowed sin(dec) range. - U : float + U The upper value of the allowed sin(dec) range. Returns ------- - S : 1D numpy ndarray + S The sin(dec) shift of the sin(dec) values of the given sources, such that ``sin(dec_src) + S`` is the new sin(dec) of the source, and ``sin(dec_src) + S +/- w`` is always within the sin(dec) range [L, U]. @@ -53,7 +56,7 @@ def source_sin_dec_shift_linear(x, w, L, U): return S -def source_sin_dec_shift_cubic(x, w, L, U): +def source_sin_dec_shift_cubic(x: np.ndarray, w: float, L: float, U: float) -> np.ndarray: """Calculates the shift of the sine of the source declination, in order to allow the construction of the source sine declination band with sin(dec_src) +/- w. This shift function, S(x), is implemented as a cubic @@ -65,18 +68,18 @@ def source_sin_dec_shift_cubic(x, w, L, U): Parameters ---------- - x : 1D numpy ndarray + x The sine of the source declination for each source. - w : float + w The half size of the sin(dec)-window. - L : float + L The lower value of the allowed sin(dec) range. - U : float + U The upper value of the allowed sin(dec) range. Returns ------- - S : 1D numpy ndarray + S The sin(dec) shift of the sin(dec) values of the given sources, such that ``sin(dec_src) + S`` is the new sin(dec) of the source, and ``sin(dec_src) + S +/- w`` is always within the sin(dec) range [L, U]. @@ -99,10 +102,10 @@ class PointLikeSourceI3SignalGenerationMethod(SignalGenerationMethod): def __init__( self, - src_sin_dec_half_bandwidth=_DEFAULT_SRC_SIN_DEC_HALF_BANDWIDTH, - src_sin_dec_shift_func=None, - energy_range=None, - src_batch_size=128, + src_sin_dec_half_bandwidth: float = _DEFAULT_SRC_SIN_DEC_HALF_BANDWIDTH, + src_sin_dec_shift_func: Callable | None = None, + energy_range: tuple[float, float] | None = None, + src_batch_size: int = 128, **kwargs, ): """Constructs a new signal generation method instance for a point-like @@ -110,19 +113,19 @@ def __init__( Parameters ---------- - src_sin_dec_half_bandwidth : float + src_sin_dec_half_bandwidth The half-width of the sin(dec) band to take MC events from around a source. The default is sin(1deg), i.e. a 1deg half-bandwidth. - src_sin_dec_shift_func : callable | None + src_sin_dec_shift_func The function that provides the source sin(dec) shift needed for constructing the source declination bands from where to draw monte-carlo events from. If set to None, the default function ``source_sin_dec_shift_linear`` will be used. - energy_range : 2-element tuple of float | None + energy_range The energy range from which to take MC events into account for signal event generation, specified in true neutrino energy (GeV). If set to None, the entire energy range [0, +inf] is used. - src_batch_size : int + src_batch_size The source processing batch size used for the signal event flux calculation. """ @@ -174,27 +177,29 @@ def src_batch_size(self, v): v = int_cast(v, 'The src_batch_size property must be cast-able to type int!') self._src_batch_size = v - def _get_src_dec_bands(self, src_dec, max_sin_dec_range): + def _get_src_dec_bands( + self, src_dec: np.ndarray, max_sin_dec_range: tuple + ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: """Calculates the minimum and maximum sin(dec) values for each source to use with a specified maximal sin(dec) range, which should get determined from the available MC data itself. Parameters ---------- - src_dec : 1D ndarray + src_dec The declination values of the sources. - max_sin_dec_range : 2-element tuple of floats + max_sin_dec_range The maximal sin(dec) range from where MC events are available. Returns ------- - src_sin_dec_band_min : (N_sources,)-shaped 1D ndarray + src_sin_dec_band_min The array holding the lower value of the sin(dec) band for each source. - src_sin_dec_band_max : (N_sources,)-shaped 1D ndarray + src_sin_dec_band_max The array holding the upper value of the sin(dec) band for each source. - src_dec_band_omega : (N_sources,)-shaped 1D ndarray + src_dec_band_omega The solid angle of the declination band for each source. """ # Shift the source declination in order to be able to always create the @@ -210,28 +215,30 @@ def _get_src_dec_bands(self, src_dec, max_sin_dec_range): return (src_sin_dec_band_min, src_sin_dec_band_max, src_dec_band_omega) - def calc_source_signal_mc_event_flux(self, data_mc, shg): + def calc_source_signal_mc_event_flux( # pyright: ignore[reportIncompatibleMethodOverride] + self, data_mc: np.ndarray, shg: SourceHypoGroup + ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: """Calculates the signal flux of each given MC event for each source hypothesis of the given source hypothesis group. Parameters ---------- - data_mc : numpy record ndarray + data_mc The numpy record array holding the MC events of a dataset. - shg : SourceHypoGroup instance + shg The source hypothesis group, which defines the list of sources, and their flux model. Returns ------- - ev_idx_arr : ndarray + ev_idx_arr The (N_selected_signal_events,)-shaped 1D ndarray holding the index of the MC event. - shg_src_idx_arr : ndarray + shg_src_idx_arr The (N_selected_signal_events,)-shaped 1D ndarray holding the index of the source within the given source hypothesis group for each signal candidate event. - flux_arr : ndarray + flux_arr The (N_selected_signal_events,)-shaped 1D ndarray holding the flux value of each signal candidate event. """ @@ -308,21 +315,21 @@ def calc_source_signal_mc_event_flux(self, data_mc, shg): return (ev_idx_arr, shg_src_idx_arr, flux_arr) - def signal_event_post_sampling_processing( + def signal_event_post_sampling_processing( # pyright: ignore[reportIncompatibleMethodOverride] self, - shg, - shg_sig_events_meta, - shg_sig_events, - ): + shg: SourceHypoGroup, + shg_sig_events_meta: np.ndarray, + shg_sig_events: np.ndarray, + ) -> np.ndarray: """Rotates the generated signal events to their source location for a given source hypothesis group. Parameters ---------- - shg : SourceHypoGroup instance + shg The source hypothesis group instance holding the sources and their locations. - shg_sig_events_meta : numpy record ndarray + shg_sig_events_meta The numpy record ndarray holding meta information about the generated signal events for the given source hypothesis group. The length of this array must be the same as shg_sig_events. @@ -331,14 +338,14 @@ def signal_event_post_sampling_processing( - 'shg_src_idx': int The source index within the source hypothesis group. - shg_sig_events : numpy record ndarray + shg_sig_events The numpy record ndarray holding the generated signal events for the given source hypothesis group and in the format of the original MC events. Returns ------- - shg_sig_events : numpy record ndarray + shg_sig_events The numpy record ndarray with the processed MC signal events. """ # Get the unique source indices of that source hypo group. diff --git a/skyllh/i3/signalpdf.py b/skyllh/i3/signalpdf.py index 3c2e27f670..643ac9ac4a 100644 --- a/skyllh/i3/signalpdf.py +++ b/skyllh/i3/signalpdf.py @@ -3,6 +3,7 @@ from skyllh.core.binning import ( BinningDefinition, ) +from skyllh.core.config import Config from skyllh.core.flux_model import ( FluxModel, ) @@ -19,18 +20,20 @@ IsSignalPDF, PDFSet, ) +from skyllh.core.progressbar import ProgressBar from skyllh.core.py import ( classname, ) from skyllh.core.smoothing import ( SmoothingFilter, ) +from skyllh.core.storage import DataFieldRecordArray from skyllh.i3.pdf import ( I3EnergyPDF, ) -class SignalI3EnergyPDFSet( +class SignalI3EnergyPDFSet( # pyright: ignore[reportIncompatibleMethodOverride] PDFSet, IsSignalPDF, PDF, @@ -43,15 +46,15 @@ class SignalI3EnergyPDFSet( def __init__( self, - cfg, - data_mc, - log10_energy_binning, - sin_dec_binning, - fluxmodel, - param_grid_set, - smoothing_filter=None, - ncpu=None, - ppbar=None, + cfg: Config, + data_mc: DataFieldRecordArray, + log10_energy_binning: BinningDefinition, + sin_dec_binning: BinningDefinition, + fluxmodel: FluxModel, + param_grid_set: ParameterGridSet | ParameterGrid, + smoothing_filter: SmoothingFilter | None = None, + ncpu: int | None = None, + ppbar: ProgressBar | None = None, **kwargs, ): """Creates a new IceCube energy signal PDF for a given flux model and @@ -62,44 +65,43 @@ def __init__( Parameters ---------- - cfg : instance of Config + cfg The instance of Config holding the local configuration. - data_mc : instance of DataFieldRecordArray + data_mc The instance of DataFieldRecordArray holding the monte-carlo data. The following data fields must exist: - true_energy : float + true_energy The true energy value of the data event. - log_energy : float + log_energy The base10-logarithm of the reconstructed energy value of the data event. - sin_dec : float + sin_dec The declination of the data event. - mcweight : float + mcweight The monte-carlo weight value of the data events in unit GeV cm^2 sr. - log10_energy_binning : instance of BinningDefinition + log10_energy_binning The binning definition for the reconstructed energy binning in log10(E). - sin_dec_binning : instance of BinningDefinition + sin_dec_binning The binning definition for the binning in sin(declination). - fluxmodel : instance of FluxModel + fluxmodel The flux model to use to create the signal energy PDF. - param_grid_set : instance of ParameterGridSet | - instance of ParameterGrid + param_grid_set The set of parameter grids. A ParameterGrid instance for each energy parameter, for which an I3EnergyPDF object needs to be created. - smoothing_filter : instance of SmoothingFilter | None + smoothing_filter The smoothing filter to use for smoothing the energy histogram. If ``None``, no smoothing will be applied. - ncpu : int | None + ncpu The number of CPUs to use to create the different I3EnergyPDF instances for the different parameter grid values. If set to ``None``, the configured default number of CPUs will be used. - ppbar : instance of ProgressBar | None + ppbar The instance of ProgressBar of the optional parent progress bar. """ if isinstance(param_grid_set, ParameterGrid): @@ -145,51 +147,51 @@ def __init__( # Create I3EnergyPDF objects for all permutations of the parameter # grid values. def create_I3EnergyPDF( - cfg, - data_log10_energy, - data_sin_dec, - data_mcweight, - data_true_energy, - log10_energy_binning, - sin_dec_binning, - smoothing_filter, - fluxmodel, - flux_unit_conv_factor, - gridparams, - ): + cfg: Config, + data_log10_energy: np.ndarray, + data_sin_dec: np.ndarray, + data_mcweight: np.ndarray, + data_true_energy: np.ndarray, + log10_energy_binning: BinningDefinition, + sin_dec_binning: BinningDefinition, + smoothing_filter: SmoothingFilter | None, + fluxmodel: FluxModel, + flux_unit_conv_factor: float, + gridparams: dict, + ) -> I3EnergyPDF: """Creates an I3EnergyPDF object for the given flux model and flux parameters. Parameters ---------- - cfg : instance of Config + cfg The instance of Config holding the local configuration. - data_log10_energy : 1d ndarray + data_log10_energy The base-10 logarithm of the reconstructed energy value of the data events. - data_sin_dec : 1d ndarray + data_sin_dec The sin(dec) value of the the data events. - data_mcweight : 1d ndarray + data_mcweight The monte-carlo weight value of the data events. - data_true_energy : 1d ndarray + data_true_energy The true energy value of the data events. - log10_energy_binning : instance of BinningDefinition + log10_energy_binning The binning definition for the binning in log10(E). - sin_dec_binning : instance of BinningDefinition + sin_dec_binning The binning definition for the sin(declination). - smoothing_filter : instance of SmoothingFilter | None + smoothing_filter The smoothing filter to use for smoothing the energy histogram. If ``None``, no smoothing will be applied. - fluxmodel : instance of FluxModel + fluxmodel The flux model to use to create the signal event weights. - flux_unit_conv_factor : float + flux_unit_conv_factor The factor to convert the flux unit into the internal flux unit. - gridparams : dict + gridparams The dictionary holding the specific signal flux parameters. Returns ------- - i3energypdf : instance of I3EnergyPDF + i3energypdf The created I3EnergyPDF instance for the given flux model and flux parameters. """ diff --git a/skyllh/i3/utils/analysis.py b/skyllh/i3/utils/analysis.py index de44e348cc..33b0a62ff9 100644 --- a/skyllh/i3/utils/analysis.py +++ b/skyllh/i3/utils/analysis.py @@ -1,11 +1,13 @@ import numpy as np +from skyllh.core.analysis import SingleSourceMultiDatasetLLHRatioAnalysis from skyllh.core.logging import ( get_logger, ) from skyllh.core.progressbar import ( ProgressBar, ) +from skyllh.core.random import RandomStateService from skyllh.core.source_model import ( PointLikeSource, ) @@ -16,8 +18,16 @@ def generate_ps_sin_dec_h0_ts_values( - ana, rss, sin_dec_min, sin_dec_max, sin_dec_step, n_bkg_trials=10000, n_iter=1, bkg_kwargs=None, ppbar=None -): + ana: SingleSourceMultiDatasetLLHRatioAnalysis, + rss: RandomStateService, + sin_dec_min: float, + sin_dec_max: float, + sin_dec_step: float, + n_bkg_trials: int = 10000, + n_iter: int = 1, + bkg_kwargs: dict | None = None, + ppbar: ProgressBar | None = None, +) -> tuple[np.ndarray, np.ndarray]: """Generates sets of null-hypothesis, i.e. background-only trial data events, test-statistic values for the given point-source analysis for a grid of sin(dec) values. @@ -27,40 +37,40 @@ def generate_ps_sin_dec_h0_ts_values( Parameters ---------- - ana : Analysis instance + ana The instance of Analysis to use for generating trials. - rss : RandomStateService instance + rss The instance of RandomStateService to use for generating random numbers from. - sin_dec_min : float + sin_dec_min The minimum sin(dec) value for creating the sin(dec) value grid. - sin_dec_max : float + sin_dec_max The maximum sin(dec) value for creating the sin(dec) value grid. - sin_dec_step : float + sin_dec_step The step size in sin(dec) for creating the sin(dec) value grid. - n_bkg_trials : int, optional + n_bkg_trials The number of background trials to generate. Default is 10000. - n_iter : int, optional + n_iter Each set of ts values can be calculated several times to be able to estimate the variance of the sensitivity / discovery potential. This parameter specifies the number of iterations to perform. For each iteration the RandomStateService is re-seeded with a seed that is incremented by 1. Default is 1. - bkg_kwargs : dict | None + bkg_kwargs Additional keyword arguments for the `generate_events` method of the background generation method class. An usual keyword argument is `poisson`. - ppbar : ProgressBar instance | None + ppbar The optional parent progress bar. Returns ------- - sin_dec_arr : (n_sin_dec,)-shaped 1D ndarray of float + sin_dec_arr The numpy ndarray holding the sin(dec) values for which the ts values have been calculated. - h0_ts_vals_arr : (n_sin_dec,n_iter,n_bkg_trials)-shaped 3D ndarray of float + h0_ts_vals_arr The numpy ndarray holding the null-hypothesis ts values for all sin(dec) values and iterations. """ @@ -96,6 +106,7 @@ def generate_ps_sin_dec_h0_ts_values( pbar_iter.increment() + assert rss.seed is not None rss.reseed(rss.seed + 1) pbar_iter.finish() @@ -103,66 +114,66 @@ def generate_ps_sin_dec_h0_ts_values( def estimate_ps_sin_dec_sensitivity_curve( - ana, - rss, - sin_dec_arr, - h0_ts_vals_arr, - eps_p=0.0075, - mu_min=0, - mu_max=20, - n_iter=1, - bkg_kwargs=None, - sig_kwargs=None, - ppbar=None, + ana: SingleSourceMultiDatasetLLHRatioAnalysis, + rss: RandomStateService, + sin_dec_arr: np.ndarray, + h0_ts_vals_arr: np.ndarray, + eps_p: float = 0.0075, + mu_min: float = 0, + mu_max: float = 20, + n_iter: int = 1, + bkg_kwargs: dict | None = None, + sig_kwargs: dict | None = None, + ppbar: ProgressBar | None = None, **kwargs, -): +) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: """Estimates the point-source sensitivity of the given analysis as a function of sin(dec). This function places a PointLikeSource source on the given declination values and estimates its sensitivity. Parameters ---------- - ana : Analysis instance + ana The instance of Analysis to use for generating trials. - rss : RandomStateService instance + rss The instance of RandomStateService to use for generating random numbers from. - sin_dec_arr : (n_sin_dec,)-shaped 1D ndarray + sin_dec_arr The ndarray holding the sin(dec) values for which to estimate the point-source sensitivity. - h0_ts_vals_arr : (n_sin_dec,n_iter,n_bkg_trials)-shaped 3D ndarray of float + h0_ts_vals_arr The numpy ndarray holding the null-hypothesis ts values for all sin(dec) values and iterations. - eps_p : float, optional + eps_p The precision in probability used for the `estimate_sensitivity` function. Default is 0.0075. - mu_min : float, optional + mu_min The minimum value for the mean number of injected signal events as a seed for the mu range in which the sensitivity is located. Default is 0. - mu_max : float, optional + mu_max The maximum value for the mean number of injected signal events as a seed for the mu range in which the sensitivity is located. Default is 20. - n_iter : int, optional + n_iter Each sensitivity can be estimated several times to be able to estimate the variance of the sensitivity. This parameter specifies the number of iterations to perform for each sensitivity. For each iteration the RandomStateService is re-seeded with a seed that is incremented by 1. Default is 1. - bkg_kwargs : dict | None + bkg_kwargs Additional keyword arguments for the `generate_events` method of the background generation method class. An usual keyword argument is `poisson`. - sig_kwargs : dict | None + sig_kwargs Additional keyword arguments for the `generate_signal_events` method of the `SignalGenerator` class. An usual keyword argument is `poisson`. If `poisson` is set to True, the actual number of generated signal events will be drawn from a Poisson distribution with the given mean number of signal events. - ppbar : ProgressBar instance | None + ppbar The optional parent progress bar. Additional Keyword Arguments @@ -172,16 +183,16 @@ def estimate_ps_sin_dec_sensitivity_curve( Returns ------- - sin_dec_arr : (n_sin_dec,)-shaped 1D ndarray + sin_dec_arr The ndarray holding the sin(dec) values for which the sensitivities have been estimated. - mean_ns_arr : (n_sin_dec,n_iter)-shaped 2D ndarray + mean_ns_arr The ndarray holding the mean number of signal events corresponding to the sensitivity for each sin(dec) value and iteration. - mean_ns_err_arr : (n_sin_dec,n_iter)-shaped 2D ndarray + mean_ns_err_arr The ndarray holding the estimated error in `mean_ns` for each sin(dec) value and iteration. - flux_scaling_arr : (n_sin_dec,n_iter)-shaped 2D ndarray + flux_scaling_arr The ndarray holding the scaling factor the reference flux needs to get scaled to obtain the flux for the estimated sensitivity. """ @@ -212,7 +223,7 @@ def estimate_ps_sin_dec_sensitivity_curve( rss, mu_range=(mu_min, mu_max), eps_p=eps_p, - h0_ts_vals=h0_ts_vals, + h0_trials=h0_ts_vals, bkg_kwargs=bkg_kwargs, sig_kwargs=sig_kwargs, ppbar=pbar_iter, @@ -227,6 +238,7 @@ def estimate_ps_sin_dec_sensitivity_curve( mu_min_arr = np.mean(mean_ns_arr[:, 0 : iter_idx + 1] * 0.8, axis=1) mu_max_arr = np.mean(mean_ns_arr[:, 0 : iter_idx + 1] * 1.2, axis=1) + assert rss.seed is not None rss.reseed(rss.seed + 1) pbar_iter.increment() @@ -263,7 +275,7 @@ def estimate_ps_sin_dec_sensitivity_curve( rss, mu_range=(mu_min, mu_max), eps_p=eps_p, - h0_ts_vals=h0_ts_vals, + h0_trials=h0_ts_vals, bkg_kwargs=bkg_kwargs, sig_kwargs=sig_kwargs, ppbar=pbar_sin_dec, @@ -280,72 +292,72 @@ def estimate_ps_sin_dec_sensitivity_curve( def estimate_ps_sin_dec_discovery_potential_curve( - ana, - rss, - sin_dec_arr, - h0_ts_vals_arr, - h0_ts_quantile=2.7e-3, - eps_p=0.0075, - mu_min=0, - mu_max=20, - n_iter=1, - bkg_kwargs=None, - sig_kwargs=None, - ppbar=None, + ana: SingleSourceMultiDatasetLLHRatioAnalysis, + rss: RandomStateService, + sin_dec_arr: np.ndarray, + h0_ts_vals_arr: np.ndarray, + h0_ts_quantile: float = 2.7e-3, + eps_p: float = 0.0075, + mu_min: float = 0, + mu_max: float = 20, + n_iter: int = 1, + bkg_kwargs: dict | None = None, + sig_kwargs: dict | None = None, + ppbar: ProgressBar | None = None, **kwargs, -): +) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: """Estimates the point-source discovery potential of the given analysis as a function of sin(dec). This function places a PointLikeSource source on the given declination values and estimates its discovery potential. Parameters ---------- - ana : Analysis instance + ana The instance of Analysis to use for generating trials. - rss : RandomStateService instance + rss The instance of RandomStateService to use for generating random numbers from. - sin_dec_arr : (n_sin_dec,)-shaped 1D ndarray + sin_dec_arr The ndarray holding the sin(dec) values for which to estimate the point-source sensitivity. - h0_ts_vals_arr : (n_sin_dec,n_iter,n_bkg_trials)-shaped 3D ndarray of float + h0_ts_vals_arr The numpy ndarray holding the null-hypothesis ts values for all sin(dec) values and iterations. - h0_ts_quantile : float, optional + h0_ts_quantile Null-hypothesis test statistic quantile that defines the critical value. For a 5sigma discovery potential that value is 5.733e-7. For a 3sigma discovery potential this value is 2.7e-3. Default is 2.7e-3. - eps_p : float, optional + eps_p The precision in probability used for the `estimate_discovery_potential` function. Default is 0.0075. - mu_min : float, optional + mu_min The minimum value for the mean number of injected signal events as a seed for the mu range in which the sensitivity is located. Default is 0. - mu_max : float, optional + mu_max The maximum value for the mean number of injected signal events as a seed for the mu range in which the sensitivity is located. Default is 20. - n_iter : int, optional + n_iter Each discovery potential can be estimated several times to be able to estimate its variance. This parameter specifies the number of iterations to perform for each discovery potential estimation. For each iteration the RandomStateService is re-seeded with a seed that is incremented by 1. Default is 1. - bkg_kwargs : dict | None + bkg_kwargs Additional keyword arguments for the `generate_events` method of the background generation method class. An usual keyword argument is `poisson`. - sig_kwargs : dict | None + sig_kwargs Additional keyword arguments for the `generate_signal_events` method of the `SignalGenerator` class. An usual keyword argument is `poisson`. If `poisson` is set to True, the actual number of generated signal events will be drawn from a Poisson distribution with the mean number of signal events, mu. - ppbar : ProgressBar instance | None + ppbar The optional parent progress bar. Additional Keyword Arguments @@ -355,16 +367,16 @@ def estimate_ps_sin_dec_discovery_potential_curve( Returns ------- - sin_dec_arr : (n_sin_dec,)-shaped 1D ndarray + sin_dec_arr The ndarray holding the sin(dec) values for which the discovery potential have been estimated. - mean_ns_arr : (n_sin_dec,n_iter)-shaped 2D ndarray + mean_ns_arr The ndarray holding the mean number of signal events corresponding to the discovery potential for each sin(dec) value and iteration. - mean_ns_err_arr : (n_sin_dec,n_iter)-shaped 2D ndarray + mean_ns_err_arr The ndarray holding the estimated error in `mean_ns` for each sin(dec) value and iteration. - flux_scaling_arr : (n_sin_dec,n_iter)-shaped 2D ndarray + flux_scaling_arr The ndarray holding the scaling factor the reference flux needs to get scaled to obtain the flux for the estimated discovery potential. """ @@ -393,7 +405,7 @@ def estimate_ps_sin_dec_discovery_potential_curve( h0_ts_quantile=h0_ts_quantile, mu_range=(mu_min, mu_max), eps_p=eps_p, - h0_ts_vals=h0_ts_vals, + h0_trials=h0_ts_vals, bkg_kwargs=bkg_kwargs, sig_kwargs=sig_kwargs, ppbar=pbar, @@ -413,6 +425,7 @@ def estimate_ps_sin_dec_discovery_potential_curve( pbar_iter.increment() + assert rss.seed is not None rss.reseed(rss.seed + 1) pbar_iter.finish() diff --git a/skyllh/i3/utils/coords.py b/skyllh/i3/utils/coords.py index dcddbc07b0..807c3e1d4c 100644 --- a/skyllh/i3/utils/coords.py +++ b/skyllh/i3/utils/coords.py @@ -1,9 +1,17 @@ """IceCube specific coordinate utility functions.""" +from typing import overload + import numpy as np -def azi_to_ra_transform(azi, mjd): +@overload +def azi_to_ra_transform(azi: np.ndarray, mjd: float | np.ndarray) -> np.ndarray: ... +@overload +def azi_to_ra_transform(azi: float, mjd: float) -> float: ... +@overload +def azi_to_ra_transform(azi: float | np.ndarray, mjd: float | np.ndarray) -> float | np.ndarray: ... +def azi_to_ra_transform(azi: float | np.ndarray, mjd: float | np.ndarray) -> float | np.ndarray: """Rotates the given IceCube azimuth angles into right-ascention angles for the given MJD times. This function is IceCube specific and assumes that the detector is located excently at the South Pole and neglects all astronomical @@ -11,14 +19,14 @@ def azi_to_ra_transform(azi, mjd): Parameters ---------- - azi : instance of numpy.ndarray + azi The array with the azimuth angles. - mjd : instance of numpy.ndarray + mjd The array with the MJD times for each azimuth angle. Returns ------- - ra : instance of numpy.ndarray + ra The right-ascention values. """ # sidereal day = length * solar day @@ -31,19 +39,19 @@ def azi_to_ra_transform(azi, mjd): return ra -def ra_to_azi_transform(ra, mjd): +def ra_to_azi_transform(ra: float | np.ndarray, mjd: float | np.ndarray) -> float | np.ndarray: """Rotates the given right-ascention angles to local IceCube azimuth angles. Parameters ---------- - ra : instance of numpy.ndarray + ra The array with the right-ascention angles. - mjd : instance of numpy.ndarray + mjd The array with the MJD times for each right-ascention angle. Returns ------- - azi : instance of numpy.ndarray + azi The azimuth angle for each right-ascention angle. """ # Use the azi_to_ra_transform function because it is symmetric. @@ -52,7 +60,19 @@ def ra_to_azi_transform(ra, mjd): return azi -def hor_to_equ_transform(azi, zen, mjd): +@overload +def hor_to_equ_transform( + azi: np.ndarray, zen: np.ndarray, mjd: float | np.ndarray +) -> tuple[np.ndarray, np.ndarray]: ... +@overload +def hor_to_equ_transform(azi: float, zen: float, mjd: float) -> tuple[float, float]: ... +@overload +def hor_to_equ_transform( + azi: float | np.ndarray, zen: float | np.ndarray, mjd: float | np.ndarray +) -> tuple[float | np.ndarray, float | np.ndarray]: ... +def hor_to_equ_transform( + azi: float | np.ndarray, zen: float | np.ndarray, mjd: float | np.ndarray +) -> tuple[float | np.ndarray, float | np.ndarray]: """Transforms the coordinate from the horizontal system (azimuth, zenith) into the equatorial system (right-ascention, declination) for detector at the South Pole and neglecting all astronomical effects like Earth @@ -60,18 +80,18 @@ def hor_to_equ_transform(azi, zen, mjd): Parameters ---------- - azi : instance of numpy.ndarray + azi The azimuth angle. - zen : instance of numpy.ndarray + zen The zenith angle. - mjd : instance of numpy.ndarray + mjd The time in MJD. Returns ------- - ra : instance of numpy.ndarray + ra The right-ascention angle. - dec : instance of numpy.ndarray + dec The declination angle. """ ra = azi_to_ra_transform(azi, mjd) diff --git a/skyllh/plotting/core/pdfratio.py b/skyllh/plotting/core/pdfratio.py index b3a5335a93..7d4e24fc16 100644 --- a/skyllh/plotting/core/pdfratio.py +++ b/skyllh/plotting/core/pdfratio.py @@ -1,13 +1,18 @@ """Plotting module for core PDF ratio objects.""" import itertools +from typing import cast import numpy as np from matplotlib.axes import Axes from matplotlib.colors import LogNorm +from matplotlib.image import AxesImage +from skyllh.core.parameters import ParameterModelMapper +from skyllh.core.pdf import PDF from skyllh.core.pdfratio import SigOverBkgPDFRatio from skyllh.core.py import classname +from skyllh.core.source_hypo_grouping import SourceHypoGroupManager from skyllh.core.storage import DataFieldRecordArray from skyllh.core.trialdata import TrialDataManager @@ -15,16 +20,16 @@ class SigOverBkgPDFRatioPlotter: """Plotter class to plot a SigOverBkgPDFRatio object.""" - def __init__(self, tdm, pdfratio): + def __init__(self, tdm: TrialDataManager, pdfratio: SigOverBkgPDFRatio): """Creates a new plotter object for plotting a - SpatialSigOverBkgPDFRatio object. + SigOverBkgPDFRatio object. Parameters ---------- - tdm : instance of TrialDataManager + tdm The instance of TrialDataManager that provides the data for the PDF ratio evaluation. - pdfratio : SpatialSigOverBkgPDFRatio + pdfratio The PDF ratio object to plot. """ self.tdm = tdm @@ -52,23 +57,34 @@ def tdm(self, obj): raise TypeError('The tdm property must be an instance of TrialDataManager!') self._tdm = obj - def plot(self, src_hypo_group_manager, axes, source_idx=None, log=True, **kwargs): + def plot( + self, + src_hypo_group_manager: SourceHypoGroupManager, + pmm: ParameterModelMapper, + src_params_recarray: np.ndarray, + axes: Axes, + source_idx: int | None = None, + log: bool = True, + **kwargs, + ) -> AxesImage: """Plots the spatial PDF ratio. If the signal PDF depends on the source, source_idx specifies the index of the source for which the PDF should get plotted. Parameters ---------- - src_hypo_group_manager : instance of SourceHypoGroupManager + src_hypo_group_manager The instance of SourceHypoGroupManager that defines the source hypotheses. - axes : mpl.axes.Axes + axes The matplotlib Axes object on which the PDF ratio should get drawn to. - source_idx : int | None + source_idx The index of the source for which the PDF ratio should get plotted. If set to None and the signal PDF depends on the source, index 0 will be used. + log + Whether to use a logarithmic color scale for the PDF ratio image. Additional Keyword Arguments ---------------------------- @@ -77,7 +93,7 @@ def plot(self, src_hypo_group_manager, axes, source_idx=None, log=True, **kwargs Returns ------- - img : instance of mpl.AxesImage + img The AxesImage instance showing the PDF ratio image. """ if not isinstance(axes, Axes): @@ -90,8 +106,8 @@ def plot(self, src_hypo_group_manager, axes, source_idx=None, log=True, **kwargs delta_ra_deg = 0.5 delta_dec_deg = 0.5 - raaxis = self._pdfratio.signalpdf.axes['ra'] - decaxis = self._pdfratio.signalpdf.axes['dec'] + raaxis = cast(PDF, self._pdfratio.sig_pdf).axes['ra'] + decaxis = cast(PDF, self._pdfratio.sig_pdf).axes['dec'] # Create a grid of ratio in right-ascention and declination and fill it # with PDF ratio values from events that fall into these bins. @@ -131,9 +147,9 @@ def plot(self, src_hypo_group_manager, axes, source_idx=None, log=True, **kwargs events['sin_dec'][i] = np.sin(dec) events['ang_err'][i] = np.deg2rad(0.5) - self._tdm.initialize_for_new_trial(src_hypo_group_manager, events) + self._tdm.initialize_trial(shg_mgr=src_hypo_group_manager, pmm=pmm, events=events) - event_ratios = self._pdfratio.get_ratio(self._tdm) + event_ratios = self._pdfratio.get_ratio(self._tdm, src_params_recarray) # Select only the ratios for the requested source. if event_ratios.ndim == 2: diff --git a/skyllh/plotting/core/signalpdf.py b/skyllh/plotting/core/signalpdf.py index d0beb2f709..38c1e54d6a 100644 --- a/skyllh/plotting/core/signalpdf.py +++ b/skyllh/plotting/core/signalpdf.py @@ -4,6 +4,9 @@ from matplotlib.axes import Axes from matplotlib.colors import LogNorm +from skyllh.core.parameters import ( + ParameterModelMapper, +) from skyllh.core.pdf import ( IsSignalPDF, SpatialPDF, @@ -27,7 +30,7 @@ class SignalSpatialPDFPlotter: def __init__( self, - tdm, + tdm: TrialDataManager, pdf, **kwargs, ): @@ -36,10 +39,10 @@ def __init__( Parameters ---------- - tdm : instance of TrialDataManager + tdm The instance of TrialDataManager that provides the data for the PDF evaluation. - pdf : class instance derived from SpatialPDF and IsSignalPDF + pdf The PDF object to plot. """ super().__init__(**kwargs) @@ -72,27 +75,30 @@ def tdm(self, obj): def plot( self, - src_hypo_group_manager, - axes, - source_idx=None, - sin_dec=True, - log=True, + src_hypo_group_manager: SourceHypoGroupManager, + pmm: ParameterModelMapper, + axes: Axes, + source_idx: int | None = None, + sin_dec: bool = True, + log: bool = True, **kwargs, ): """Plots the signal spatial PDF for the specified source. Parameters ---------- - axes : mpl.axes.Axes + axes The matplotlib Axes object on which the PDF ratio should get drawn to. - source_idx : int | None + source_idx The index of the source for which the PDF ratio should get plotted. If set to None and the signal PDF depends on the source, index 0 will be used. - sin_dec : bool + sin_dec Flag if the plot should be made in right-ascention vs. declination (False), or in right-ascention vs. sin(declination) (True). + log + Flag if it should be plotted in logarithmic color scale. Additional Keyword Arguments ---------------------------- @@ -101,7 +107,7 @@ def plot( Returns ------- - img : instance of mpl.AxesImage + img The AxesImage instance showing the PDF ratio image. """ if not isinstance(src_hypo_group_manager, SourceHypoGroupManager): @@ -164,9 +170,9 @@ def plot( events['dec'][i] = dec events['ang_err'][i] = np.deg2rad(sigma_deg) - self._tdm.initialize_for_new_trial(src_hypo_group_manager, events) + self._tdm.initialize_trial(shg_mgr=src_hypo_group_manager, pmm=pmm, events=events) - event_probs = self._pdf.get_prob(self._tdm) + (event_probs, _) = self._pdf.get_pd(self._tdm) # Select only the probabilities for the requested source. if event_probs.ndim == 2: diff --git a/skyllh/plotting/i3/backgroundpdf.py b/skyllh/plotting/i3/backgroundpdf.py index 5ee8de0c4e..249a5c2ece 100644 --- a/skyllh/plotting/i3/backgroundpdf.py +++ b/skyllh/plotting/i3/backgroundpdf.py @@ -8,6 +8,9 @@ LogNorm, ) +from skyllh.core.parameters import ( + ParameterModelMapper, +) from skyllh.core.py import ( classname, ) @@ -28,16 +31,16 @@ class BackgroundI3SpatialPDFPlotter: """Plotter class to plot an BackgroundI3SpatialPDF object.""" - def __init__(self, tdm, pdf): + def __init__(self, tdm: TrialDataManager, pdf: BackgroundI3SpatialPDF): """Creates a new plotter object for plotting an BackgroundI3SpatialPDF object. Parameters ---------- - tdm : instance of TrialDataManager + tdm The instance of TrialDataManager that provides the data for the PDF evaluation. - pdf : instance of BackgroundI3SpatialPDF + pdf The PDF object to plot. """ self.tdm = tdm @@ -65,21 +68,21 @@ def tdm(self, obj): raise TypeError('The tdm property must be an instance of TrialDataManager!') self._tdm = obj - def plot(self, src_hypo_group_manager, axes): + def plot(self, src_hypo_group_manager: SourceHypoGroupManager, pmm: ParameterModelMapper, axes: Axes): """Plots the spatial PDF. It uses the sin(dec) binning of the PDF to propperly represent the resolution of the PDF in the drawing. Parameters ---------- - src_hypo_group_manager : instance of SourceHypoGroupManager + src_hypo_group_manager The instance of SourceHypoGroupManager that defines the source hypotheses. - axes : mpl.axes.Axes + axes The matplotlib Axes object on which the PDF should get drawn to. Returns ------- - img : instance of mpl.AxesImage + img The AxesImage instance showing the PDF image. """ if not isinstance(src_hypo_group_manager, SourceHypoGroupManager): @@ -98,9 +101,9 @@ def plot(self, src_hypo_group_manager, axes): for i, sin_dec in enumerate(sin_dec_points): events['sin_dec'][i] = sin_dec - self._tdm.initialize_for_new_trial(src_hypo_group_manager, events) + self._tdm.initialize_trial(shg_mgr=src_hypo_group_manager, pmm=pmm, events=events) - event_probs = self._pdf.get_prob(self._tdm) + (event_probs, _) = self._pdf.get_pd(self._tdm) for i in range(len(events)): pdfprobs[0, i] = event_probs[i] diff --git a/skyllh/plotting/i3/pdf.py b/skyllh/plotting/i3/pdf.py index 7c8493232e..c6651c4dac 100644 --- a/skyllh/plotting/i3/pdf.py +++ b/skyllh/plotting/i3/pdf.py @@ -10,6 +10,9 @@ LogNorm, ) +from skyllh.core.parameters import ( + ParameterModelMapper, +) from skyllh.core.py import ( classname, ) @@ -30,15 +33,15 @@ class I3EnergyPDFPlotter: """Plotter class to plot an I3EnergyPDF object.""" - def __init__(self, tdm, pdf): + def __init__(self, tdm: TrialDataManager, pdf: I3EnergyPDF): """Creates a new plotter object for plotting an I3EnergyPDF object. Parameters ---------- - tdm : instance of TrialDataManager + tdm The instance of TrialDataManager that provides the data for the PDF evaluation. - pdf : I3EnergyPDF + pdf The PDF object to plot. """ self.tdm = tdm @@ -66,19 +69,19 @@ def tdm(self, obj): raise TypeError('The tdm property must be an instance of TrialDataManager!') self._tdm = obj - def plot(self, src_hypo_group_manager, axes, **kwargs): + def plot(self, src_hypo_group_manager: SourceHypoGroupManager, pmm: ParameterModelMapper, axes: Axes, **kwargs): """Plots the PDF object. Parameters ---------- - src_hypo_group_manager : instance of SourceHypoGroupManager + src_hypo_group_manager The instance of SourceHypoGroupManager that defines the source hypotheses. - axes : mpl.axes.Axes - The matplotlib Axes object on which the PDF ratio should get drawn - to. - fitparams : dict - The dictionary with the set of fit paramater values. + pmm + The instance of ParameterModelMapper that maps the global parameters + to the local model parameters. + axes + The matplotlib Axes object on which the PDF should get drawn to. Additional Keyword Arguments ---------------------------- @@ -87,7 +90,7 @@ def plot(self, src_hypo_group_manager, axes, **kwargs): Returns ------- - img : instance of mpl.AxesImage + img The AxesImage instance showing the PDF ratio image. """ if not isinstance(src_hypo_group_manager, SourceHypoGroupManager): @@ -113,9 +116,9 @@ def plot(self, src_hypo_group_manager, axes, **kwargs): events['iy'][i] = iy events[ybinning.name][i] = y - self._tdm.initialize_for_new_trial(src_hypo_group_manager, events) + self._tdm.initialize_trial(shg_mgr=src_hypo_group_manager, pmm=pmm, events=events) - event_pdf_values = self._pdf.get_prob(self._tdm) + (event_pdf_values, _) = self._pdf.get_pd(self._tdm) pdf_values[events['ix'], events['iy']] = event_pdf_values (left, right, bottom, top) = ( diff --git a/skyllh/plotting/i3/pdfratio.py b/skyllh/plotting/i3/pdfratio.py index 277f6b1713..886b0e3212 100644 --- a/skyllh/plotting/i3/pdfratio.py +++ b/skyllh/plotting/i3/pdfratio.py @@ -1,33 +1,37 @@ """Plotting module to plot IceCube specific PDF ratio objects.""" import itertools +from typing import cast import numpy as np from matplotlib.axes import Axes from matplotlib.colors import LogNorm +from matplotlib.image import AxesImage +from skyllh.core.parameters import ParameterModelMapper from skyllh.core.py import classname from skyllh.core.source_hypo_grouping import ( SourceHypoGroupManager, ) from skyllh.core.storage import DataFieldRecordArray from skyllh.core.trialdata import TrialDataManager +from skyllh.i3.pdf import I3EnergyPDF from skyllh.i3.pdfratio import SplinedI3EnergySigSetOverBkgPDFRatio class SplinedI3EnergySigSetOverBkgPDFRatioPlotter: - """Plotter class to plot an I3EnergySigSetOverBkgPDFRatioSpline object.""" + """Plotter class to plot an SplinedI3EnergySigSetOverBkgPDFRatio object.""" - def __init__(self, tdm, pdfratio): + def __init__(self, tdm: TrialDataManager, pdfratio: SplinedI3EnergySigSetOverBkgPDFRatio): """Creates a new plotter object for plotting an - I3EnergySigSetOverBkgPDFRatioSpline object. + SplinedI3EnergySigSetOverBkgPDFRatio object. Parameters ---------- - tdm : instance of TrialDataManager + tdm The instance of TrialDataManager that provides the data for the PDF ratio evaluation. - pdfratio : I3EnergySigSetOverBkgPDFRatioSpline + pdfratio The PDF ratio object to plot. """ self.tdm = tdm @@ -55,19 +59,30 @@ def tdm(self, obj): raise TypeError('The tdm property must be an instance of TrialDataManager!') self._tdm = obj - def plot(self, src_hypo_group_manager, axes, fitparams, **kwargs): + def plot( + self, + src_hypo_group_manager: SourceHypoGroupManager, + pmm: ParameterModelMapper, + src_params_recarray: np.ndarray, + axes: Axes, + **kwargs, + ) -> AxesImage: """Plots the PDF ratio for the given set of fit paramater values. Parameters ---------- - src_hypo_group_manager : instance of SourceHypoGroupManager + src_hypo_group_manager The instance of SourceHypoGroupManager that defines the source hypotheses. - axes : mpl.axes.Axes + pmm + The instance of ParameterModelMapper that maps the global parameters + to the local model parameters. + src_params_recarray + The numpy record ndarray holding the parameter names and values of + all sources. + axes The matplotlib Axes object on which the PDF ratio should get drawn to. - fitparams : dict - The dictionary with the set of fit paramater values. Additional Keyword Arguments ---------------------------- @@ -76,20 +91,18 @@ def plot(self, src_hypo_group_manager, axes, fitparams, **kwargs): Returns ------- - img : instance of mpl.AxesImage + img The AxesImage instance showing the PDF ratio image. """ if not isinstance(src_hypo_group_manager, SourceHypoGroupManager): raise TypeError('The src_hypo_group_manager argument must be an instance of SourceHypoGroupManager!') if not isinstance(axes, Axes): raise TypeError('The axes argument must be an instance of matplotlib.axes.Axes!') - if not isinstance(fitparams, dict): - raise TypeError('The fitparams argument must be an instance of dict!') # Get the binning for the axes. We use the background PDF to get it # from. By construction, all PDFs use the same binning. We know that # the PDFs are 2-dimensional. - (xbinning, ybinning) = self._pdfratio.backgroundpdf.binnings + (xbinning, ybinning) = cast(I3EnergyPDF, self._pdfratio.bkg_pdf).binnings # Create a 2D array with the ratio values. We put one event into each # bin. @@ -108,9 +121,9 @@ def plot(self, src_hypo_group_manager, axes, fitparams, **kwargs): events['iy'][i] = iy events[ybinning.name][i] = y - self._tdm.initialize_for_new_trial(src_hypo_group_manager, events) + self._tdm.initialize_trial(shg_mgr=src_hypo_group_manager, pmm=pmm, events=events) - event_ratios = self.pdfratio.get_ratio(self._tdm, fitparams) + event_ratios = self.pdfratio.get_ratio(self._tdm, src_params_recarray) for i in range(len(events)): ratios[events['ix'][i], events['iy'][i]] = event_ratios[i] diff --git a/skyllh/plotting/utils/trials.py b/skyllh/plotting/utils/trials.py index b852683edd..c64b16eb8a 100644 --- a/skyllh/plotting/utils/trials.py +++ b/skyllh/plotting/utils/trials.py @@ -1,89 +1,87 @@ -import matplotlib as mpl import numpy as np from matplotlib import ( pyplot as plt, ) +from matplotlib.colors import LogNorm +from matplotlib.figure import Figure from mpl_toolkits.axes_grid1.axes_divider import ( make_axes_locatable, ) def plot_ns_fit_vs_mean_ns_inj( - trials, - mean_n_sig_key='mean_n_sig', - ns_fit_key='ns', - rethist=False, - title='', - figsize=None, - line_color=None, - axis_fontsize=16, - title_fontsize=16, - tick_fontsize=16, - xlabel=None, - ylabel=None, - ylim=None, - ratio_ylim=None, -): + trials: np.ndarray, + mean_n_sig_key: str = 'mean_n_sig', + ns_fit_key: str = 'ns', + rethist: bool = False, + title: str = '', + figsize: tuple | None = None, + line_color: str | None = None, + axis_fontsize: float = 16, + title_fontsize: float = 16, + tick_fontsize: float = 16, + xlabel: str | None = None, + ylabel: str | None = None, + ylim: tuple | None = None, + ratio_ylim: tuple | None = None, +) -> tuple[Figure, np.ndarray, np.ndarray, np.ndarray] | Figure: r"""Creates a 2D histogram plot showing the fit number of signal events vs. the mean number of injected signal events. Parameters ---------- - trials : numpy record array + trials The record array holding the results of the trials. - mean_n_sig_key : str + mean_n_sig_key The name of the key for the mean number of injected signal events in the given trials record array. - Default is ``'mean_n_sig'``. - ns_fit_key : str + ns_fit_key The name of the key for the fitted number of signal events in the given trials record array. - Default is ``'ns'``. - rethist : bool + rethist If set to ``True``, the histogram data along with the histogram bin edges will be return as well. - Default is ``False``. - title : str + title The title of the plot. - figsize : tuple | None + figsize The two-element tuple (width,height) specifying the size of the figure. If set to None, the default size (12,10) will be used. - line_color : str | None + line_color The color of the lines. The default is '#E37222'. - axis_fontsize : float + axis_fontsize The font size of the axis labels. - title_fontsize : float + title_fontsize The font size of the plot title. - tick_fontsize : float + tick_fontsize The font size of the tick labels. - xlabel : str | None + xlabel The label of the x-axis in math syntax. Default is ``r'_{\mathrm{sig,inj}}}'``. - ylabel : str | None + ylabel The label of if y-axis in math syntax. Default is ``r'n_\mathrm{sig,fit}'``. - ylim : tuple | None + ylim The (low,high)-two-element tuple specifying the y-axis limits of the main plot. - ratio_ylim : tuple | None + ratio_ylim The (low,high)-two-element tuple specifying the y-axis limits of the ratio plot in percentage. If set to None, the default (-100,100) will be used. Returns ------- - fig : MPL Figure instance + fig The created matplotlib Figure instance. - hist : 2d ndarray + hist The histogram bin content. This will only be returned, when the ``rethist`` argument was set to ``True``. - xedges : 1d ndarray + xedges The histogram x-axis bin edges. This will only be returned, when the ``rethist`` argument was set to ``True``. - yedges : 1d ndarray + yedges The histogram y-axis bin edges. This will only be returned, when the ``rethist`` argument was set to ``True``. @@ -139,7 +137,7 @@ def plot_ns_fit_vs_mean_ns_inj( trials[ns_fit_key], bins=[x_bins, y_bins], weights=hist_weights, - norm=mpl.colors.LogNorm(), + norm=LogNorm(), cmap=plt.get_cmap('GnBu'), ) @@ -219,81 +217,78 @@ def plot_ns_fit_vs_mean_ns_inj( def plot_gamma_fit_vs_mean_ns_inj( - trials, - gamma_inj=2, - mean_n_sig_key='mean_n_sig', - gamma_fit_key='gamma', - rethist=False, - title='', - figsize=None, - line_color=None, - axis_fontsize=16, - title_fontsize=16, - tick_fontsize=16, - xlabel=None, - ylabel=None, - ratio_ylim=None, -): + trials: np.ndarray, + gamma_inj: float = 2, + mean_n_sig_key: str = 'mean_n_sig', + gamma_fit_key: str = 'gamma', + rethist: bool = False, + title: str = '', + figsize: tuple | None = None, + line_color: str | None = None, + axis_fontsize: float = 16, + title_fontsize: float = 16, + tick_fontsize: float = 16, + xlabel: str | None = None, + ylabel: str | None = None, + ratio_ylim: tuple | None = None, +) -> tuple[Figure, np.ndarray, np.ndarray, np.ndarray] | Figure: r"""Creates a 2D histogram plot showing the fit spectral index gamma vs. the mean number of injected signal events. Parameters ---------- - trials : numpy record array + trials The record array holding the results of the trials. - gamma_inj : float + gamma_inj The spectral index with which signal events got injected into tha trial data set. - mean_n_sig_key : str + mean_n_sig_key The name of the key for the mean number of injected signal events in the given trials record array. - Default is ``'mean_n_sig'``. - gamma_fit_key : str + gamma_fit_key The name of the key for the fitted spectral index in the given trials record array. - Default is ``'gamma'``. - rethist : bool + rethist If set to ``True``, the histogram data along with the histogram bin edges will be return as well. - Default is ``False``. - title : str + title The title of the plot. - figsize : tuple | None + figsize The two-element tuple (width,height) specifying the size of the figure. If set to None, the default size (12,10) will be used. - line_color : str | None + line_color The color of the lines. The default is '#E37222'. - axis_fontsize : float + axis_fontsize The font size of the axis labels. - title_fontsize : float + title_fontsize The font size of the plot title. - tick_fontsize : float + tick_fontsize The font size of the tick labels. - xlabel : str | None + xlabel The label of the x-axis in math syntax. Default is ``r'_{\mathrm{sig,inj}}}'``. - ylabel : str | None + ylabel The label of if y-axis in math syntax. Default is ``r'\gamma_\mathrm{fit}'``. - ratio_ylim : tuple | None + ratio_ylim The (low,high)-two-element tuple specifying the y-axis limits of the ratio plot in percentage. If set to None, the default (-100,100) will be used. Returns ------- - fig : MPL Figure instance + fig The created matplotlib Figure instance. - hist : 2d ndarray + hist The histogram bin content. This will only be returned, when the ``rethist`` argument was set to ``True``. - xedges : 1d ndarray + xedges The histogram x-axis bin edges. This will only be returned, when the ``rethist`` argument was set to ``True``. - yedges : 1d ndarray + yedges The histogram y-axis bin edges. This will only be returned, when the ``rethist`` argument was set to ``True``. @@ -349,7 +344,7 @@ def plot_gamma_fit_vs_mean_ns_inj( trials[gamma_fit_key], bins=[x_bins, y_bins], weights=hist_weights, - norm=mpl.colors.LogNorm(), + norm=LogNorm(), cmap=plt.get_cmap('GnBu'), ) diff --git a/skyllh/scripting/argparser.py b/skyllh/scripting/argparser.py index 7e94b3313d..b5c3fe3b15 100644 --- a/skyllh/scripting/argparser.py +++ b/skyllh/scripting/argparser.py @@ -5,20 +5,24 @@ import argparse -def create_argparser(description=None, options=True): +def create_argparser(description: str | None = None, options: bool | dict | None = True) -> argparse.ArgumentParser: """Creates an argparser with the given description and adds common options useful for analysis scripts. Parameters ---------- - description : str | None + description The description for the argparser. - options : bool | dict | None + options If set to None or False, no options will be added. If set to True, all common analysis script options will be added. If set to a dictionary, individual options can be turned on and off. See the :func:`add_argparser_options` for possible options. Default is ``True``. + + Returns + ------- + An instance of ArgumentParser with the given description and options. """ parser = argparse.ArgumentParser(description=description, formatter_class=argparse.RawTextHelpFormatter) @@ -32,48 +36,54 @@ def create_argparser(description=None, options=True): def add_argparser_options( - parser, config=True, data_basepath=True, debug_logfile=True, enable_tracing=True, n_cpu=True, seed=True -): + parser: argparse.ArgumentParser, + config: bool = True, + data_basepath: bool = True, + debug_logfile: bool = True, + enable_tracing: bool = True, + n_cpu: bool = True, + seed: bool = True, +) -> None: """Adds common argparser options to the given argparser instance, useful for analysis scripts. Parameters ---------- - parser : instance of ArgumentParser + parser The instance of ArgumentParser to which options should get added. - config : bool + config If set to ``True``, the ``--config`` option of type ``str`` will be added. It specifies the configuration file. The default value is ``None``. The option destination is ``config``. - data_basepath : bool + data_basepath If set to ``True``, the ``--data-basepath`` option of type ``str`` will be added. It specifies the base path to the data samples. The default value is ``None``. The option destination is ``data_basepath``. - debug_logfile : bool + debug_logfile If set to ``True``, the ``--debug-logfile`` option of type ``str`` will be added. If not ``None``, it specifies the log file for dubug information. The default value is ``None``. The option destination is ``debug_logfile``. - enable_tracing : bool + enable_tracing If set to ``True``, the ``--enable-tracing`` option of type ``bool`` will be added. If specified, enables the logging on the tracing level, i.e. a lot of DEBUG messages. The default value is ``False``. The option destination is ``enable_tracing``. - n_cpu : bool + n_cpu If set to ``True``, the ``--n-cpu`` option of type ``int`` will be added. It specifies the number of CPUs to utilize where parallelization is possible. The default value is ``1``. The option destination is ``n_cpu``. - seed : bool + seed If set to ``True``, the ``--seed`` option of type ``int`` will be added. It specifies the seed for the random number generator. diff --git a/tests/core/test_dataset.py b/tests/core/test_dataset.py index b4cf0d136c..99babfdd6b 100644 --- a/tests/core/test_dataset.py +++ b/tests/core/test_dataset.py @@ -83,6 +83,7 @@ def setUp(self): ) def test_transfer(self): + assert self.ds.origin is not None password = os.environ.get('ICECUBE_PASSWORD', None) if password is None: self.skipTest(f'No password for username "{self.ds.origin.username}" provided via the environment!') @@ -120,6 +121,7 @@ def setUp(self): ) def test_transfer(self): + assert self.ds.origin is not None password = os.environ.get('ICECUBE_PASSWORD', None) if password is None: self.skipTest(f'No password for username "{self.ds.origin.username}" provided via the environment!') @@ -152,6 +154,8 @@ def test_get_data_subset(self): livetime_data = Livetime(self.livetime_datafile) (dataset_data_subset, livetime_subset) = get_data_subset(dataset_data, livetime_data, t_start, t_end) + assert dataset_data_subset.exp is not None + assert dataset_data_subset.mc is not None self.assertEqual(len(dataset_data_subset.exp), 4) self.assertEqual(len(dataset_data_subset.mc), 4) self.assertAlmostEqual(livetime_subset.livetime, 1) @@ -163,6 +167,8 @@ def test_get_data_subset(self): livetime_data = Livetime(self.livetime_datafile) (dataset_data_subset, livetime_subset) = get_data_subset(dataset_data, livetime_data, t_start, t_end) + assert dataset_data_subset.exp is not None + assert dataset_data_subset.mc is not None self.assertEqual(len(dataset_data_subset.exp), 2) self.assertEqual(len(dataset_data_subset.mc), 2) self.assertAlmostEqual(livetime_subset.livetime, 0.5) @@ -172,6 +178,8 @@ def test_get_data_subset(self): t_end = 58444.75 (dataset_data_subset, livetime_subset) = get_data_subset(dataset_data, livetime_data, t_start, t_end) + assert dataset_data_subset.exp is not None + assert dataset_data_subset.mc is not None self.assertEqual(len(dataset_data_subset.exp), 3) self.assertEqual(len(dataset_data_subset.mc), 3) self.assertAlmostEqual(livetime_subset.livetime, 0.9) @@ -181,6 +189,8 @@ def test_get_data_subset(self): t_end = 58444.6 (dataset_data_subset, livetime_subset) = get_data_subset(dataset_data, livetime_data, t_start, t_end) + assert dataset_data_subset.exp is not None + assert dataset_data_subset.mc is not None self.assertEqual(len(dataset_data_subset.exp), 4) self.assertEqual(len(dataset_data_subset.mc), 4) self.assertAlmostEqual(livetime_subset.livetime, 0.85) @@ -190,6 +200,8 @@ def test_get_data_subset(self): t_end = 58444.6 (dataset_data_subset, livetime_subset) = get_data_subset(dataset_data, livetime_data, t_start, t_end) + assert dataset_data_subset.exp is not None + assert dataset_data_subset.mc is not None self.assertEqual(len(dataset_data_subset.exp), 3) self.assertEqual(len(dataset_data_subset.mc), 3) self.assertAlmostEqual(livetime_subset.livetime, 0.75) diff --git a/tests/core/test_event_selection.py b/tests/core/test_event_selection.py index ab1cb15f63..adf40ec732 100644 --- a/tests/core/test_event_selection.py +++ b/tests/core/test_event_selection.py @@ -35,7 +35,7 @@ def shgm_setup(n_sources=1): # Mock SourceHypoGroupManager class in order to pass isinstance checks and # set its properties used by event selection methods. shgm = Mock(spec_set=['__class__', 'source_list', 'n_sources']) - shgm.__class__ = SourceHypoGroupManager + shgm.__class__ = SourceHypoGroupManager # pyright: ignore[reportAttributeAccessIssue] rng = np.random.default_rng(0) x = rng.random((n_sources, 2)) @@ -74,6 +74,7 @@ def test_change_shg_mgr(self): shg_mgr = shgm_setup(n_sources=n_sources) evt_sel_method = AllEventSelectionMethod(shg_mgr) + assert evt_sel_method.shg_mgr is not None self.assertEqual( evt_sel_method.shg_mgr.source_list, shg_mgr.source_list, @@ -88,6 +89,7 @@ def test_change_shg_mgr(self): shg_mgr_new = shgm_setup(n_sources=n_sources) evt_sel_method.change_shg_mgr(shg_mgr_new) + assert evt_sel_method.shg_mgr is not None self.assertEqual( evt_sel_method.shg_mgr.source_list, shg_mgr_new.source_list, diff --git a/tests/core/test_interpolate.py b/tests/core/test_interpolate.py index c7df81cf8b..5acbd9ba17 100644 --- a/tests/core/test_interpolate.py +++ b/tests/core/test_interpolate.py @@ -90,10 +90,10 @@ def create_tdm(n_sources, n_selected_events): ) def tdm_broadcast_params_recarray_to_values_array(params_recarray): - return TrialDataManager.broadcast_params_recarray_to_values_array(tdm, params_recarray) + return TrialDataManager.broadcast_params_recarray_to_values_array(tdm, params_recarray) # pyright: ignore[reportAttributeAccessIssue] def tdm_broadcast_arrays_to_values_array(arrays): - return TrialDataManager.broadcast_arrays_to_values_array(tdm, arrays) + return TrialDataManager.broadcast_arrays_to_values_array(tdm, arrays) # pyright: ignore[reportAttributeAccessIssue] def tdm_broadcast_sources_array_to_values_array(*args, **kwargs): return TrialDataManager.broadcast_sources_array_to_values_array(tdm, *args, **kwargs) @@ -101,7 +101,7 @@ def tdm_broadcast_sources_array_to_values_array(*args, **kwargs): def tdm_broadcast_sources_arrays_to_values_arrays(*args, **kwargs): return TrialDataManager.broadcast_sources_arrays_to_values_arrays(tdm, *args, **kwargs) - tdm.__class__ = TrialDataManager + tdm.__class__ = TrialDataManager # pyright: ignore[reportAttributeAccessIssue] tdm.trial_data_state_id = 1 tdm.get_n_values = lambda: n_sources * n_selected_events tdm.src_evt_idxs = ( diff --git a/tests/core/test_logging.py b/tests/core/test_logging.py index 774bd5b19a..6cea2c178b 100644 --- a/tests/core/test_logging.py +++ b/tests/core/test_logging.py @@ -5,19 +5,21 @@ import tempfile import unittest +from skyllh.core.config import Config from skyllh.core.logging import setup_logging class SetupLoggingTestCase(unittest.TestCase): def setUp(self): # Configure a base log_format for testing - self.cfg = { - 'logging': { - 'log_level': 'INFO', - 'log_format': '%(levelname)s:%(name)s:%(message)s', - }, - 'project': {'working_directory': '.'}, - } + self.cfg = Config.from_dict( + { + 'logging': { + 'log_level': 'INFO', + 'log_format': '%(levelname)s:%(name)s:%(message)s', + }, + } + ) self.user_logger_name = 'skyllh.tests.setup_logging' self._reset_logger('skyllh') diff --git a/tests/core/test_model.py b/tests/core/test_model.py index eb74ab41c0..346c05c815 100644 --- a/tests/core/test_model.py +++ b/tests/core/test_model.py @@ -49,9 +49,9 @@ def test_cast(self): # Test that non-Model instances raises a TypeError. with self.assertRaises(TypeError): - modelcoll = ModelCollection.cast('A str instance.') + modelcoll = ModelCollection.cast('A str instance.') # pyright: ignore[reportArgumentType] with self.assertRaises(TypeError): - modelcoll = ModelCollection.cast(('str1', 'str2')) + modelcoll = ModelCollection.cast(('str1', 'str2')) # pyright: ignore[reportArgumentType] def test_model_type(self): self.assertTrue(issubclass(self.modelcoll.model_type, Model)) diff --git a/tests/core/test_parameters.py b/tests/core/test_parameters.py index 2c341d218c..fab74f67e1 100644 --- a/tests/core/test_parameters.py +++ b/tests/core/test_parameters.py @@ -103,11 +103,11 @@ def test_isfixed(self): def test_valmin(self): self.assertEqual(self.fixed_param.valmin, None) - np.testing.assert_almost_equal(self.floating_param.valmin, self.floating_param_valmin) + np.testing.assert_almost_equal(self.floating_param.valmin, self.floating_param_valmin) # pyright: ignore[reportArgumentType] def test_valmax(self): self.assertEqual(self.fixed_param.valmax, None) - np.testing.assert_almost_equal(self.floating_param.valmax, self.floating_param_valmax) + np.testing.assert_almost_equal(self.floating_param.valmax, self.floating_param_valmax) # pyright: ignore[reportArgumentType] def test_value(self): np.testing.assert_almost_equal(self.fixed_param.value, self.fixed_param_initial) @@ -164,8 +164,8 @@ def test_make_floating(self): ) np.testing.assert_almost_equal(self.fixed_param.initial, self.floating_param_initial) np.testing.assert_almost_equal(self.fixed_param.value, self.floating_param_initial) - np.testing.assert_almost_equal(self.fixed_param.valmin, self.floating_param_valmin) - np.testing.assert_almost_equal(self.fixed_param.valmax, self.floating_param_valmax) + np.testing.assert_almost_equal(self.fixed_param.valmin, self.floating_param_valmin) # pyright: ignore[reportArgumentType] + np.testing.assert_almost_equal(self.fixed_param.valmax, self.floating_param_valmax) # pyright: ignore[reportArgumentType] class ParameterSet_TestCase(unittest.TestCase): @@ -351,7 +351,7 @@ def test_copy(self): def test_add_param(self): with self.assertRaises(TypeError): - self.paramset.add_param('p2') + self.paramset.add_param('p2') # pyright: ignore[reportArgumentType] with self.assertRaises(KeyError): param = Parameter('p0', 42.0) self.paramset.add_param(param) @@ -419,16 +419,16 @@ def test_round_to_nearest_grid_point(self): np.testing.assert_almost_equal(gp, [3.5]) # Test a value between two grid points. - x = [2.1, 2.4, 2.2, 2.3] + x = np.array([2.1, 2.4, 2.2, 2.3]) gp = self.paramgrid_gamma1.round_to_nearest_grid_point(x) np.testing.assert_almost_equal(gp, [2.0, 2.5, 2.0, 2.5]) - x = [1.051, 1.14] + x = np.array([1.051, 1.14]) gp = self.paramgrid_gamma3.round_to_nearest_grid_point(x) np.testing.assert_almost_equal(gp, [1.05, 1.15]) # Test a value on a grid point. - x = [1.05, 1.35] + x = np.array([1.05, 1.35]) gp = self.paramgrid_gamma3.round_to_nearest_grid_point(x) np.testing.assert_almost_equal(gp, [1.05, 1.35]) @@ -447,7 +447,7 @@ def test_round_to_lower_grid_point(self): gp = self.paramgrid_gamma2.round_to_lower_grid_point(x) np.testing.assert_almost_equal(gp, 1.6) - x = [1.05, 1.15, 1.25, 1.35] + x = np.array([1.05, 1.15, 1.25, 1.35]) gp = self.paramgrid_gamma3.round_to_lower_grid_point(x) np.testing.assert_almost_equal(gp, [1.05, 1.15, 1.25, 1.35]) @@ -466,7 +466,7 @@ def test_round_to_upper_grid_point(self): gp = self.paramgrid_gamma2.round_to_upper_grid_point(x) np.testing.assert_almost_equal(gp, 1.7) - x = [1.05, 1.15, 1.25, 1.35] + x = np.array([1.05, 1.15, 1.25, 1.35]) gp = self.paramgrid_gamma3.round_to_upper_grid_point(x) np.testing.assert_almost_equal(gp, [1.15, 1.25, 1.35, 1.45]) diff --git a/tests/core/test_signal_generator.py b/tests/core/test_signal_generator.py index 1696c58e47..3f4ef58b37 100644 --- a/tests/core/test_signal_generator.py +++ b/tests/core/test_signal_generator.py @@ -156,6 +156,7 @@ def testSigCandidatesArray(self): self.assertTrue(isinstance(arr, np.ndarray)) # Check field names. + assert arr.dtype.fields is not None field_names = arr.dtype.fields.keys() self.assertTrue( ('ds_idx' in field_names) @@ -237,10 +238,10 @@ class TestMultiDatasetSignalGeneratorEnergyRangeConsistency(unittest.TestCase): @staticmethod def _make_sig_gen(configured_ranges): sig_gen = MultiDatasetSignalGenerator.__new__(MultiDatasetSignalGenerator) - sig_gen._shg_mgr = _DummyShgMgr(n_sources=2) + sig_gen._shg_mgr = _DummyShgMgr(n_sources=2) # pyright: ignore[reportAttributeAccessIssue] sig_gen._src_params_recarray = np.zeros((2,), dtype=_GAMMA_RECARRAY_DTYPE) src_service = _DummySrcDetSigYieldWeightsService(a_jk=np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float64)) - sig_gen._ds_sig_weight_factors_service = _DummyDatasetSignalWeightFactorsService(src_service=src_service) + sig_gen._ds_sig_weight_factors_service = _DummyDatasetSignalWeightFactorsService(src_service=src_service) # pyright: ignore[reportAttributeAccessIssue] sig_gen._sig_generator_list = [ _DummyDatasetSigGenerator(configured_energy_range=r, correction_factors=[1.0, 1.0]) for r in configured_ranges @@ -280,7 +281,7 @@ def test_generate_signal_events_rejects_energy_range_kwarg(self): sig_gen._src_params_recarray = np.zeros((1,), dtype=_GAMMA_RECARRAY_DTYPE) src_service = _DummySrcDetSigYieldWeightsService(a_jk=np.array([[1.0], [1.0]], dtype=np.float64)) - sig_gen._ds_sig_weight_factors_service = _DummyDatasetSignalWeightFactorsService( + sig_gen._ds_sig_weight_factors_service = _DummyDatasetSignalWeightFactorsService( # pyright: ignore[reportAttributeAccessIssue] src_service=src_service, ds_weights=np.array([0.5, 0.5], dtype=np.float64), ) @@ -302,7 +303,7 @@ class _DummyRSS: rss = _DummyRSS() with self.assertRaises(TypeError): - sig_gen.generate_signal_events(rss=rss, mean=10, poisson=False, energy_range=(1e2, 1e3)) + sig_gen.generate_signal_events(rss=rss, mean=10, poisson=False, energy_range=(1e2, 1e3)) # pyright: ignore[reportCallIssue] self.assertIsNone(ds_gen_0.last_generate_kwargs) self.assertIsNone(ds_gen_1.last_generate_kwargs) diff --git a/tests/core/test_signalpdf.py b/tests/core/test_signalpdf.py index 85bd34e2e3..59fcab6ceb 100644 --- a/tests/core/test_signalpdf.py +++ b/tests/core/test_signalpdf.py @@ -53,7 +53,7 @@ def tdm_get_data(key): return np.array([0, 5, 9.7]) raise ValueError(f'Value n_selected_events={n_selected_events} is not supported!') - tdm.__class__ = TrialDataManager + tdm.__class__ = TrialDataManager # pyright: ignore[reportAttributeAccessIssue] tdm.trial_data_state_id = 1 tdm.get_n_values = lambda: n_sources * n_selected_events tdm.src_evt_idxs = ( @@ -102,7 +102,7 @@ def test__calculate_sum_of_ontime_time_flux_profile_integrals(self): def test_get_pd(self): tdm = create_tdm(n_sources=self.pmm.n_sources, n_selected_events=3) - src_params_recarray = self.pmm.create_src_params_recarray(gflp_values=[]) + src_params_recarray = self.pmm.create_src_params_recarray(gflp_values=np.array([])) (pd, grads) = self.sig_time_pdf.get_pd(tdm=tdm, params_recarray=src_params_recarray) diff --git a/tests/core/test_source_model.py b/tests/core/test_source_model.py index cdf9f07fd3..3b5bb93514 100644 --- a/tests/core/test_source_model.py +++ b/tests/core/test_source_model.py @@ -36,8 +36,8 @@ def setUp(self): self.dec = 1 def test_SourceModelCollection(self): - source_model1 = SourceModel(self.ra, self.dec) - source_model2 = SourceModel(self.ra, self.dec) + source_model1 = SourceModel(self.ra, self.dec) # pyright: ignore[reportArgumentType] + source_model2 = SourceModel(self.ra, self.dec) # pyright: ignore[reportArgumentType] source_collection_casted = SourceModelCollection.cast( source_model1, 'Could not cast SourceModel to SourceCollection' @@ -57,8 +57,8 @@ def setUp(self): self.name = 'MySourceCatalog' self.ra = 0.1 self.dec = 1.1 - self.source1 = SourceModel(self.ra, self.dec) - self.source2 = SourceModel(self.ra, self.dec) + self.source1 = SourceModel(self.ra, self.dec) # pyright: ignore[reportArgumentType] + self.source2 = SourceModel(self.ra, self.dec) # pyright: ignore[reportArgumentType] self.catalog = SourceCatalog(name=self.name, sources=[self.source1, self.source2], source_type=SourceModel) diff --git a/tests/core/test_weights.py b/tests/core/test_weights.py index b3ed2f5f2d..fcc168875e 100644 --- a/tests/core/test_weights.py +++ b/tests/core/test_weights.py @@ -41,13 +41,13 @@ class SimpleDetSigYieldWithoutGrads(DetSigYield): def __init__(self, scale=1, **kwargs): self._scale = scale - def sources_to_recarray(self, sources): - recarr = np.empty((len(sources),), dtype=[('dec', np.double)]) - for i, src in enumerate(sources): + def sources_to_recarray(self, sources): # pyright: ignore[reportIncompatibleMethodOverride] + recarr = np.empty((len(sources),), dtype=[('dec', np.double)]) # pyright: ignore[reportArgumentType] + for i, src in enumerate(sources): # pyright: ignore[reportArgumentType] recarr[i]['dec'] = src.dec return recarr - def __call__(self, src_recarray, src_params_recarray): + def __call__(self, src_recarray, src_params_recarray): # pyright: ignore[reportIncompatibleMethodOverride] """ Parameters ---------- @@ -162,7 +162,7 @@ class NoDetSigYieldBuilder(DetSigYieldBuilder): def __init__(self, **kwargs): super().__init__(**kwargs) - def construct_detsigyield(self, **kwargs): + def construct_detsigyield(self, **kwargs): # pyright: ignore[reportIncompatibleMethodOverride] pass @@ -207,7 +207,7 @@ def create_DetSigYieldService(shg_mgr, detsigyield_arr): ] ) - detsigyield_service.__class__ = DetSigYieldService + detsigyield_service.__class__ = DetSigYieldService # pyright: ignore[reportAttributeAccessIssue] detsigyield_service.arr = detsigyield_arr detsigyield_service.shg_mgr = shg_mgr detsigyield_service.n_datasets = detsigyield_arr.shape[0] @@ -238,6 +238,8 @@ def test_without_grads(self): src_params_recarray = type(self)._pmm.create_src_params_recarray(gflp_values) src_detsigyield_weights_service.calculate(src_params_recarray) (a_jk, a_jk_grads) = src_detsigyield_weights_service.get_weights() + assert a_jk is not None + assert a_jk_grads is not None self.assertIsInstance(a_jk, np.ndarray, 'instance of a_jk') @@ -266,6 +268,8 @@ def test_with_grads_p1(self): src_params_recarray = type(self)._pmm.create_src_params_recarray(gflp_values) src_detsigyield_weights_service.calculate(src_params_recarray) (a_jk, a_jk_grads) = src_detsigyield_weights_service.get_weights() + assert a_jk is not None + assert a_jk_grads is not None self.assertIsInstance(a_jk, np.ndarray, 'instance of a_jk') @@ -300,6 +304,8 @@ def test_with_grads_p2(self): src_params_recarray = type(self)._pmm.create_src_params_recarray(gflp_values) src_detsigyield_weights_service.calculate(src_params_recarray) (a_jk, a_jk_grads) = src_detsigyield_weights_service.get_weights() + assert a_jk is not None + assert a_jk_grads is not None self.assertIsInstance(a_jk, np.ndarray, 'instance of a_jk') @@ -347,6 +353,8 @@ def test_without_grads(self): src_detsigyield_weights_service.calculate(src_params_recarray) ds_sig_weight_factors_service.calculate() (f_j, f_j_grads) = ds_sig_weight_factors_service.get_weights() + assert f_j is not None + assert f_j_grads is not None self.assertIsInstance(f_j, np.ndarray, 'instance of f_j') @@ -380,6 +388,8 @@ def test_with_grads_p1(self): src_detsigyield_weights_service.calculate(src_params_recarray) ds_sig_weight_factors_service.calculate() (f_j, f_j_grads) = ds_sig_weight_factors_service.get_weights() + assert f_j is not None + assert f_j_grads is not None self.assertIsInstance(f_j, np.ndarray, 'instance of f_j') @@ -417,6 +427,8 @@ def test_with_grads_p2(self): src_detsigyield_weights_service.calculate(src_params_recarray) ds_sig_weight_factors_service.calculate() (f_j, f_j_grads) = ds_sig_weight_factors_service.get_weights() + assert f_j is not None + assert f_j_grads is not None self.assertIsInstance(f_j, np.ndarray, 'instance of f_j') diff --git a/tests/core/testdata/testdata_generator.py b/tests/core/testdata/testdata_generator.py index 6aed7a9952..1de83b19a3 100644 --- a/tests/core/testdata/testdata_generator.py +++ b/tests/core/testdata/testdata_generator.py @@ -136,6 +136,6 @@ def generate_testdata(): if __name__ == '__main__': testdata = generate_testdata() - np.save('exp_testdata.npy', testdata.get('exp_testdata')) - np.save('mc_testdata.npy', testdata.get('mc_testdata')) - np.save('livetime_testdata.npy', testdata.get('livetime_testdata')) + np.save('exp_testdata.npy', testdata['exp_testdata']) + np.save('mc_testdata.npy', testdata['mc_testdata']) + np.save('livetime_testdata.npy', testdata['livetime_testdata']) diff --git a/tests/i3/test_scrambling.py b/tests/i3/test_scrambling.py index 2952c00fdf..5385a80417 100644 --- a/tests/i3/test_scrambling.py +++ b/tests/i3/test_scrambling.py @@ -60,7 +60,7 @@ def test_scramble(self): i3timescramblingmethod = I3TimeScramblingMethod(timegen) rss = RandomStateService(seed=1) - data = i3timescramblingmethod.scramble(rss=rss, dataset=None, data=exp_data) + data = i3timescramblingmethod.scramble(rss=rss, dataset=None, data=exp_data) # pyright: ignore[reportArgumentType] np.testing.assert_allclose(data['time'], self.scrambled_data['time']) np.testing.assert_allclose(data['ra'], self.scrambled_data['ra']) diff --git a/tests/i3/testdata/testdata_generator.py b/tests/i3/testdata/testdata_generator.py index 3c595402e3..838303cb86 100644 --- a/tests/i3/testdata/testdata_generator.py +++ b/tests/i3/testdata/testdata_generator.py @@ -127,6 +127,6 @@ def generate_testdata(): if __name__ == '__main__': testdata = generate_testdata() - np.save('exp_testdata.npy', testdata.get('exp_testdata')) - np.save('mc_testdata.npy', testdata.get('mc_testdata')) - np.save('grl_testdata.npy', testdata.get('grl_testdata')) + np.save('exp_testdata.npy', testdata['exp_testdata']) + np.save('mc_testdata.npy', testdata['mc_testdata']) + np.save('grl_testdata.npy', testdata['grl_testdata']) diff --git a/tests/publicdata_ps/test_time_integrated_ps_dr1.py b/tests/publicdata_ps/test_time_integrated_ps_dr1.py index 9749ccb508..f1d6d06c09 100644 --- a/tests/publicdata_ps/test_time_integrated_ps_dr1.py +++ b/tests/publicdata_ps/test_time_integrated_ps_dr1.py @@ -8,6 +8,7 @@ from skyllh.core.config import Config from skyllh.core.logging import setup_logging from skyllh.core.random import RandomStateService +from skyllh.core.signal_generator import MultiDatasetSignalGenerator from skyllh.core.source_model import PointLikeSource from skyllh.core.timing import TimeLord from skyllh.datasets.i3 import PublicData_10y_ps @@ -182,6 +183,8 @@ def setUpClass(cls): source=cls.source, ) cls.ana_direct_sig_gen.construct_signal_generator() + assert isinstance(cls.ana_direct_sig_gen._sig_generator, MultiDatasetSignalGenerator) + assert cls.ana_direct_sig_gen._sig_generator.sig_generator_list is not None for gen in cls.ana_direct_sig_gen._sig_generator.sig_generator_list: if hasattr(gen, 'energy_range'): gen.energy_range = cls.ENERGY_RANGE