Skip to content
Snippets Groups Projects
Commit 537992f2 authored by Stefan Reck's avatar Stefan Reck
Browse files

Merge branch 'v2.0.5'

parents 29148c6e 8bc1d256
No related branches found
No related tags found
No related merge requests found
#!/usr/bin/env python
# -*- coding: utf-8 -*-
import numpy as np
import km3pipe as kp
import matplotlib.pyplot as plt
import orcasong.modules as modules
__author__ = 'Stefan Reck'
class TimePlotter:
"""
For plotting the time distribution of hits, in a histogram,
and finding a good binning.
The plot will have some bins attached in both directions for
better overview.
Attributes:
-----------
files : list
The .h5 files to read hits from.
do_mc_hits : bool
Read mchits instead of hits.
det_file : str, optional
Path to a .detx detector geometry file, which can be used to
calibrate the hits.
center_time : bool
Subtract time of first triggered hit from all hit times. Will
also be done for McHits if they are in the blob [default: True].
add_t0 : bool
If true, add t0 to the time of hits. If using a det_file,
this will already have been done automatically [default: False].
inactive_du : int, optional
Dont plot hits from this du.
"""
def __init__(self, files,
do_mchits=False,
det_file=None,
center_time=True,
add_t0=False,
inactive_du=None):
if isinstance(files, str):
self.files = [files]
else:
self.files = files
self.do_mchits = do_mchits
self.det_file = det_file
self.add_t0 = add_t0
self.center_time = center_time
self.inactive_du = inactive_du
self.data = np.array([])
for file in self.files:
self._read(file, self.det_file)
def hist(self, bins=50, padding=0.5, **kwargs):
"""
Plot the hits as a histogram.
Parameters
----------
bins : int or np.array
Number of bins, or bin edges. If bin edges are given, some
bins are attached left and right (next to red vertical lines)
for better overview.
padding : float
Fraction of total number of bins to attach left and right
in the plot.
"""
plt.grid(True, zorder=-10)
plt.xlabel("time")
if self.do_mchits:
plt.ylabel("mchits / bin")
else:
plt.ylabel("hits / bin")
if isinstance(bins, int):
plt.hist(self.data, bins=bins, zorder=10, **kwargs)
else:
self.print_binstats(bins)
plt.hist(
self.data,
bins=_get_padded_bins(bins, padding),
zorder=10,
**kwargs
)
for bin_line_x in (bins[0], bins[-1]):
plt.axvline(
x=bin_line_x, color='firebrick', linestyle='--', zorder=20
)
def print_binstats(self, bin_edges):
print(f"Cutoff left: {np.mean(self.data < bin_edges[0]):.2%}")
print(f"Cutoff right: {np.mean(self.data > bin_edges[-1]):.2%}")
print(f"Avg. time per bin: {np.mean(np.diff(bin_edges)):.2f}")
def _read(self, file, det_file=None):
if self.det_file is not None:
cal = modules.DetApplier(det_file=det_file)
else:
cal = None
if self.center_time or self.add_t0:
time_pp = modules.TimePreproc(add_t0=self.add_t0, center_time=self.center_time)
else:
time_pp = None
pump = kp.io.hdf5.HDF5Pump(filename=file)
for blob in pump:
if cal is not None:
blob = cal.process(blob)
if time_pp is not None:
blob = time_pp.process(blob)
self.data = np.concatenate(
[self.data, self._get_data_one_event(blob)])
pump.close()
def _get_data_one_event(self, blob):
if self.do_mchits:
fld = "McHits"
else:
fld = "Hits"
blob_data = blob[fld]["time"]
if self.inactive_du is not None:
dus = blob[fld]["du"]
blob_data = blob_data[dus != self.inactive_du]
return blob_data
def _get_padded_bins(bin_edges, padding):
""" Add fraction of total # of bins, with same width. """
bin_width = bin_edges[1] - bin_edges[0]
n_bins = len(bin_edges) - 1
bins_to_add = np.ceil(padding * n_bins)
return np.linspace(
bin_edges[0] - bins_to_add * bin_width,
bin_edges[-1] + bins_to_add * bin_width,
n_bins + 2 * bins_to_add + 1,
)
\ No newline at end of file
#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
For investigating the ideal binning, based on the info in calibrated
.h5 files.
Specialized classes TimePlotter and ZPlotter are available for plotting
the time/ Z-Coordinate.
"""
import numpy as np
import km3pipe as kp
import matplotlib.pyplot as plt
from orcasong.modules import TimePreproc, DetApplier
__author__ = 'Stefan Reck'
class FieldPlotter:
"""
Baseclass for investigating the ideal binning, based on the info in
a field of calibrated .h5 files.
Intended for 1d binning, like for fields "time" or "pos_z".
Workflow:
1. Initialize with files, then run .plot() to extract and store
the data, and show the plot interactively.
2. Choose a binning via .set_binning.
3. Run .plot() again to show the plot with the adjusted binning on the
stored data.
4. Repeat step 2 and 3 unitl happy with binning.
(5.) Save plot via .plot(savepath), or get the bin edges via .get_bin_edges()
The plot will have some bins attached in both directions for
better overview.
Attributes:
-----------
files : list or str
The .h5 file(s).
field : str
The field to look stuff up, e.g. "time", "pos_z", ...
det_file : str, optional
Path to a detx file for calibrating.
filter_for_du : int, optional
Only get hits from one specific du, specified by the int.
hits : ndarray
The extracted Hits.
mc_hits : ndarray
The extracted McHits, if present.
n_events : int
The number of events in the extracted data.
limits : List
Left- and right-most edge of the binning.
n_bins : int
The number of bins.
plot_padding : List
Fraction of bins to append to left and right direction
(only in the plot for better overview).
x_label : str
X label of the plot.
y_label : str
Y label of the plot.
hist_kwargs : dict
Kwargs for plt.hist
xlim : List
The xlimits of the hist plot.
show_plots : bool
If True, auto plt.show() the plot.
"""
def __init__(self, files, field, det_file=None):
self.files = files
self.field = field
self.filter_for_du = None
self.det_file = det_file
self.hits = None
self.mc_hits = None
self.n_events = None
self.limits = None
self.n_bins = 100
self.plot_padding = [0.2, 0.2]
# Plotting stuff
self.xlabel = None
self.ylabel = 'Fraction of hits'
self.hist_kwargs = {
"histtype": "stepfilled",
"density": True,
}
self.xlim = None
self.ylim = None
self.show_plots = True
self.last_ylim = None
def plot(self, only_mc_hits=False, save_path=None):
"""
Generate and store or load the data, then make the plot.
Parameters
----------
only_mc_hits : bool
If true, plot the McHits instead of the Hits.
save_path : str, optional
Save plot to here.
Returns
-------
fig, ax : pyplot figure
The plot.
"""
if self.hits is None:
self.extract()
fig, ax = self.make_histogram(only_mc_hits, save_path)
return fig, ax
def set_binning(self, limits, n_bins):
"""
Set the desired binning.
Parameters
----------
limits : List
Left- and right-most edge of the binning.
n_bins : int
The number of bins.
"""
self.limits = limits
self.n_bins = n_bins
def get_binning(self):
"""
Set the desired binning.
Returns
-------
limits : List
Left- and right-most edge of the binning.
n_bins : int
The number of bins.
"""
return self.limits, self.n_bins
def get_bin_edges(self):
"""
Get the bin edges as a ndarray.
"""
limits, n_bins = self.get_binning()
if limits is None:
raise ValueError("Can not return bin edges: No binning limits set")
bin_edges = np.linspace(limits[0], limits[1], n_bins + 1)
return bin_edges
def extract(self):
"""
Extract the content of a field from all events in the file(s) and
store it.
"""
data_all_events = None
mc_all_events = None
self.n_events = 0
if not isinstance(self.files, list):
files = [self.files]
else:
files = self.files
event_pump = kp.io.hdf5.HDF5Pump(filenames=files)
if self.det_file is not None:
cal = DetApplier(det_file=self.det_file)
for i, event_blob in enumerate(event_pump):
self.n_events += 1
if self.det_file is not None:
event_blob = cal.process(event_blob)
if i % 2000 == 0:
print("Blob no. "+str(i))
data_one_event = self._get_hits(event_blob, get_mc_hits=False)
if data_all_events is None:
data_all_events = data_one_event
else:
data_all_events = np.concatenate(
[data_all_events, data_one_event], axis=0)
if "McHits" in event_blob:
mc_one_event = self._get_hits(event_blob, get_mc_hits=True)
if mc_all_events is None:
mc_all_events = mc_one_event
else:
mc_all_events = np.concatenate(
[mc_all_events, mc_one_event], axis=0)
event_pump.close()
print("Number of events: " + str(self.n_events))
self.hits = data_all_events
self.mc_hits = mc_all_events
def make_histogram(self, only_mc_hits=False, save_path=None):
"""
Plot the hist data. Can also save it if given a save path.
Parameters
----------
only_mc_hits : bool
If true, plot the McHits instead of the Hits.
save_path : str, optional
Save the fig to this path.
Returns
-------
fig, ax : pyplot figure
The plot.
"""
if only_mc_hits:
data = self.mc_hits
else:
data = self.hits
if data is None:
raise ValueError("Can not make histogram, no data extracted yet.")
bin_edges = self._get_padded_bin_edges()
fig, ax = plt.subplots()
bins = plt.hist(data, bins=bin_edges, **self.hist_kwargs)[1]
print("Size of first bin: " + str(bins[1] - bins[0]))
plt.grid(True, zorder=0, linestyle='dotted')
if self.limits is not None:
for bin_line_x in self.limits:
plt.axvline(x=bin_line_x, color='firebrick', linestyle='--')
if self.xlabel is None:
plt.xlabel(self._get_xlabel())
if self.xlim is not None:
plt.xlim(self.xlim)
if self.ylim is not None:
plt.ylim(self.ylim)
plt.ylabel(self.ylabel)
plt.tight_layout()
if save_path is not None:
print("Saving plot to "+str(save_path))
plt.savefig(save_path)
if self.show_plots:
plt.show()
return fig, ax
def _get_padded_bin_edges(self):
"""
Get the padded bin edges.
"""
limits, n_bins = self.get_binning()
if limits is None:
bin_edges = n_bins
else:
total_range = limits[1] - limits[0]
bin_size = total_range / n_bins
addtnl_bins = [
int(self.plot_padding[0] * n_bins),
int(self.plot_padding[1] * n_bins)
]
padded_range = [
limits[0] - bin_size * addtnl_bins[0],
limits[1] + bin_size * addtnl_bins[1]
]
padded_n_bins = n_bins + addtnl_bins[0] + addtnl_bins[1]
bin_edges = np.linspace(padded_range[0], padded_range[1],
padded_n_bins + 1)
return bin_edges
def _get_hits(self, blob, get_mc_hits):
"""
Get desired attribute from an event blob.
Parameters
----------
blob
The blob.
get_mc_hits : bool
If true, will get the "McHits" instead of the "Hits".
Returns
-------
blob_data : ndarray
The data.
"""
if get_mc_hits:
field_name = "McHits"
else:
field_name = "Hits"
blob_data = blob[field_name][self.field]
if self.filter_for_du is not None:
dus = blob[field_name]["du"]
blob_data = blob_data[dus == self.filter_for_du]
return blob_data
def _get_xlabel(self):
"""
Some saved xlabels.
"""
if self.field == "time":
xlabel = "Time [ns]"
elif self.field == "pos_z":
xlabel = "Z position [m]"
else:
xlabel = None
return xlabel
def __repr__(self):
return "<FieldPlotter: {}>".format(self.files)
class TimePlotter(FieldPlotter):
"""
For plotting the time.
"""
def __init__(self, files, add_t0=True, det_file=None):
field = "time"
FieldPlotter.__init__(self, files, field, det_file=det_file)
self.tp = TimePreproc(add_t0=add_t0)
def _get_hits(self, blob, get_mc_hits):
blob = self.tp.process(blob)
if get_mc_hits:
field_name = "McHits"
else:
field_name = "Hits"
blob_data = blob[field_name][self.field]
if self.filter_for_du is not None:
dus = blob[field_name]["du"]
blob_data = blob_data[dus == self.filter_for_du]
return blob_data
class ZPlotter(FieldPlotter):
"""
For plotting the z dim.
"""
def __init__(self, files, det_file=None):
field = "pos_z"
FieldPlotter.__init__(self, files, field, det_file=det_file)
self.plotting_bins = 100
def _get_padded_bin_edges(self):
"""
Get the padded bin edges.
"""
return self.plotting_bins
def set_binning(self, limits, n_bins):
"""
Set the desired binning.
Parameters
----------
limits : List
Left- and right-most edge of the binning.
n_bins : int
The number of bins.
"""
bin_edges = np.linspace(limits[0], limits[1],
n_bins + 1)
self.limits = bin_edges
self.n_bins = n_bins
def get_bin_edges(self):
return self.limits
......@@ -35,6 +35,9 @@ def plot_hists(hists, save_to, plot_bin_edges=True):
if isinstance(hists, list):
hists = combine_hists(hists)
if os.path.exists(save_to):
raise FileExistsError(f"File {save_to} exists already!")
with PdfPages(save_to) as pdf_file:
for bin_name, hists_data in hists.items():
hist_bin_edges = hists_data["hist_bin_edges"]
......
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