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 _getitem(array, where, is_momentum):
if isinstance(where, str):
if is_momentum:
where = _repr_momentum_to_generic.get(where, where)
return array.view(numpy.ndarray)[where]
else:
out = numpy.ndarray.__getitem__(array, where)
if isinstance(out, numpy.void):
azimuthal, longitudinal, temporal = None, None, None
if hasattr(array, "_azimuthal_type"):
azimuthal = array._azimuthal_type.ObjectClass(
*[out[x] for x in _coordinate_class_to_names[_aztype(array)]]
)
if hasattr(array, "_longitudinal_type"):
longitudinal = array._longitudinal_type.ObjectClass(
*[out[x] for x in _coordinate_class_to_names[_ltype(array)]]
)
if hasattr(array, "_temporal_type"):
temporal = array._temporal_type.ObjectClass(
*[out[x] for x in _coordinate_class_to_names[_ttype(array)]]
)
if temporal is not None:
return array.ObjectClass(azimuthal, longitudinal, temporal)
elif longitudinal is not None:
return array.ObjectClass(azimuthal, longitudinal)
elif azimuthal is not None:
return array.ObjectClass(azimuthal)
else:
return array.ObjectClass(*out.view(numpy.ndarray))
else:
return out
def __repr__(self):
aznames = _coordinate_class_to_names[_aztype(self)]
out = []
for x in aznames:
y = _repr_generic_to_momentum.get(x, x)
out.append(f"{y}={getattr(self.azimuthal, x)}")
return "vector.obj(" + ", ".join(out) + ")"
def dispatch(obj, v):
function, *returns = _from_signature(
__name__,
dispatch_map,
(
_aztype(v),
_ltype(v),
),
)
with numpy.errstate(all="ignore"):
return v._wrap_result(
function(
v.lib,
obj["xx"],
obj["xy"],
obj["xz"],
obj["yx"],
obj["yy"],
obj["yz"],
obj["zx"],
obj["zy"],
obj["zz"],
*v.azimuthal.elements,
*v.longitudinal.elements
),
returns,
)
def __repr__(self):
aznames = _coordinate_class_to_names[_aztype(self)]
lnames = _coordinate_class_to_names[_ltype(self)]
out = []
for x in aznames:
out.append(f"{x}={getattr(self.azimuthal, x)}")
for x in lnames:
out.append(f"{x}={getattr(self.longitudinal, x)}")
return "vector.obj(" + ", ".join(out) + ")"
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 dispatch(angle, v):
function, *returns = _from_signature(
__name__,
dispatch_map,
(
_aztype(v),
_ltype(v),
),
)
with numpy.errstate(all="ignore"):
return v._wrap_result(
function(v.lib, angle, *v.azimuthal.elements, *v.longitudinal.elements),
returns,
)
def dispatch(u, i, j, k, vec):
function, *returns = _from_signature(
__name__,
dispatch_map,
(
_aztype(vec),
_ltype(vec),
),
)
with numpy.errstate(all="ignore"):
return vec._wrap_result(
function(
vec.lib, u, i, j, k, *vec.azimuthal.elements, *vec.longitudinal.elements
),
returns,
)
def dispatch(factor, v):
function, *returns = _from_signature(__name__, dispatch_map, (_aztype(v),))
with numpy.errstate(all="ignore"):
return v._wrap_result(function(v.lib, factor, *v.azimuthal.elements), returns)
