search

scikit-hep / uproot3-methods

Pythonic behaviors for non-I/O related ROOT classes.
BSD 3-Clause "New" or "Revised" License
21 stars 28 forks source link

rotate method for TVector2 and TVector2Array #61

Closed phylsix closed 5 years ago

phylsix commented 5 years ago

I see rotate method for TVector2 and TVector2Array missing and marked as TODO: https://github.com/scikit-hep/uproot-methods/blob/master/uproot_methods/classes/TVector2.py#L20-L21

Is it possible to add that? I see you mentioned in #27 . Thanks!

jpivarski commented 5 years ago

Hi! Sorry I'm getting back to this late—I'm catching up to issues from the weekend.

The idea of uproot-methods was to be user-contributed—not because I'm trying to offload work, but because the scope that it represents (Pythonic method for any ROOT classes that users need) is very open-ended.

You've seen the code: each of these methods assumes Numpy-like arrays (which can be awkward arrays, thanks to their shared behavior) and performs vectorized calculations on them. The corresponding rotation code for TVector3 (https://github.com/scikit-hep/uproot-methods/blob/master/uproot_methods/classes/TVector3.py#L32-L80) is much more complicated than it would be for TVector2, since that just takes a scalar angle.

It would be something like this:

    def rotate(self, angle):
        cosangle = self.awkward.numpy.cos(angle)
        sinangle = self.awkward.numpy.sin(angle)
        xprime = self.x*cosangle - self.y*sinangle
        yprime = self.x*sinangle + self.y*cosangle
        return TVector2Array.from_cartesian(xprime, yprime)

Okay, I guess there's a slight complication from the fact that we want to get numpy from an object attached to the array itself (so that in the future, this can be swapped for cupy for GPU processing).

Also, looking into it, I wonder how deep your use-case goes—if you're only interested in Numpy arrays of 2-vectors or any level of jaggedness. The TLorentzVector.from_cartesian constructor has been expanded to handle arbitrary jaggedness, but not the TVector2.from_cartesian constructor. (Oh, TVector3 uses a different trick to get both flat-Numpy and jagged arrays; it uses empty_like(): https://github.com/scikit-hep/uproot-methods/blob/master/uproot_methods/classes/TVector3.py#L119-L125 . Maybe we could do the same here.)

See what I mean about the open-ended scope? If you try to implement this for your case (and open a pull request), I'll help you with any issues you run into. Thanks!

phylsix commented 5 years ago

Hi, I open a PR #64, mimicking what TVector3 is implemented. What do you think? Thanks!