Skip to content
Snippets Groups Projects
Commit b847e375 authored by Mikhail Smirnov's avatar Mikhail Smirnov
Browse files

some changes

parent 4e4ea156
Branches 1-addition-to-readme
Tags v3.2.3
3 merge requests!8another-dev to main,!7Another dev,!6Another dev
Pipeline #31429 passed
......@@ -38,6 +38,7 @@ install_requires =
importlib_resources
prettytable
astropy
setuptools
setuptools_scm
uproot
python_requires = >=3.6
......
......@@ -15,11 +15,11 @@ import astropy.units as u
# from gammapy.irf import EnergyDispersion2D
from scipy.stats import binned_statistic
# from scipy.ndimage import gaussian_filter1d, gaussian_filter
# from scipy.ndimage import gaussian_filter1d, gaussian_filter
# from collections import defaultdict
# from collections import defaultdict
# import sys
# sys.path.append('../')
......@@ -28,71 +28,68 @@ from scipy.stats import binned_statistic
# from python_scripts.func import WriteAeff
# from python_scripts.func import WritePSF
def build_aeff(
input=(file_nu, file_nubar),
no_bdt=False,
cuts=False,
power_degree=-2.5
):
input=(file_nu, file_nubar), no_bdt=False, cuts=False, power_degree=-2.5
):
"""
Create Aeff .fits from dist files
Create Aeff .fits from dist files
input: a tuple with pathes to nu_dst file and nubar_dst file
input: a tuple with pathes to nu_dst file and nubar_dst file
no_bdt: include or exclude bdt, default False
no_bdt: include or exclude bdt, default False
cuts: apply cuts, default False
cuts: apply cuts, default False
power_degree: re-weight data, default value -2.5
power_degree: re-weight data, default value -2.5
"""
#Read data files using km3io
# Read data files using km3io
f_nu_km3io = OfflineReader(input[0])
f_nubar_km3io = OfflineReader(input[1])
#Read data files using uproot
# Read data files using uproot
f_nu_uproot = ur.open(input[0])
f_nubar_uproot = ur.open(input[1])
def unpack_data(no_bdt):
"""
some words
some words
return tuple with two pandas data frames (nu, nubar)
return tuple with two pandas data frames (nu, nubar)
"""
# Access data arrays
data_km3io = dict()
for l,f in zip(['nu', 'nubar'], [f_nu_km3io, f_nubar_km3io]):
for l, f in zip(["nu", "nubar"], [f_nu_km3io, f_nubar_km3io]):
data_km3io[l] = dict()
data_km3io[l]['E'] = f.tracks.E[:,0]
data_km3io[l]['dir_x'] = f.tracks.dir_x[:,0]
data_km3io[l]['dir_y'] = f.tracks.dir_y[:,0]
data_km3io[l]['dir_z'] = f.tracks.dir_z[:,0]
data_km3io[l]["E"] = f.tracks.E[:, 0]
data_km3io[l]["dir_x"] = f.tracks.dir_x[:, 0]
data_km3io[l]["dir_y"] = f.tracks.dir_y[:, 0]
data_km3io[l]["dir_z"] = f.tracks.dir_z[:, 0]
data_km3io[l]['energy_mc'] = f.mc_tracks.E[:,0]
data_km3io[l]['dir_x_mc'] = f.mc_tracks.dir_x[:,0]
data_km3io[l]['dir_y_mc'] = f.mc_tracks.dir_y[:,0]
data_km3io[l]['dir_z_mc'] = f.mc_tracks.dir_z[:,0]
data_km3io[l]["energy_mc"] = f.mc_tracks.E[:, 0]
data_km3io[l]["dir_x_mc"] = f.mc_tracks.dir_x[:, 0]
data_km3io[l]["dir_y_mc"] = f.mc_tracks.dir_y[:, 0]
data_km3io[l]["dir_z_mc"] = f.mc_tracks.dir_z[:, 0]
data_km3io[l]['weight_w2'] = f.w[:,1]
data_km3io[l]["weight_w2"] = f.w[:, 1]
#extracting bdt information
# extracting bdt information
if not no_bdt:
for l,f in zip(['nu', 'nubar'], [f_nu_uproot, f_nubar_uproot]):
T = f['T;1']
bdt = T['bdt'].array()
data_km3io[l]['bdt0'] = bdt[:,0]
data_km3io[l]['bdt1'] = bdt[:,1]
for l, f in zip(["nu", "nubar"], [f_nu_uproot, f_nubar_uproot]):
T = f["T;1"]
bdt = T["bdt"].array()
data_km3io[l]["bdt0"] = bdt[:, 0]
data_km3io[l]["bdt1"] = bdt[:, 1]
# create Data Frames
df_nu = pd.DataFrame(data_km3io['nu'])
df_nubar = pd.DataFrame(data_km3io['nubar'])
df_nu = pd.DataFrame(data_km3io["nu"])
df_nubar = pd.DataFrame(data_km3io["nubar"])
data_tuple = (df_nu, df_nubar)
return data_tuple
\ No newline at end of file
return data_tuple
......@@ -3,12 +3,13 @@
import numpy as np
import pandas as pd
import numba as nb
from numba import jit, prange
from numba import jit, prange
from km3pipe.math import azimuth, zenith
import astropy.coordinates as ac
from astropy.time import Time
import astropy.units as u
def calc_theta(table, mc=True):
"""
Calculate the zenith angle theta of events given the direction coordinates x,y,z
......@@ -25,56 +26,59 @@ def calc_theta(table, mc=True):
"""
if not mc:
dir_x=table['dir_x'].to_numpy()
dir_y=table['dir_y'].to_numpy()
dir_z=table['dir_z'].to_numpy()
dir_x = table["dir_x"].to_numpy()
dir_y = table["dir_y"].to_numpy()
dir_z = table["dir_z"].to_numpy()
else:
dir_x=table['dir_x_mc'].to_numpy()
dir_y=table['dir_y_mc'].to_numpy()
dir_z=table['dir_z_mc'].to_numpy()
dir_x = table["dir_x_mc"].to_numpy()
dir_y = table["dir_y_mc"].to_numpy()
dir_z = table["dir_z_mc"].to_numpy()
nu_directions=np.vstack([dir_x, dir_y, dir_z]).T
nu_directions = np.vstack([dir_x, dir_y, dir_z]).T
theta = zenith(nu_directions) # zenith angles in rad [0:pi]
theta = zenith(nu_directions) # zenith angles in rad [0:pi]
return theta
def edisp_3D(e_bins, m_bins, t_bins, dataset, weights=1):
"""
Calculate the 3-dimensional energy dispersion matrix.
This is just a historgram with the simulated events.
The normalization needs to be done afterwards.
Calculate the 3-dimensional energy dispersion matrix.
This is just a historgram with the simulated events.
The normalization needs to be done afterwards.
e_bins: Array of energy bins in GeV
e_bins: Array of energy bins in GeV
m_bins: Array of energy migration bins (reconstructed energy / true energy)
m_bins: Array of energy migration bins (reconstructed energy / true energy)
t_bins: Array of zenith angle bins in rad
t_bins: Array of zenith angle bins in rad
dataset: pandas.Dataframe with the events
dataset: pandas.Dataframe with the events
weights: Array of weights for each event
weights: Array of weights for each event
------
------
Returns: 3-dimensional energy dispersion matrix. Binned in energy, energy migration and zenith angle
Returns: 3-dimensional energy dispersion matrix. Binned in energy, energy migration and zenith angle
"""
if 'theta_mc' not in dataset.keys():
dataset['theta_mc'] = calc_theta(dataset, mc=True)
if 'migra' not in dataset.keys():
dataset['migra'] = dataset.E / dataset.energy_mc
if "theta_mc" not in dataset.keys():
dataset["theta_mc"] = calc_theta(dataset, mc=True)
if "migra" not in dataset.keys():
dataset["migra"] = dataset.E / dataset.energy_mc
theta_bins = pd.cut(dataset.theta_mc, t_bins, labels=False).to_numpy()
energy_bins = pd.cut(dataset.energy_mc, e_bins, labels=False).to_numpy()
migra_bins = pd.cut(dataset.migra, m_bins, labels=False).to_numpy()
edisp = fill_edisp_3D(e_bins, m_bins, t_bins, energy_bins, migra_bins, theta_bins, weights)
edisp = fill_edisp_3D(
e_bins, m_bins, t_bins, energy_bins, migra_bins, theta_bins, weights
)
return edisp
@jit(nopython=True, fastmath=False, parallel=True)
def fill_edisp_3D(e_bins, m_bins, t_bins, energy_bins, migra_bins, theta_bins, weights):
"""
......@@ -82,70 +86,79 @@ def fill_edisp_3D(e_bins, m_bins, t_bins, energy_bins, migra_bins, theta_bins, w
Needed because numba does not work with pandas but needs numpy arrays.
fastmath is disabled because it gives different results.
"""
edisp = np.zeros((len(t_bins)-1, len(m_bins)-1, len(e_bins)-1))
for i in prange(len(t_bins)-1):
for j in range(len(m_bins)-1):
for k in range(len(e_bins)-1):
edisp = np.zeros((len(t_bins) - 1, len(m_bins) - 1, len(e_bins) - 1))
for i in prange(len(t_bins) - 1):
for j in range(len(m_bins) - 1):
for k in range(len(e_bins) - 1):
mask = (energy_bins == k) & (migra_bins == j) & (theta_bins == i)
edisp[i,j,k] = np.sum(mask * weights)
edisp[i, j, k] = np.sum(mask * weights)
return edisp
def psf_3D(e_bins, r_bins, t_bins, dataset, weights=1):
"""
Calculate the 3-dimensional PSF matrix.
This is a historgram with the simulated evenets.
The normalization needs to be done afterwards.
Calculate the 3-dimensional PSF matrix.
This is a historgram with the simulated evenets.
The normalization needs to be done afterwards.
e_bins: Array of energy bins in GeV
e_bins: Array of energy bins in GeV
r_bins: Array of radial bins in deg (angle between true and reconstructed direction)
r_bins: Array of radial bins in deg (angle between true and reconstructed direction)
t_bins: Array of zenith angle bins in rad
t_bins: Array of zenith angle bins in rad
dataset: pandas.Dataframe with the events
dataset: pandas.Dataframe with the events
weights: Array of weights for each event
weights: Array of weights for each event
------
------
Returns: 3-dimensional PSF matrix. Binned in energy, zenith angle and radial bins
Returns: 3-dimensional PSF matrix. Binned in energy, zenith angle and radial bins
"""
if 'theta_mc' not in dataset.keys():
dataset['theta_mc'] = calc_theta(dataset, mc=True)
if "theta_mc" not in dataset.keys():
dataset["theta_mc"] = calc_theta(dataset, mc=True)
scalar_prod = dataset.dir_x*dataset.dir_x_mc + dataset.dir_y*dataset.dir_y_mc + dataset.dir_z*dataset.dir_z_mc
scalar_prod[scalar_prod>1.0] = 1.0
rad = np.arccos(scalar_prod) * 180/np.pi # in deg now
dataset['rad'] = rad
scalar_prod = (
dataset.dir_x * dataset.dir_x_mc
+ dataset.dir_y * dataset.dir_y_mc
+ dataset.dir_z * dataset.dir_z_mc
)
scalar_prod[scalar_prod > 1.0] = 1.0
rad = np.arccos(scalar_prod) * 180 / np.pi # in deg now
dataset["rad"] = rad
theta_bins = pd.cut(dataset.theta_mc, t_bins, labels=False).to_numpy()
energy_bins = pd.cut(dataset.energy_mc, e_bins, labels=False).to_numpy()
rad_bins = pd.cut(rad, r_bins, labels=False).to_numpy()
psf = fill_psf_3D(e_bins, r_bins, t_bins, energy_bins, rad_bins, theta_bins, weights)
psf = fill_psf_3D(
e_bins, r_bins, t_bins, energy_bins, rad_bins, theta_bins, weights
)
return psf
@jit(nopython=True, fastmath=False, parallel=True) # do not use fastmath=True -- gives 300 aditional entries and no gain
@jit(
nopython=True, fastmath=False, parallel=True
) # do not use fastmath=True -- gives 300 aditional entries and no gain
def fill_psf_3D(e_bins, r_bins, t_bins, energy_bins, rad_bins, theta_bins, weights):
psf = np.zeros((len(r_bins)-1, len(t_bins)-1, len(e_bins)-1))
for j in prange(len(r_bins)-1):
for i in range(len(t_bins)-1):
for k in range(len(e_bins)-1):
psf = np.zeros((len(r_bins) - 1, len(t_bins) - 1, len(e_bins) - 1))
for j in prange(len(r_bins) - 1):
for i in range(len(t_bins) - 1):
for k in range(len(e_bins) - 1):
mask = (energy_bins == k) & (rad_bins == j) & (theta_bins == i)
psf[j,i,k] = np.sum(mask * weights)
psf[j, i, k] = np.sum(mask * weights)
return psf
def aeff_2D(e_bins, t_bins, dataset, gamma=1.4, nevents=2e7):
"""
Calculate the effective area in energy and zenith angle bins.
Calculate the effective area in energy and zenith angle bins.
e_bins: Array of energy bins in GeV
......@@ -163,8 +176,8 @@ def aeff_2D(e_bins, t_bins, dataset, gamma=1.4, nevents=2e7):
"""
if 'theta_mc' not in dataset.keys():
dataset['theta_mc'] = calc_theta(dataset, mc=True)
if "theta_mc" not in dataset.keys():
dataset["theta_mc"] = calc_theta(dataset, mc=True)
theta_bins = pd.cut(dataset.theta_mc, t_bins, labels=False).to_numpy()
energy_bins = pd.cut(dataset.energy_mc, e_bins, labels=False).to_numpy()
......@@ -175,15 +188,20 @@ def aeff_2D(e_bins, t_bins, dataset, gamma=1.4, nevents=2e7):
return aeff
@jit(nopython=True, fastmath=False, parallel=True)
def fill_aeff_2D(e_bins, t_bins, energy_bins, theta_bins, w2, E, gamma, nevents):
T = 365*24*3600
aeff = np.empty((len(e_bins)-1, len(t_bins)-1))
for k in prange(len(e_bins)-1):
for i in range(len(t_bins)-1):
T = 365 * 24 * 3600
aeff = np.empty((len(e_bins) - 1, len(t_bins) - 1))
for k in prange(len(e_bins) - 1):
for i in range(len(t_bins) - 1):
mask = (energy_bins == k) & (theta_bins == i)
d_omega = -(np.cos(t_bins[i+1]) - np.cos(t_bins[i]))
d_E = (e_bins[k+1])**(1-gamma) - (e_bins[k])**(1-gamma)
aeff[k,i] = (1-gamma) * np.sum(E[mask]**(-gamma)*w2[mask]) / (T*d_omega*d_E*nevents*2*np.pi)
d_omega = -(np.cos(t_bins[i + 1]) - np.cos(t_bins[i]))
d_E = (e_bins[k + 1]) ** (1 - gamma) - (e_bins[k]) ** (1 - gamma)
aeff[k, i] = (
(1 - gamma)
* np.sum(E[mask] ** (-gamma) * w2[mask])
/ (T * d_omega * d_E * nevents * 2 * np.pi)
)
return aeff
......@@ -36,7 +36,7 @@ def merge_fits(
edisp_fits: path to Edisp .fits file
bkg_fits: path to Background .fits file
output_file: name of the merged .fits file in data foledr of the package
"""
hdu_list = []
......
......@@ -5,10 +5,7 @@ from os import path, listdir, curdir, remove
import astropy
from km3irf import Calculator
from km3irf.utils import (
merge_fits,
list_data
)
from km3irf.utils import merge_fits, list_data
class TestCalculator(unittest.TestCase):
......@@ -40,6 +37,7 @@ class TestCalculator(unittest.TestCase):
for b in range(100):
assert 0 == self.calculator.multiply(0, b)
class TestUtils(unittest.TestCase):
def setUp(self):
self.test_path = path.join(path.abspath(curdir), "src", "km3irf", "data")
......@@ -49,4 +47,4 @@ class TestUtils(unittest.TestCase):
assert "all_in_one.fits" in listdir(self.test_path)
def test_list_data(self):
assert len(list_data()) == len(listdir(self.test_path))
\ No newline at end of file
assert len(list_data()) == len(listdir(self.test_path))
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment