def dispatch(v1, v2):
function, *returns = _from_signature(
__name__,
dispatch_map,
(
_aztype(v1),
_ltype(v1),
_ttype(v1),
_aztype(v2),
_ltype(v2),
_ttype(v2),
),
)
with numpy.errstate(all="ignore"):
return _handler((v1, v2))._wrap_result(
function(
_lib_of((v1, v2)),
*v1.azimuthal.elements,
*v1.longitudinal.elements,
*v1.temporal.elements,
*v2.azimuthal.elements,
*v2.longitudinal.elements,
*v2.temporal.elements,
),
returns,
)
def dispatch(rtol, atol, equal_nan, v1, v2):
function, *returns = _from_signature(
__name__,
dispatch_map,
(
_aztype(v1),
_aztype(v2),
),
)
with numpy.errstate(all="ignore"):
return _handler((v1, v2))._wrap_result(
function(
_lib_of((v1, v2)),
rtol,
atol,
equal_nan,
*v1.azimuthal.elements,
*v2.azimuthal.elements,
),
returns,
)
def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
if not isinstance(_handler(inputs), VectorObject):
# Let the array-of-vectors object handle it.
return NotImplemented
if isinstance(self, Vector2D):
from vector.compute.planar import add, dot, equal, not_equal
from vector.compute.planar import rho as absolute
from vector.compute.planar import rho2 as absolute2
from vector.compute.planar import scale, subtract
elif isinstance(self, Vector3D):
from vector.compute.spatial import add, dot, equal
from vector.compute.spatial import mag as absolute
from vector.compute.spatial import mag2 as absolute2
from vector.compute.spatial import not_equal, scale, subtract
elif isinstance(self, Vector4D):
from vector.compute.lorentz import (
add,
dot,
equal,
not_equal,
scale,
subtract,
)
from vector.compute.lorentz import tau as absolute
from vector.compute.lorentz import tau2 as absolute2
outputs = kwargs.get("out", ())
if any(not isinstance(x, VectorObject) for x in outputs):
raise TypeError(
"ufunc operating on VectorObjects can only fill another VectorObject "
"with 'out' keyword")
if ufunc is numpy.absolute and len(inputs) == 1:
result = absolute.dispatch(inputs[0])
for output in outputs:
_replace_data(output, result)
return result
elif ufunc is numpy.add and len(inputs) == 2:
result = add.dispatch(inputs[0], inputs[1])
for output in outputs:
_replace_data(output, result)
return result
elif ufunc is numpy.subtract and len(inputs) == 2:
result = subtract.dispatch(inputs[0], inputs[1])
for output in outputs:
_replace_data(output, result)
return result
elif (ufunc is numpy.multiply
and not isinstance(inputs[0],
(Vector, Coordinates)) and len(inputs) == 2):
result = scale.dispatch(inputs[0], inputs[1])
for output in outputs:
_replace_data(output, result)
return result
elif (ufunc is numpy.multiply
and not isinstance(inputs[1],
(Vector, Coordinates)) and len(inputs) == 2):
result = scale.dispatch(inputs[1], inputs[0])
for output in outputs:
_replace_data(output, result)
return result
elif ufunc is numpy.negative and len(inputs) == 1:
result = scale.dispatch(-1, inputs[0])
for output in outputs:
_replace_data(output, result)
return result
elif ufunc is numpy.positive and len(inputs) == 1:
result = inputs[0]
for output in outputs:
_replace_data(output, result)
return result
elif (ufunc is numpy.true_divide
and not isinstance(inputs[1],
(Vector, Coordinates)) and len(inputs) == 2):
result = scale.dispatch(1 / inputs[1], inputs[0])
for output in outputs:
_replace_data(output, result)
return result
elif (ufunc is numpy.power and not isinstance(inputs[1],
(Vector, Coordinates))
and len(inputs) == 2):
result = absolute.dispatch(inputs[0])**inputs[1]
for output in outputs:
_replace_data(output, result)
return result
elif ufunc is numpy.square and len(inputs) == 1:
result = absolute2.dispatch(inputs[0])
for output in outputs:
_replace_data(output, result)
return result
elif ufunc is numpy.sqrt and len(inputs) == 1:
result = numpy.sqrt(absolute.dispatch(inputs[0]))
for output in outputs:
_replace_data(output, result)
return result
elif ufunc is numpy.cbrt and len(inputs) == 1:
result = numpy.cbrt(absolute.dispatch(inputs[0]))
for output in outputs:
_replace_data(output, result)
return result
elif ufunc is numpy.matmul and len(inputs) == 2:
result = dot.dispatch(inputs[0], inputs[1])
for output in outputs:
_replace_data(output, result)
return result
elif ufunc is numpy.equal and len(inputs) == 2:
result = equal.dispatch(inputs[0], inputs[1])
for output in outputs:
_replace_data(output, result)
return result
elif ufunc is numpy.not_equal and len(inputs) == 2:
result = not_equal.dispatch(inputs[0], inputs[1])
for output in outputs:
_replace_data(output, result)
return result
else:
return NotImplemented
