diff --git a/src/vector/_backends/numpy_.py b/src/vector/_backends/numpy_.py index 65cab36a..a432c74c 100644 --- a/src/vector/_backends/numpy_.py +++ b/src/vector/_backends/numpy_.py @@ -430,6 +430,15 @@ def __eq__(self, other: typing.Any) -> typing.Any: def __ne__(self, other: typing.Any) -> typing.Any: return numpy.not_equal(self, other) # type: ignore + def __reduce__(self) -> typing.Union[str, typing.Tuple[typing.Any, ...]]: + pickled_state = super().__reduce__() + new_state = (*pickled_state[2], self.__dict__) + return pickled_state[0], pickled_state[1], new_state + + def __setstate__(self, state: typing.Any) -> None: + self.__dict__.update(state[-1]) + super().__setstate__(state[0:-1]) # type: ignore + def __array_ufunc__( self, ufunc: typing.Any, @@ -665,6 +674,9 @@ def __new__(cls, *args: typing.Any, **kwargs: typing.Any) -> "VectorNumpy2D": return array.view(cls) def __array_finalize__(self, obj: typing.Any) -> None: + if obj is None: + return + if _has(self, ("x", "y")): self._azimuthal_type = AzimuthalNumpyXY elif _has(self, ("rho", "phi")): @@ -840,6 +852,9 @@ class MomentumNumpy2D(PlanarMomentum, VectorNumpy2D): # type: ignore dtype: "numpy.dtype[typing.Any]" def __array_finalize__(self, obj: typing.Any) -> None: + if obj is None: + return + self.dtype.names = tuple( _repr_momentum_to_generic.get(x, x) for x in (self.dtype.names or ()) ) @@ -882,6 +897,9 @@ def __new__(cls, *args: typing.Any, **kwargs: typing.Any) -> "VectorNumpy3D": return array.view(cls) def __array_finalize__(self, obj: typing.Any) -> None: + if obj is None: + return + if _has(self, ("x", "y")): self._azimuthal_type = AzimuthalNumpyXY elif _has(self, ("rho", "phi")): @@ -1076,6 +1094,9 @@ class MomentumNumpy3D(SpatialMomentum, VectorNumpy3D): # type: ignore dtype: "numpy.dtype[typing.Any]" def __array_finalize__(self, obj: typing.Any) -> None: + if obj is None: + return + self.dtype.names = tuple( _repr_momentum_to_generic.get(x, x) for x in (self.dtype.names or ()) ) @@ -1132,6 +1153,9 @@ def __new__(cls, *args: typing.Any, **kwargs: typing.Any) -> "VectorNumpy4D": return array.view(cls) def __array_finalize__(self, obj: typing.Any) -> None: + if obj is None: + return + if _has(self, ("x", "y")): self._azimuthal_type = AzimuthalNumpyXY elif _has(self, ("rho", "phi")): @@ -1349,6 +1373,9 @@ class MomentumNumpy4D(LorentzMomentum, VectorNumpy4D): # type: ignore dtype: "numpy.dtype[typing.Any]" def __array_finalize__(self, obj: typing.Any) -> None: + if obj is None: + return + self.dtype.names = tuple( _repr_momentum_to_generic.get(x, x) for x in (self.dtype.names or ()) ) diff --git a/tests/backends/test_numpy.py b/tests/backends/test_numpy.py index 55dddb33..f3ebfc01 100644 --- a/tests/backends/test_numpy.py +++ b/tests/backends/test_numpy.py @@ -4,6 +4,7 @@ # or https://github.com/scikit-hep/vector for details. import math +import pickle import numpy @@ -29,3 +30,99 @@ def test_rhophi(): assert numpy.allclose(array.y, [0, 1, 4]) assert numpy.allclose(array.rho, [0, 1, 5]) assert numpy.allclose(array.phi, [10, math.atan2(1, 0), math.atan2(4, 3)]) + + +def test_pickle_vector_numpy_2d(): + array = vector._backends.numpy_.VectorNumpy2D( + [(0, 0), (0, 1), (3, 4)], dtype=[("x", numpy.float64), ("y", numpy.float64)] + ) + + array_pickled = pickle.dumps(array) + array_new = pickle.loads(array_pickled) + + assert numpy.allclose(array_new.x, array.x) + assert numpy.allclose(array_new.y, array.y) + + +def test_pickle_momentum_numpy_2d(): + array = vector._backends.numpy_.MomentumNumpy2D( + [(0, 0), (0, 1), (3, 4)], dtype=[("rho", numpy.float64), ("phi", numpy.float64)] + ) + + array_pickled = pickle.dumps(array) + array_new = pickle.loads(array_pickled) + + assert numpy.allclose(array_new.rho, array.rho) + assert numpy.allclose(array_new.phi, array.phi) + + +def test_pickle_vector_numpy_3d(): + array = vector._backends.numpy_.VectorNumpy3D( + [(0, 0, 0), (0, 1, 1), (3, 4, 5)], + dtype=[("x", numpy.float64), ("y", numpy.float64), ("z", numpy.float64)], + ) + + array_pickled = pickle.dumps(array) + array_new = pickle.loads(array_pickled) + + assert numpy.allclose(array_new.x, array.x) + assert numpy.allclose(array_new.y, array.y) + assert numpy.allclose(array_new.z, array.z) + + +def test_pickle_momentum_numpy_3d(): + array = vector._backends.numpy_.MomentumNumpy3D( + [(0, 0, 0), (0, 1, 1), (3, 4, 5)], + dtype=[ + ("rho", numpy.float64), + ("phi", numpy.float64), + ("theta", numpy.float64), + ], + ) + + array_pickled = pickle.dumps(array) + array_new = pickle.loads(array_pickled) + + assert numpy.allclose(array_new.rho, array.rho) + assert numpy.allclose(array_new.phi, array.phi) + assert numpy.allclose(array_new.theta, array.theta) + + +def test_pickle_vector_numpy_4d(): + array = vector._backends.numpy_.VectorNumpy4D( + [(0, 0, 0, 0), (0, 1, 1, 1), (3, 4, 5, 6)], + dtype=[ + ("x", numpy.float64), + ("y", numpy.float64), + ("z", numpy.float64), + ("t", numpy.float64), + ], + ) + + array_pickled = pickle.dumps(array) + array_new = pickle.loads(array_pickled) + + assert numpy.allclose(array_new.x, array.x) + assert numpy.allclose(array_new.y, array.y) + assert numpy.allclose(array_new.z, array.z) + assert numpy.allclose(array_new.t, array.t) + + +def test_pickle_momentum_numpy_4d(): + array = vector._backends.numpy_.MomentumNumpy4D( + [(0, 0, 0, 0), (0, 1, 1, 1), (3, 4, 5, 6)], + dtype=[ + ("rho", numpy.float64), + ("phi", numpy.float64), + ("theta", numpy.float64), + ("tau", numpy.float64), + ], + ) + + array_pickled = pickle.dumps(array) + array_new = pickle.loads(array_pickled) + + assert numpy.allclose(array_new.rho, array.rho) + assert numpy.allclose(array_new.phi, array.phi) + assert numpy.allclose(array_new.theta, array.theta) + assert numpy.allclose(array_new.tau, array.tau)