From af27859b20633b4f865143bd5038cf17152404db Mon Sep 17 00:00:00 2001
From: zineb aly <aly.zineb.az@gmail.com>
Date: Mon, 20 Apr 2020 18:23:17 +0200
Subject: [PATCH] add fitparams function
---
km3io/offline.py | 30 +++++++++++++++++++++++++++++-
1 file changed, 29 insertions(+), 1 deletion(-)
diff --git a/km3io/offline.py b/km3io/offline.py
index fb2c05a..83c392e 100644
--- a/km3io/offline.py
+++ b/km3io/offline.py
@@ -1,7 +1,8 @@
from collections import namedtuple
import uproot
import warnings
-from .definitions import mc_header
+import awkward1 as ak1
+from .definitions import mc_header, fitparameters
from .tools import Branch, BranchMapper, cached_property, _to_num, _unfold_indices
MAIN_TREE_NAME = "E"
@@ -17,6 +18,33 @@ def _nested_mapper(key):
return '_'.join(key.split('.')[1:])
+def fitinf(fitparam, tracks):
+ """Access fit parameters in tracks.fitinf.
+
+ Parameters
+ ----------
+ fitparam : str
+ the fit parameter name according to fitparameters defined in
+ km3net-dataformat.
+ tracks : class km3io.offline.OfflineBranch
+ the tracks class .
+
+ Returns
+ -------
+ awkward array
+ awkward array of the fit parameters.
+ """
+ fit = tracks.fitinf
+ counts = ak1.Array([ak1.num(i, axis=1) for i in fit])
+ index = fitparameters[fitparam]
+ params = fit[counts > index]
+ return ak1.Array([i[:, index] for i in params])
+
+
+def fitparams():
+ return fitparameters.keys()
+
+
EVENTS_MAP = BranchMapper(name="events",
key="Evt",
extra={
--
GitLab