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)] lnames = _coordinate_class_to_names[_ltype(self)] tnames = _coordinate_class_to_names[_ttype(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)}") for x in tnames: out.append(f"{x}={getattr(self.temporal, x)}") return "vector.obj(" + ", ".join(out) + ")"
def dispatch(v): function, *returns = _from_signature( __name__, dispatch_map, ( _aztype(v), _ltype(v), _ttype(v), ), ) with numpy.errstate(all="ignore"): return v._wrap_result( function(v.lib, *v.azimuthal.elements, *v.longitudinal.elements, *v.temporal.elements), returns, )
def dispatch(obj, v): function, *returns = _from_signature( __name__, dispatch_map, ( _aztype(v), _ltype(v), _ttype(v), ), ) with numpy.errstate(all="ignore"): return v._wrap_result( function(v.lib, obj["xx"], obj["xy"], obj["xz"], obj["xt"], obj["yx"], obj["yy"], obj["yz"], obj["yt"], obj["zx"], obj["zy"], obj["zz"], obj["zt"], obj["tx"], obj["ty"], obj["tz"], obj["tt"], *v.azimuthal.elements, *v.longitudinal.elements, *v.temporal.elements), returns, )
def _wrap_result(self, result, returns): if returns == [float] or returns == [bool]: return result elif returns == [AzimuthalXY] or returns == [AzimuthalRhoPhi]: result = _toarrays(result) dtype = [] for i, name in enumerate(_coordinate_class_to_names[returns[0]]): dtype.append((name, result[i].dtype)) for name in _coordinate_class_to_names[_ltype(self)]: dtype.append((name, self.dtype[name])) for name in _coordinate_class_to_names[_ttype(self)]: dtype.append((name, self.dtype[name])) out = numpy.empty(_shape_of(result), dtype=dtype) for i, name in enumerate(_coordinate_class_to_names[returns[0]]): out[name] = result[i] for name in _coordinate_class_to_names[_ltype(self)]: out[name] = self[name] for name in _coordinate_class_to_names[_ttype(self)]: out[name] = self[name] return out.view(type(self)) elif ( len(returns) == 2 and isinstance(returns[0], type) and issubclass(returns[0], Azimuthal) and isinstance(returns[1], type) and issubclass(returns[1], Longitudinal) ): result = _toarrays(result) dtype = [] i = 0 for name in _coordinate_class_to_names[returns[0]]: dtype.append((name, result[i].dtype)) i += 1 for name in _coordinate_class_to_names[returns[1]]: dtype.append((name, result[i].dtype)) i += 1 for name in _coordinate_class_to_names[_ttype(self)]: dtype.append((name, self.dtype[name])) out = numpy.empty(_shape_of(result), dtype=dtype) i = 0 for name in _coordinate_class_to_names[returns[0]]: out[name] = result[i] i += 1 for name in _coordinate_class_to_names[returns[1]]: out[name] = result[i] i += 1 for name in _coordinate_class_to_names[_ttype(self)]: out[name] = self[name] return out.view(type(self)) elif ( len(returns) == 3 and isinstance(returns[0], type) and issubclass(returns[0], Azimuthal) and isinstance(returns[1], type) and issubclass(returns[1], Longitudinal) ): result = _toarrays(result) dtype = [] i = 0 for name in _coordinate_class_to_names[returns[0]]: dtype.append((name, result[i].dtype)) i += 1 for name in _coordinate_class_to_names[returns[1]]: dtype.append((name, result[i].dtype)) i += 1 is_4d = False if isinstance(returns[2], type) and issubclass(returns[2], Temporal): is_4d = True for name in _coordinate_class_to_names[returns[2]]: dtype.append((name, result[i].dtype)) i += 1 elif returns[2] is not None: raise AssertionError(repr(type(returns[2]))) out = numpy.empty(_shape_of(result), dtype=dtype) i = 0 for name in _coordinate_class_to_names[returns[0]]: out[name] = result[i] i += 1 for name in _coordinate_class_to_names[returns[1]]: out[name] = result[i] i += 1 if is_4d: for name in _coordinate_class_to_names[returns[2]]: out[name] = result[i] i += 1 return out.view(type(self)) else: return out.view(self.ProjectionClass3D) else: raise AssertionError(repr(returns))