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