feat: validate coordinate names in one place, for every backend

[ikrommyd]

🤖 AI text below 🤖

Description

Rebased on main and reimplemented. Solves the half of #660 that vector can fix on its own.

The problem

Each backend had its own rules for which coordinate names make a vector — obj(), the VectorObject*D constructors, array(), the six __array_finalize__, _check_names, the VectorSympy*D constructors, and the Numba vector.obj — and they disagreed. All of these work on main and silently build the wrong vector:

vector.array({"pt": ..., "phi": ..., "eta": ..., "mass": ..., "energy": ...})  # energy wins, mass is dead
vector.array({"x": ..., "y": ..., "rho": ..., "phi": ...})                     # rho, phi are dead
vector.array({"x": ..., "y": ..., "z": ..., "eta": ...})                       # eta is dead
vector.zip({"rho": ..., "phi": ..., "pt": ...})                                # pt is dead
vector.VectorObject2D(x=1, px=2, y=3)                                          # x == 2

The fix

_methods._check_coordinate_names is now the only implementation. It returns (is_momentum, dimension, coordinates, extra) or raises, and takes the dimension and momentum-ness a constructor is restricted to, so it serves both vector.obj (deduces both) and MomentumNumpy4D (fixes both). Every backend calls it, and the error messages come from one table instead of the nine hand-written copies. Source diff: −1016/+567.

The awkward behaviors validate through __awkward_validation__ (awkward#3710, needs awkward>=2.8.11; older versions skip it), so what never went through a vector constructor is caught too:

>>> ak.zip({"pt": pt, "eta": eta, "phi": phi, "mass": mass, "energy": energy}, with_name="Momentum4D")
TypeError: MomentumArray4D with fields [...]: specify t= or tau=, but not more than one

>>> jets["rho"] = rho
TypeError: MomentumArray4D with fields [...]: 'pt' and 'rho' both map to 'rho'

Subclassed behaviors (coffea's) inherit it. Cost is ~0.5 µs per array creation against ~130 µs for one v.pt on 100k entries; _check_coordinate_names is lru_cached on the field-name tuple.

Intended breaks

1. A field that is a coordinate under another name is no longer carried as payload. Ordinary extra fields (charge, btag) are unaffected.
2. Generic vectors reject momentum-aliases: VectorObject2D(px=1, py=2) used to return a non-momentum vector while vector.obj(px=1, py=2) returns MomentumObject2D. Momentum classes still take generic names.
3. An array whose fields stop describing its record name is rejected — del jets["mass"], or a column-projected form that kept __record__ but dropped coordinates. On main those gave a half-broken vector that failed at first coordinate access.
4. Error messages changed. No test asserted on their text.

Tests

tests/test_pr_659.py: all 511 subsets of the generic names and all 222 combinations with momentum-aliases, per backend, against a reference implementation written independently of the one under test; plus the record names accepted and rejected by __awkward_validation__, and the numba path. 3969 tests in ~5 s, 100% coverage of the new code.

Coffea's validation and this agree on all 48,597 combinations I compared, and coffea's suite passes against this branch unchanged (its own hooks shadow vector's, so nothing there needs to change).

Checklist

• Have you followed the guidelines in our Contributing document?
• Have you checked to ensure there aren't any other open Pull Requests for the required change?
• Does your submission pass pre-commit? ($ pre-commit run --all-files or $ nox -s lint)
• Does your submission pass tests? ($ pytest or $ nox -s tests)
• Does the documentation build with your changes? ($ cd docs; make clean; make html or $ nox -s docs)
• Does your submission pass the doctests? ($ pytest --doctest-plus src/vector/ or $ nox -s doctests)

[Saransh-cpp]

Good catch, @ikrommyd! Thanks for working on this!!

Please see my comments below:

[Saransh-cpp]

Should this also have the and "t" not in generic_coordinates condition?

[ikrommyd]

Isn't this unnecessary? This is the first if condition where generic_coordinates might be populated with t. So and "t" not in generic_coordinates is a useless check no? Same goes for tau in your other comment below.

[ikrommyd]

I guess we can add it for "visual OCD reasons" but it's a check that's always going to be false right?

[Saransh-cpp]

Ditto, for tau

[Saransh-cpp]

vector.array is just a wrapper around individual constructors of Vector/MomentumNumpy*D, which can be used to construct vectors (unlike the Awkward backend). Hence, it would be better if we move this check to the __array_finalize__ method of each class:

vector/src/vector/backends/numpy.py

Line 1168 in e374915

def __array_finalize__(self, obj: typing.Any) -> None:

[ikrommyd]

I vaguely remember that I had some issue with __array_finalize__ regarding when is it being ran when I was making these edits but I will take a look again.

[ikrommyd]

I've moved calling this function inside the __array_finalize__ of the 2d, 3d, and 4d classes. The tests pass so I'll need to play with it a bit more in a terminal in case there's something I'm not remembering here.

[ikrommyd]

Oh..now I think I remember. Look at this

vector/src/vector/backends/numpy.py

Lines 2160 to 2169 in 776f405

	is_momentum = any(x in _repr_momentum_to_generic for x in names)
	if any(x in ("t", "E", "e", "energy", "tau", "M", "m", "mass") for x in names):
	cls = MomentumNumpy4D if is_momentum else VectorNumpy4D
	elif any(x in ("z", "pz", "theta", "eta") for x in names):
	cls = MomentumNumpy3D if is_momentum else VectorNumpy3D
	else:
	cls = MomentumNumpy2D if is_momentum else VectorNumpy2D
	return cls(*args, **kwargs)

It makes a decision whether it is momentum or not and which class constructor to use based on the field names. That's why I think I wanted it in the vector.array implementation and not in the finalize. I think it's best to validate the fields first and then decide which constructor to use. I think it's best to move it back to where it was. I have reverted the change.

[ikrommyd]

Actually, it needs to be in both places because classes like MomentumNumpy4D can be directly instantiated. And therefore this works fine and it shouldn't.

In [1]: import vector
   ...:
   ...: vec = vector.MomentumNumpy4D([(1.1, 2.1, 3.1, 4.1, 4.1), (1.2, 2.2, 3.2, 4.2, 4.2), (1.3, 2.3, 3.3, 4.3, 4.3), (1.4, 2.4, 3.4, 4.4, 4.4), (1.5, 2.5, 3.5, 4.5, 4.5)],
   ...:
   ...:               dtype=[('x', float), ('y', float), ('z', float), ('e', float), ('mass', float)])
   ...:
   ...: vec
Out[1]:
MomentumNumpy4D([(1.1, 2.1, 3.1, 4.1, 4.1), (1.2, 2.2, 3.2, 4.2, 4.2),
                 (1.3, 2.3, 3.3, 4.3, 4.3), (1.4, 2.4, 3.4, 4.4, 4.4),
                 (1.5, 2.5, 3.5, 4.5, 4.5)],
                dtype=[('x', '<f8'), ('y', '<f8'), ('z', '<f8'), ('t', '<f8'), ('tau', '<f8')])

So I think we need to validate in vector.array and in __array_finalize__ too because both are valid instantiation methods.

[Saransh-cpp]

Given that these tests are related to constructing vectors, we should move them to the appropriate test backend files. The object ones can go in:

vector/tests/backends/test_object.py

Line 14 in e374915

def test_constructors_2D():

The NumPy ones (we don't have much constructor tests for NumPy at the moment, so this is perfect 😅):

vector/tests/backends/test_numpy.py

Line 85 in e374915

def test_xy():

The Awkward ones (same as NumPy):

vector/tests/backends/test_awkward.py

Line 161 in e374915

def test_basic(backend):

[ikrommyd]

Sure. I just thought that maybe it should be a new test because it's like 1K lines but yeah I can move them under their respective backend tests.

[Saransh-cpp]

Ah, I see, it might make sense to create a new function under test_issues.py for this then 🤔

[ikrommyd]

Any final suggestion here? To me, putting 1K lines under test_issues.py is not a great idea. My understanding is that this file is for small individual issues. The only refactoring that I think may be good here is just splitting this 1K file into multiple files per backend. Any takes?

[Saransh-cpp]

Having a different file sounds good, @ikrommyd. But, can we rename it to test_issue_xxx.py with either the issue/PR number such that it is easy to spot and stays somewhat aligned with the issue_xxx tests?

[ikrommyd]

I've renamed to tests/test_pr_659.py because there is no exact issue, only the long discussion. It's best to have the pr number I think as this partially solves the discussion (the other half done by the validation hook Peter added in awkward).

[ikrommyd]

Thanks @Saransh-cpp for the feedback. I'll tackle it as soon as find a time slot.

There is one more case that I'd like to solve which was the original inspiration for this. That's when people do ak.zip and use with_name but I don't see at all how we can solve this easily. For example

In [12]: x = ak.zip({"pt": events.Electron.pt, "eta": events.Electron.eta, "phi": events.Electron.phi, "mass": events.Electron.mass, "energy": -999, "rho": 1001}, with_name="Momentum4D")

In [13]: x.pt
Out[13]: <Array [[1001], [], [], ..., [1001, 1001], [1001]] type='55342 * var * int64'>

In [14]: x.mass
Out[14]: <Array [[-3.98e+03], [], ..., [...], [-1.53e+03]] type='55342 * var * float64'>

People can assign nonsense fields together like rho and energy here without knowing that rho it is an alias for pt for example and that entirely messes up your 4-vector math (silently).
Or people can do jets["rho"] = something. But I really don't know how to prevent this. Awkward just blindly assigns a name without caring at all what that implies and without knowing at all what behaviors it carries over.
Even worse things happen when it comes straight from the root file. In CMS, we have a flavor of nanoaod that has pt, eta, phi, phi, mass and energy. But "energy" is not the "correct" energy you'd get from the pt-eta-phi-mass 4-vector. It's a different energy measurement. And when you define a 4-vector with those 5 using with_name, that's bad.
But I really don't know if it's even possible to solve this because it's awkward who is assigning the names and behaviors without any knowledge.

[Saransh-cpp]

I think @pfackeldey might be able to help with this - can we add constructor checks to our custom Awkward behaviors?

[pfackeldey]

I think @pfackeldey might be able to help with this - can we add constructor checks to our custom Awkward behaviors?

I'm not sure about the following:
I think you can't really change the way __init__ works through behaviors because awkward updates it's base class after the constructor has been run.

If my understanding here is wrong, you can do the following:

class MyBehavior:
  def __init__(self, *args, **kwargs):
    super().__init__(*args, **kwargs)
    # check
    if set(self._layout.fields) != {"pt", "eta", "phi", "mass"}:
      raise ValueError('got different fields than this behavior expected')

However, @ikrommyd and me talked about this and these constructor checks are indeed helpful, but do not prevent to break the behaviors. What would instead prevent this thoroughly would be:

1. Awkward Arrays are fully immutable (you can't even set/update/delete fields)
2. we disallow getattr to retrieve fields so it's clear to the user if they invoke a property or access a field
3. vector constructors are only allowed to have access to certain fixed fields

(1.) and (3.) solves this issue entirely I think, (2.) is in addition to make people aware of when they invoke the mass property and when they access the mass field.

This is of course major break in API, so I'm not sure if we can/want to do this. I'm adding this here to write down the fix that solves this issue thoroughly.
There could be ways to enforce this immutability wrt fields per array instance. Then we can preserve the current functionality in general awkward, but vector could enforce true immutability which does not allow to update/add/del fields to arrays with vector behaviors.
About (2.) I don't think there's really a way to solve this except for a API breakage in awkward. But this is not the issue, it's just hard to understand when one accesses a property and when a field through dot-access.

I could add a way in awkward-array to make them truly immutable that you can define through behaviors in vector. And in theory I could also think of a way to make post init constructor checks in awkward arrays possible for behaviors.

If you want @ikrommyd & @Saransh-cpp I can prepare a PR and you can have a look? Or what do you think?

[Saransh-cpp]

Sounds good to me, @pfackeldey! I'll be happy to review the PR in awkward and subsequently propagate the new immutability way to vector 🙂

[eduardo-rodrigues]

Hi all. The subject sounded interesting/perplexing hence I had a read.

Indeed good catches, @ikrommyd 👍.

Now, there are 2 different matters being mixed together, as it were, it seems to me:

The real problem of one being able to create a vector with some nonsensical combination is the main thing (to be) addressed in this PR, which is important to get fixed while providing relevant error messages to users so that they can see the issues in the way they wrote things. All good here and I won't comment on the details as @Saransh-cpp knows them infinitely better than me,

The other issue mentioned relates to constructions such as In [12]: x = ak.zip({"pt": events.Electron.pt, "eta": events.Electron.eta, "phi": events.Electron.phi, "mass": events.Electron.mass, "energy": -999, "rho": 1001}, with_name="Momentum4D"), which relates to vector only by the time the actual 4-vector is put together behind the scenes. At construction level a user could in principle assign events.Electron.pt to px, events.Electron.eta to phi, etc., if they had some crazy idea for renaming variables. It is only the with_name that will impose constraints on what variables to assign to what. Admittedly, in the case you @ikrommyd provide, it seems to me that the issue should be trivial to catch because 5 variables to make a 4-vector is over-constraining the system and should not be allowed (this comment holds for the "real problem" just above). Or am I reading too quickly and misundestanding things? In any case I don't think you can in general leverage 4-vector knowledge in the Vector package at the level of ak.zip since the latter is generic. Unless you explicitly make the connection awkward-vector with a with_name parameter that is a Vector class. Else you implicitly make an awkward coupling with subtle behaviours, I feel.

[ikrommyd]

Hi @eduardo-rodrigues, it's not trivial to catch when you do ak.zip though and you don't use any of the vector constructor methods. Because ak.zip does not know what Momentum4D means. It does not know or care what behaviors it carries over. More than that, we can hard-code the vector behavior checks explicitly in awkward-array sure but that's not the full story. You can subclass and define your own behaviors (coffea does that a lot). So then you can do with_name="PtEtaPhiMCandidate" which is something that coffea defines. So what behaviors do we check? We can't hard-code all the user-defined stuff? And with_name, all it does is it assigns a __record__ parameter to the array's layout to a string. It is completely blind as to what this implies.

[eduardo-rodrigues]

You are emphasising of my comments in some way, meaning that it is hard and/or dangerous to couple the 2 packages, but if you have no knowledge, then you can easily get into trouble. That's why I mentioned the possibility of having an explicit coupling using a class rather than a string, so with_name=vector.MomentumObject4D rather than with_name="Momentum4D". At least then somebody subclassing would be on its own, taking responsibility, and for other standard use cases, there would be some checks implemented. But I'm still likely being naive here and am showing how little I know about the internals of awkward, sorry if this is not helpful.

[ikrommyd]

Did a full write-up here of the issue: #660
This PR solves part of it. For further discussion, I suggest we move there to not make the PR discussion messy.

[eduardo-rodrigues]

It's a picky comment but it's a bit sub-optimal that we duplicate these "complaint strings" in several submodules. For better maintainability it's probably best to move them to a trival submodule and import the strings in the several places, such as here but also in awkward_constructions.py for example.

[ikrommyd]

is there a suggestion regarding where to put these complaint strings?

[ikrommyd]

maybe @Saransh-cpp ?

[Saransh-cpp]

I think we can put it in _methods.py as all the common methods/base classes and protocols live there. This exact string is used at several other places as well, so those can be refactored too (perhaps in another PR if this PR gets too big).

[ikrommyd]

Honestly, all "types" of vector constructors and backend literally validate the coordinates in different places each with their own implementations. I'd rather not do this here and just a minimal working case and do a whole refactoring of validation in a separate PR. obj numpy, awkward, numba, sympy vectors all validate in different places with slightly different mechanisms each. We need a general refactoring of the coordinate presence validation that is shared amongst all of those types of vectors.

[pfackeldey]

Linking my proposed solution here as well: #660 (comment)

[codecov]

Codecov Report

❌ Patch coverage is 98.38710% with 3 lines in your changes missing coverage. Please review.
✅ Project coverage is 88.85%. Comparing base (e79158a) to head (991bce7).

Files with missing lines	Patch %	Lines
src/vector/backends/_numba_object.py	87.50%	3 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #659      +/-   ##
==========================================
+ Coverage   87.72%   88.85%   +1.13%     
==========================================
  Files          96       96              
  Lines       11194    10769     -425     
==========================================
- Hits         9820     9569     -251     
+ Misses       1374     1200     -174     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[ikrommyd]

The more I investigate, the more horrible it becomes. There are a ton of inconsistencies between backends and some bugs regarding which coordinate combinations are allowed. Converting to draft. Need to come back with a different approach.

[ikrommyd]

@henryiii send an agent over for review. I did the rebasing and consolidation with a coder and an adverserial reviewer myself. The goal was to have the validation only in one place for all backends and to have 100% test coverage.