Unable to access an otherwise available field of TLorentzVector in expressions

[alpetukhov]

Dear uproot developers!

I'm trying to read a .root file that contains branches with TLorentzVectors in the following form

import uproot

f = uproot.open('file:tree:)
p = f.arrays(["truth_p4"])
print(p.truth_p4[0])

{fP: {fX: 2.12e+04, fY: -3.81e+03, fZ: 1.38e+04}, fE: 2.56e+04}

After I've requested this branch, I have access to any of its fields, e.g. this

p.truth_p4.fE

yields

<Array [2.56e+04, 9.36e+04, ... 2.17e+04] type='54465 * float64'>

I've seen in discussions here that there's possibility to access the fields directly in the expressions fields of arrays or iterate. I want it to specify the aliases upfront to simplify the request for branches like transverse momentum of pseudorapidity.

But if I try to run

f.arrays(["truth_p4.fE"])

I get the following errors

Unexpected exception formatting exception. Falling back to standard exception
Traceback (most recent call last):
  File "/home/alex/trt_pid_venv/lib/python3.8/site-packages/uproot/language/python.py", line 382, in free_symbols
    )
  File "/home/alex/trt_pid_venv/lib/python3.8/site-packages/uproot/language/python.py", line 92, in _walk_ast_yield_symbols
  File "/home/alex/trt_pid_venv/lib/python3.8/site-packages/uproot/language/python.py", line 97, in _walk_ast_yield_symbols
    else:
  File "/home/alex/trt_pid_venv/lib/python3.8/site-packages/uproot/language/python.py", line 92, in _walk_ast_yield_symbols
  File "/home/alex/trt_pid_venv/lib/python3.8/site-packages/uproot/language/python.py", line 87, in _walk_ast_yield_symbols
KeyError: 'truth_p4.fE'

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/home/alex/trt_pid_venv/lib/python3.8/site-packages/IPython/core/interactiveshell.py", line 3508, in run_code
    exec(code_obj, self.user_global_ns, self.user_ns)
  File "/tmp/ipykernel_1902/203329954.py", line 1, in <module>
    f.arrays(["truth_p4.fE"])
  File "/home/alex/trt_pid_venv/lib/python3.8/site-packages/uproot/behaviors/TBranch.py", line 1095, in arrays
  File "/home/alex/trt_pid_venv/lib/python3.8/site-packages/uproot/behaviors/TBranch.py", line 3285, in _regularize_expressions
  File "/home/alex/trt_pid_venv/lib/python3.8/site-packages/uproot/behaviors/TBranch.py", line 3151, in _regularize_expression
  File "/home/alex/trt_pid_venv/lib/python3.8/site-packages/uproot/language/python.py", line 388, in free_symbols
uproot.exceptions.KeyInFileError: not found: 'truth_p4.fE'
in file /file.root
in object /tree;1

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/home/alex/trt_pid_venv/lib/python3.8/site-packages/IPython/core/interactiveshell.py", line 2105, in showtraceback
    stb = self.InteractiveTB.structured_traceback(
  File "/home/alex/trt_pid_venv/lib/python3.8/site-packages/IPython/core/ultratb.py", line 1396, in structured_traceback
    return FormattedTB.structured_traceback(
  File "/home/alex/trt_pid_venv/lib/python3.8/site-packages/IPython/core/ultratb.py", line 1287, in structured_traceback
    return VerboseTB.structured_traceback(
  File "/home/alex/trt_pid_venv/lib/python3.8/site-packages/IPython/core/ultratb.py", line 1140, in structured_traceback
    formatted_exception = self.format_exception_as_a_whole(etype, evalue, etb, number_of_lines_of_context,
  File "/home/alex/trt_pid_venv/lib/python3.8/site-packages/IPython/core/ultratb.py", line 1055, in format_exception_as_a_whole
    frames.append(self.format_record(record))
  File "/home/alex/trt_pid_venv/lib/python3.8/site-packages/IPython/core/ultratb.py", line 955, in format_record
    frame_info.lines, Colors, self.has_colors, lvals
  File "/home/alex/trt_pid_venv/lib/python3.8/site-packages/IPython/core/ultratb.py", line 778, in lines
    return self._sd.lines
  File "/home/alex/trt_pid_venv/lib/python3.8/site-packages/stack_data/utils.py", line 144, in cached_property_wrapper
    value = obj.__dict__[self.func.__name__] = self.func(obj)
  File "/home/alex/trt_pid_venv/lib/python3.8/site-packages/stack_data/core.py", line 734, in lines
    pieces = self.included_pieces
  File "/home/alex/trt_pid_venv/lib/python3.8/site-packages/stack_data/utils.py", line 144, in cached_property_wrapper
    value = obj.__dict__[self.func.__name__] = self.func(obj)
  File "/home/alex/trt_pid_venv/lib/python3.8/site-packages/stack_data/core.py", line 681, in included_pieces
    pos = scope_pieces.index(self.executing_piece)
  File "/home/alex/trt_pid_venv/lib/python3.8/site-packages/stack_data/utils.py", line 144, in cached_property_wrapper
    value = obj.__dict__[self.func.__name__] = self.func(obj)
  File "/home/alex/trt_pid_venv/lib/python3.8/site-packages/stack_data/core.py", line 660, in executing_piece
    return only(
  File "/home/alex/trt_pid_venv/lib/python3.8/site-packages/executing/executing.py", line 190, in only
    raise NotOneValueFound('Expected one value, found 0')
executing.executing.NotOneValueFound: Expected one value, found 0

I've got it with uproot 4.1.0 and awkward 1.4.0, updated to uproot 5.0.11 and awkward 2.4.2 but got the same results.

Is there a way to access the fields of the TLorentzVector in the expressions or aliases?

Thanks in advance,
Aleksandr

[alpetukhov]

I've got another question about the expressions that I think I'll also post here to not create too many new topic.
I want to get the sums of all the variables in the array (simple array of ints). I can get it with

import uproot
import awkward as ak

f = uproot.open(file_path)
a = f.arrays(["InDetTrackParticles_hit_gasType"])
ak.sum(a['InDetTrackParticles_hit_gasType'], axis=1)

but if I try to run it as

f.arrays(['sum(InDetTrackParticles_hit_gasType, axis=1)'])

I get the error

KeyInFileError: not found: 'sum'
in file file.root
in object tree;1

Shouldn't the sum translate to ak.sum here? I've got the ak.where to work in expression, so it is strange to me that sum is not working properely.
Sorry if those are basic questions, but I wasn't able to find the info in the docs or to Google it.

[jpivarski]

In the case of #131, the things with names like "Tracks.fCoordinates.fX" are TBranches. Using a dotted name in expressions makes it look for a TBranch with that name. A TLorentzVector is a single TBranch; the structure is inside of it—there is no TBranch in your file named "truth_p4.fE".¹

Outside of the expressions, if you have an Awkward Array named truth_p4, then asking for .fE extracts a field from it as a projection. That's something where a different library (Awkward vs Uproot) uses the same syntax for a different purpose.

The set of functions available in the expressions argument of uproot.TTree.arrays and similar is a closed-off set, like NumExpr. Here's the default set:

uproot5/src/uproot/language/python.py

Lines 258 to 319 in ac23179

	default_functions = {
	"abs": numpy.absolute,
	"absolute": numpy.absolute,
	"acos": numpy.arccos,
	"arccos": numpy.arccos,
	"acosh": numpy.arccosh,
	"arccosh": numpy.arccosh,
	"asin": numpy.arcsin,
	"arcsin": numpy.arcsin,
	"asinh": numpy.arcsinh,
	"arcsinh": numpy.arcsinh,
	"atan": numpy.arctan,
	"atan2": numpy.arctan2,
	"arctan": numpy.arctan,
	"arctan2": numpy.arctan2,
	"atanh": numpy.arctanh,
	"arctanh": numpy.arctanh,
	"cbrt": numpy.cbrt,
	"ceil": numpy.ceil,
	"conj": numpy.conjugate,
	"conjugate": numpy.conjugate,
	"copysign": numpy.copysign,
	"cos": numpy.cos,
	"cosh": numpy.cosh,
	"erf": _vectorized_erf(False),
	"erfc": _vectorized_erf(True),
	"exp": numpy.exp,
	"exp2": numpy.exp2,
	"expm1": numpy.expm1,
	"fabs": numpy.fabs,
	"factorial": lambda x: numpy.round(numpy.exp(_lgamma(numpy.round(x) + 1))),
	"floor": numpy.floor,
	"fmax": numpy.fmax,
	"fmin": numpy.fmin,
	"gamma": _vectorized_gamma(False),
	"hypot": numpy.hypot,
	"imag": numpy.imag,
	"isfinite": numpy.isfinite,
	"isinf": numpy.isinf,
	"isnan": numpy.isnan,
	"lgamma": _lgamma,
	"log": numpy.log,
	"log10": numpy.log10,
	"log1p": numpy.log1p,
	"log2": numpy.log2,
	"logical_and": numpy.logical_and,
	"logical_or": numpy.logical_or,
	"neg": numpy.negative,
	"nextafter": numpy.nextafter,
	"real": numpy.real,
	"rint": numpy.rint,
	"round": numpy.round,
	"signbit": numpy.signbit,
	"sin": numpy.sin,
	"sinh": numpy.sinh,
	"sqrt": numpy.sqrt,
	"tan": numpy.tan,
	"tanh": numpy.tanh,
	"tgamma": lambda x: numpy.exp(_lgamma(x)),
	"trunc": numpy.trunc,
	"where": numpy.where,
	}

We could add numpy.sum to that—not a bad idea! It would involve adding "sum": numpy.sum to the list and adding a test to the tests directory, using one of the existing tests as an example, like test_0930-....

It's also possible to use more functions than the default set without changing Uproot. Functions like uproot.TTree.arrays have a language argument, which take an uproot.language.Language object. There's only one concrete subclass, uproot.language.python.PythonLanguage, which can be constructed with a new set of functions:

PythonLanguage(PythonLanguage.default_functions | {"sum": np.sum})

(Note: d1 | d2 is Python 3.9+ syntax; for older Pythons, use dict(PythonLanguage.default_functions, sum=np.sum).)

It's interesting that you bring this up now because @aryan26roy is currently working on a second uproot.language.python.PythonLanguage with the traditional TTree::Draw syntax.

Footnotes

1. Maybe it would be possible to change Uproot so that it can do this. Currently, it treats all dotted expressions as the names of TBranches, but if the _attribute_to_dotted_name function were given a list of all TBranch names, it could stop descending into attributes if there are no TBranches with that name. For instance, if there are TBranches named "a.b" and "a.b.c" but no "a.b.c.d" and the expressions encounters a.b.c.d.e, then perhaps it should assume that you want to take the TBranch named "a.b.c" and let Awkward Array project field d and subfield e from it? Or would that rule make it too hard to predict what it's attempting to do? When we provide an interface, we need to make sure that it's easily understood by its users. ↩