Coverage for src/km3pipe/calib.py: 69%
262 statements
« prev ^ index » next coverage.py v7.4.4, created at 2024-04-20 03:14 +0000
1# Filename: calib.py
2# pylint: disable=locally-disabled
3"""
4Calibration.
6"""
7import awkward as ak
8import numba as nb
9import numpy as np
11import km3db
12import km3io
14from thepipe import Module
15from km3pipe.hardware import Detector
16from km3pipe.dataclasses import Table
17from km3pipe.tools import istype
18from km3pipe.logger import get_logger
20__author__ = "Tamas Gal"
21__copyright__ = "Copyright 2016, Tamas Gal and the KM3NeT collaboration."
22__credits__ = ["Thomas Heid"]
23__license__ = "MIT"
24__maintainer__ = "Tamas Gal"
25__email__ = "tgal@km3net.de"
26__status__ = "Development"
28log = get_logger(__name__)
31class Calibration(Module):
32 """A module which applies time, position and rotation corrections to hits.
34 This module also calibrates MC hits, but be aware, t0s are not appended to
35 the MC hit times.
36 Additionally, the global PMT ID is added to regular hits as ``pmt_id`` and
37 in case of MC hits, the ``dom_id`` and ``channel_id`` (DAQ) are set.
39 Parameters
40 ----------
41 apply: bool, optional [default=True]
42 Apply the calibration to the hits (add position/direction/t0)?
43 filename: str, optional [default=None]
44 DetX file with detector description.
45 det_id: int, optional
46 .detx ID of detector (when retrieving from database).
47 t0set: optional
48 t0set (when retrieving from database).
49 calibset: optional
50 calibset (when retrieving from database).
51 key: str, optional [default="Hits"]
52 the blob key of the hits
53 outkey: str, optional [default="CalibHits"]
54 the output blob key of the calibrated hits
55 key_mc: str, optional [default="McHits"]
56 the blob key of the MC hits (if present)
57 outkey_mc: str, optional [default="CalibMcHits"]
58 the output blob key of the calibrated MC hits
59 """
61 __name__ = "Calibration"
62 name = "Calibration"
64 def configure(self):
65 self._should_apply = self.get("apply", default=True)
66 self.filename = self.get("filename")
67 self.det_id = self.get("det_id")
68 self.run = self.get("run")
69 self.t0set = self.get("t0set")
70 self.calibset = self.get("calibset")
71 self.detector = self.get("detector")
72 self.key = self.get("key", default="Hits")
73 self.outkey = self.get("outkey", default="CalibHits")
74 self.key_mc = self.get("key_mc", default="McHits")
75 self.outkey_mc = self.get("outkey_mc", default="CalibMcHits")
76 self._pos_dom_channel = None
77 self._dir_dom_channel = None
78 self._t0_dom_channel = None
79 self._pos_pmt_id = None
80 self._dir_pmt_id = None
81 self._t0_pmt_id = None
82 self._lookup_tables = None # for Numba
84 if self.det_id and self.run:
85 self.cprint(
86 "Grabbing the calibration for Det ID {} and run {}".format(
87 self.det_id, self.run
88 )
89 )
90 raw_detx = km3db.tools.detx_for_run(self.det_id, self.run)
91 self.detector = Detector(string=raw_detx)
92 self._create_dom_channel_lookup()
93 self._create_pmt_id_lookup()
94 return
96 if self.filename or self.det_id:
97 if self.filename is not None:
98 self.detector = Detector(filename=self.filename)
99 if self.det_id:
100 self.detector = Detector(
101 det_id=self.det_id, t0set=self.t0set, calibset=self.calibset
102 )
104 if self.detector is not None:
105 self.log.debug("Creating lookup tables")
106 self._create_dom_channel_lookup()
107 self._create_pmt_id_lookup()
108 else:
109 self.log.critical("No detector information loaded.")
111 def process(self, blob):
112 if self._should_apply:
113 blob[self.outkey] = self.apply(blob[self.key])
114 if self.key_mc in blob:
115 blob[self.outkey_mc] = self.apply(blob[self.key_mc])
116 return blob
118 def get_detector(self):
119 """Return the detector"""
120 return self.detector
122 def dus(self, hits):
123 """Return the DUs for given hits"""
124 dom_ids = hits["dom_id"]
125 return _dus(dom_ids, len(dom_ids), self._calib_by_dom_and_channel)
127 def floors(self, hits):
128 """Return the floors for given hits"""
129 dom_ids = hits["dom_id"]
130 return _floors(dom_ids, len(dom_ids), self._calib_by_dom_and_channel)
132 def apply_t0(self, hits):
133 """Apply only t0s"""
134 apply_t0_nb(hits.time, hits.dom_id, hits.channel_id, self._lookup_tables)
135 return hits
137 def apply(self, hits, no_copy=False, correct_slewing=True, slewing_variant=3):
138 """Add x, y, z, t0 (and du, floor if DataFrame) columns to the hits."""
139 if not no_copy:
140 try:
141 hits = hits.copy()
142 except AttributeError: # probably a km3io object
143 pass
145 if isinstance(hits, (ak.Array, ak.Record, km3io.rootio.Branch)):
146 if hasattr(hits, "dom_id"):
147 hits = Table(
148 dict(
149 dom_id=hits.dom_id,
150 channel_id=hits.channel_id,
151 time=hits.t,
152 tot=hits.tot,
153 triggered=hits.trig,
154 )
155 )
156 else: # mc_hits in km3io
157 hits = Table(
158 dict(
159 pmt_id=hits.pmt_id,
160 time=hits.t,
161 a=hits.a,
162 # TODO: Not all MC files have these two fields
163 # pure_a=hits.pure_a,
164 # pure_t=hits.pure_t,
165 origin=hits.origin,
166 )
167 )
169 if istype(hits, "DataFrame"):
170 # do we ever see McHits here?
171 hits = Table.from_template(hits, "Hits")
173 is_mc = None
174 if hasattr(hits, "dom_id") and hasattr(hits, "channel_id"):
175 try:
176 (
177 dir_x,
178 dir_y,
179 dir_z,
180 du,
181 floor,
182 pos_x,
183 pos_y,
184 pos_z,
185 t0,
186 pmt_id,
187 ) = _get_calibration_for_hits(hits, self._calib_by_dom_and_channel)
188 except KeyError as e:
189 self.log.critical("Wrong calibration (DETX) data provided.")
190 raise
191 is_mc = False
192 elif hasattr(hits, "pmt_id"):
193 try:
194 (
195 dir_x,
196 dir_y,
197 dir_z,
198 du,
199 floor,
200 pos_x,
201 pos_y,
202 pos_z,
203 t0,
204 dom_id,
205 channel_id,
206 ) = _get_calibration_for_mchits(hits, self._calib_by_pmt_id)
207 except KeyError as e:
208 self.log.critical("Wrong calibration (DETX) data provided.")
209 raise
210 is_mc = True
211 else:
212 raise TypeError(
213 "Don't know how to apply calibration to '{0}'. "
214 "We need at least 'dom_id' and 'channel_id', or "
215 "'pmt_id'.".format(hits.name)
216 )
218 if hasattr(hits, "time") and not is_mc:
219 if hits.time.dtype != t0.dtype:
220 time = hits.time.astype("f4") + t0.astype("f4")
221 hits = hits.drop_columns(["time"])
222 hits = hits.append_columns(["time"], [time])
223 else:
224 hits.time += t0
226 hits_data = {}
227 for colname in hits.dtype.names:
228 hits_data[colname] = hits[colname]
229 calib = {
230 "dir_x": dir_x,
231 "dir_y": dir_y,
232 "dir_z": dir_z,
233 "du": du.astype(np.uint8),
234 "floor": floor.astype(np.uint8),
235 "pos_x": pos_x,
236 "pos_y": pos_y,
237 "pos_z": pos_z,
238 "t0": t0,
239 }
241 if is_mc:
242 calib["dom_id"] = dom_id.astype(np.int32)
243 calib["channel_id"] = channel_id.astype(np.int32)
244 else:
245 calib["pmt_id"] = pmt_id.astype(np.int32)
247 hits_data.update(calib)
249 if correct_slewing and not is_mc:
250 hits_data["time"] -= slew(hits_data["tot"], variant=slewing_variant)
251 return Table(
252 hits_data, h5loc=hits.h5loc, split_h5=hits.split_h5, name=hits.name
253 )
255 def _create_dom_channel_lookup(self):
256 data = nb.typed.Dict.empty(
257 key_type=nb.types.i8, value_type=nb.types.float64[:, :]
258 )
259 for pmt in self.detector.pmts:
260 if pmt.dom_id not in data:
261 data[pmt.dom_id] = np.zeros((31, 10))
262 data[pmt.dom_id][pmt.channel_id] = np.asarray(
263 [
264 pmt.pos_x,
265 pmt.pos_y,
266 pmt.pos_z,
267 pmt.dir_x,
268 pmt.dir_y,
269 pmt.dir_z,
270 pmt.t0,
271 pmt.du,
272 pmt.floor,
273 pmt.pmt_id,
274 ],
275 dtype=np.float64,
276 )
277 self._calib_by_dom_and_channel = data
278 self._lookup_tables = [(dom, cal) for dom, cal in data.items()]
280 def _create_pmt_id_lookup(self):
281 data = nb.typed.Dict.empty(key_type=nb.types.i8, value_type=nb.types.float64[:])
282 for pmt in self.detector.pmts:
283 data[pmt.pmt_id] = np.asarray(
284 [
285 pmt.pos_x,
286 pmt.pos_y,
287 pmt.pos_z,
288 pmt.dir_x,
289 pmt.dir_y,
290 pmt.dir_z,
291 pmt.t0,
292 pmt.du,
293 pmt.floor,
294 pmt.dom_id,
295 pmt.channel_id,
296 ],
297 dtype=np.float64,
298 )
299 self._calib_by_pmt_id = data
301 def __repr__(self):
302 return self.__str__()
304 def __str__(self):
305 return "Calibration: det_id({0})".format(self.det_id)
308@nb.njit
309def apply_t0_nb(times, dom_ids, channel_ids, lookup_tables):
310 """Apply t0s using a lookup table of tuples (dom_id, calib)"""
311 dom_id = 0
312 lookup = np.empty((31, 9))
313 for i in range(len(times)):
314 cur_dom_id = dom_ids[i]
315 if cur_dom_id != dom_id:
316 dom_id = cur_dom_id
317 for (d, m) in lookup_tables:
318 if d == dom_id:
319 np.copyto(lookup, m)
320 t0 = lookup[channel_ids[i]][6]
321 times[i] += t0
324@nb.jit(nopython=True)
325def _get_calibration_for_hits(hits, lookup):
326 """Append the position, direction and t0 columns and add t0 to time"""
327 n = len(hits)
328 cal = np.empty((n, 10))
329 for i in range(n):
330 calib = lookup[hits["dom_id"][i]][hits["channel_id"][i]]
331 cal[i] = calib
332 dir_x = cal[:, 3]
333 dir_y = cal[:, 4]
334 dir_z = cal[:, 5]
335 du = cal[:, 7]
336 floor = cal[:, 8]
337 pos_x = cal[:, 0]
338 pos_y = cal[:, 1]
339 pos_z = cal[:, 2]
340 pmt_id = cal[:, 9]
342 t0 = cal[:, 6]
344 return [dir_x, dir_y, dir_z, du, floor, pos_x, pos_y, pos_z, t0, pmt_id]
347@nb.jit(nopython=True)
348def _get_calibration_for_mchits(hits, lookup):
349 """Append the position, direction and t0 columns and add t0 to time"""
350 n_hits = len(hits)
351 cal = np.empty((n_hits, 11))
352 for i in range(n_hits):
353 cal[i] = lookup[hits["pmt_id"][i]]
354 dir_x = cal[:, 3]
355 dir_y = cal[:, 4]
356 dir_z = cal[:, 5]
357 du = cal[:, 7]
358 floor = cal[:, 8]
359 pos_x = cal[:, 0]
360 pos_y = cal[:, 1]
361 pos_z = cal[:, 2]
362 t0 = cal[:, 6]
363 dom_id = cal[:, 9]
364 channel_id = cal[:, 10]
366 return [dir_x, dir_y, dir_z, du, floor, pos_x, pos_y, pos_z, t0, dom_id, channel_id]
369@nb.njit
370def _dus(dom_ids, n, lookup):
371 dus = np.empty(n, dtype=np.uint8)
372 for i in range(n):
373 # DU is at index 7
374 dus[i] = lookup[dom_ids[i]][0][7] # looking only at channel ID 0
375 return dus
378@nb.njit
379def _floors(dom_ids, n, lookup):
380 dus = np.empty(n, dtype=np.uint8)
381 for i in range(n):
382 # floor is at index 8
383 dus[i] = lookup[dom_ids[i]][0][8] # looking only at channel ID 0
384 return dus
387class CalibrationService(Module):
388 """A service which provides calibration routines for hits
390 Parameters
391 ----------
392 filename: str, optional [default=None]
393 DetX file with detector description.
394 det_id: int, optional
395 .detx ID of detector (when retrieving from database).
396 t0set: optional
397 t0set (when retrieving from database).
398 calibset: optional
399 calibset (when retrieving from database).
400 detector: kp.hardware.Detector, optional
401 """
403 __name__ = "Calibration"
404 name = "Calibration"
406 def configure(self):
407 self.filename = self.get("filename")
408 self.det_id = self.get("det_id")
409 self.t0set = self.get("t0set")
410 self.calibset = self.get("calibset")
412 self._detector = self.get("detector")
414 if self._detector is not None:
415 self._calibration = Calibration(detector=self._detector)
417 self._calibration = None
419 self.expose(self.calibrate, "calibrate")
420 self.expose(self.get_detector, "get_detector")
421 self.expose(self.get_calibration, "get_calibration")
422 self.expose(self.load_calibration, "load_calibration")
423 self.expose(self.correct_slewing, "correct_slewing")
425 def load_calibration(self, filename=None, det_id=None, t0set=None, calibset=None):
426 """Load another calibration"""
427 self.filename = filename
428 self.det_id = det_id
429 self.t0set = t0set
430 self.calibset = calibset
431 self._detector = None
432 self._calibration = None
434 def calibrate(self, hits, correct_slewing=True):
435 return self.calibration.apply(hits, correct_slewing=correct_slewing)
437 @property
438 def detector(self):
439 if self._detector is None:
440 self._detector = self.calibration.detector
441 return self._detector
443 def get_detector(self):
444 """Extra getter to be as lazy as possible (expose triggers otherwise"""
445 return self.detector
447 @property
448 def calibration(self):
449 if self._calibration is None:
450 self._calibration = Calibration(
451 filename=self.filename,
452 det_id=self.det_id,
453 t0set=self.t0set,
454 calibset=self.calibset,
455 )
456 return self._calibration
458 def get_calibration(self):
459 """Extra getter to be as lazy as possible (expose triggers otherwise"""
460 return self.calibration
462 def correct_slewing(self, hits):
463 """Apply time slewing correction to the hit times"""
464 hits.time -= slew(hits.tot)
467def slew(tot, variant=3):
468 """Calculate the time slewing of a PMT response for a given ToT
471 Parameters
472 ----------
473 tot: int or np.array(int)
474 Time over threshold value of a hit
475 variant: int, optional
476 The variant of the slew calculation.
477 1: The first parametrisation approach
478 2: Jannik's improvement of the parametrisation
479 3: The latest lookup table approach based on lab measurements.
481 Returns
482 -------
483 time: int
484 Time slewing, which has to be subtracted from the original hit time.
485 """
486 if variant == 1:
487 return _slew_parametrised(7.70824, 0.00879447, -0.0621101, -1.90226, tot)
488 if variant == 2:
489 return _slew_parametrised(
490 13.6488662517, -0.128744123166, -0.0174837749244, -4.47119633965, tot
491 )
492 if variant == 3:
493 if isinstance(tot, (int, np.integer)):
494 return _SLEWS[tot]
495 return _slew_tabulated(np.array(_SLEWS), tot)
497 raise ValueError("Unknown slew calculation variant '{}'".format(variant))
500@nb.jit(nopython=True)
501def _slew_parametrised(p0, p1, p2, p3, tot):
502 return p0 * np.exp(p1 * np.sqrt(tot) + p2 * tot) + p3
505@nb.jit(nopython=True)
506def _slew_tabulated(tab, tots):
507 n = len(tots)
508 out = np.empty(n)
509 for i in range(n):
510 out[i] = tab[tots[i]]
511 return out
514_SLEWS = np.array(
515 [
516 8.01,
517 7.52,
518 7.05,
519 6.59,
520 6.15,
521 5.74,
522 5.33,
523 4.95,
524 4.58,
525 4.22,
526 3.89,
527 3.56,
528 3.25,
529 2.95,
530 2.66,
531 2.39,
532 2.12,
533 1.87,
534 1.63,
535 1.40,
536 1.19,
537 0.98,
538 0.78,
539 0.60,
540 0.41,
541 0.24,
542 0.07,
543 -0.10,
544 -0.27,
545 -0.43,
546 -0.59,
547 -0.75,
548 -0.91,
549 -1.08,
550 -1.24,
551 -1.41,
552 -1.56,
553 -1.71,
554 -1.85,
555 -1.98,
556 -2.11,
557 -2.23,
558 -2.35,
559 -2.47,
560 -2.58,
561 -2.69,
562 -2.79,
563 -2.89,
564 -2.99,
565 -3.09,
566 -3.19,
567 -3.28,
568 -3.37,
569 -3.46,
570 -3.55,
571 -3.64,
572 -3.72,
573 -3.80,
574 -3.88,
575 -3.96,
576 -4.04,
577 -4.12,
578 -4.20,
579 -4.27,
580 -4.35,
581 -4.42,
582 -4.49,
583 -4.56,
584 -4.63,
585 -4.70,
586 -4.77,
587 -4.84,
588 -4.90,
589 -4.97,
590 -5.03,
591 -5.10,
592 -5.16,
593 -5.22,
594 -5.28,
595 -5.34,
596 -5.40,
597 -5.46,
598 -5.52,
599 -5.58,
600 -5.63,
601 -5.69,
602 -5.74,
603 -5.80,
604 -5.85,
605 -5.91,
606 -5.96,
607 -6.01,
608 -6.06,
609 -6.11,
610 -6.16,
611 -6.21,
612 -6.26,
613 -6.31,
614 -6.36,
615 -6.41,
616 -6.45,
617 -6.50,
618 -6.55,
619 -6.59,
620 -6.64,
621 -6.68,
622 -6.72,
623 -6.77,
624 -6.81,
625 -6.85,
626 -6.89,
627 -6.93,
628 -6.98,
629 -7.02,
630 -7.06,
631 -7.09,
632 -7.13,
633 -7.17,
634 -7.21,
635 -7.25,
636 -7.28,
637 -7.32,
638 -7.36,
639 -7.39,
640 -7.43,
641 -7.46,
642 -7.50,
643 -7.53,
644 -7.57,
645 -7.60,
646 -7.63,
647 -7.66,
648 -7.70,
649 -7.73,
650 -7.76,
651 -7.79,
652 -7.82,
653 -7.85,
654 -7.88,
655 -7.91,
656 -7.94,
657 -7.97,
658 -7.99,
659 -8.02,
660 -8.05,
661 -8.07,
662 -8.10,
663 -8.13,
664 -8.15,
665 -8.18,
666 -8.20,
667 -8.23,
668 -8.25,
669 -8.28,
670 -8.30,
671 -8.32,
672 -8.34,
673 -8.37,
674 -8.39,
675 -8.41,
676 -8.43,
677 -8.45,
678 -8.47,
679 -8.49,
680 -8.51,
681 -8.53,
682 -8.55,
683 -8.57,
684 -8.59,
685 -8.61,
686 -8.62,
687 -8.64,
688 -8.66,
689 -8.67,
690 -8.69,
691 -8.70,
692 -8.72,
693 -8.74,
694 -8.75,
695 -8.76,
696 -8.78,
697 -8.79,
698 -8.81,
699 -8.82,
700 -8.83,
701 -8.84,
702 -8.86,
703 -8.87,
704 -8.88,
705 -8.89,
706 -8.90,
707 -8.92,
708 -8.93,
709 -8.94,
710 -8.95,
711 -8.96,
712 -8.97,
713 -8.98,
714 -9.00,
715 -9.01,
716 -9.02,
717 -9.04,
718 -9.04,
719 -9.04,
720 -9.04,
721 -9.04,
722 -9.04,
723 -9.04,
724 -9.04,
725 -9.04,
726 -9.04,
727 -9.04,
728 -9.04,
729 -9.04,
730 -9.04,
731 -9.04,
732 -9.04,
733 -9.04,
734 -9.04,
735 -9.04,
736 -9.04,
737 -9.04,
738 -9.04,
739 -9.04,
740 -9.04,
741 -9.04,
742 -9.04,
743 -9.04,
744 -9.04,
745 -9.04,
746 -9.04,
747 -9.04,
748 -9.04,
749 -9.04,
750 -9.04,
751 -9.04,
752 -9.04,
753 -9.04,
754 -9.04,
755 -9.04,
756 -9.04,
757 -9.04,
758 -9.04,
759 -9.04,
760 -9.04,
761 -9.04,
762 -9.04,
763 -9.04,
764 -9.04,
765 -9.04,
766 -9.04,
767 -9.04,
768 -9.04,
769 -9.04,
770 -9.04,
771 -9.04,
772 ]
773)