def _ttype_of(recordarraytype: typing.Any) -> typing.Any: import numba cls: typing.Union[typing.Type[TemporalObjectT], typing.Type[TemporalObjectTau], ] try: t_index = recordarraytype.recordlookup.index("t") except ValueError: t_index = None try: tau_index = recordarraytype.recordlookup.index("tau") except ValueError: tau_index = None if t_index is not None: coord1 = recordarraytype.contenttypes[t_index].arraytype.dtype cls = TemporalObjectT elif tau_index is not None: coord1 = recordarraytype.contenttypes[tau_index].arraytype.dtype cls = TemporalObjectTau else: raise numba.TypingError( f"{recordarraytype} is missing temporal fields: t or tau") return numba.typeof(cls(coord1.cast_python_value(0)))
def ga_or(a, b): if isinstance(a, MultiVectorType) and isinstance(b, MultiVectorType): if a.layout_type != b.layout_type: raise numba.TypingError( "MultiVector objects belong to different layouts") imt_func = a.layout_type.obj.imt_func def impl(a, b): return a.layout.MultiVector(imt_func(a.value, b.value)) return impl elif isinstance(a, types.abstract.Number) and isinstance( b, MultiVectorType): ret_type = np.result_type(_numpy_support.as_dtype(a), _numpy_support.as_dtype(b.value_type.dtype)) def impl(a, b): return b.layout.MultiVector(np.zeros_like(b.value, dtype=ret_type)) return impl elif isinstance(a, MultiVectorType) and isinstance(b, types.abstract.Number): ret_type = np.result_type(_numpy_support.as_dtype(a.value_type.dtype), _numpy_support.as_dtype(b)) def impl(a, b): return a.layout.MultiVector(np.zeros_like(a.value, dtype=ret_type)) return impl
def ga_xor(a, b): if isinstance(a, MultiVectorType) and isinstance(b, MultiVectorType): if a.layout_type != b.layout_type: raise numba.TypingError( "MultiVector objects belong to different layouts") omt_func = a.layout_type.obj.omt_func def impl(a, b): return a.layout.MultiVector(omt_func(a.value, b.value)) return impl elif isinstance(a, types.abstract.Number) and isinstance( b, MultiVectorType): def impl(a, b): return b.layout.MultiVector(b.value * a) return impl elif isinstance(a, MultiVectorType) and isinstance(b, types.abstract.Number): def impl(a, b): return a.layout.MultiVector(a.value * b) return impl
def ga_sub(a, b): if isinstance(a, MultiVectorType) and isinstance(b, MultiVectorType): if a.layout_type != b.layout_type: raise numba.TypingError( "MultiVector objects belong to different layouts") def impl(a, b): return a.layout.MultiVector(a.value - b.value) return impl elif isinstance(a, types.abstract.Number) and isinstance( b, MultiVectorType): scalar_index = b.layout_type.obj._basis_blade_order.bitmap_to_index[0] ret_type = np.result_type(_numpy_support.as_dtype(a), _numpy_support.as_dtype(b.value_type.dtype)) def impl(a, b): op = -b.value.astype(ret_type) op[scalar_index] += a return b.layout.MultiVector(op) return impl elif isinstance(a, MultiVectorType) and isinstance(b, types.abstract.Number): scalar_index = a.layout_type.obj._basis_blade_order.bitmap_to_index[0] ret_type = np.result_type(_numpy_support.as_dtype(a.value_type.dtype), _numpy_support.as_dtype(b)) def impl(a, b): op = a.value.astype(ret_type) op[scalar_index] -= b return a.layout.MultiVector(op) return impl
def _aztype_of(recordarraytype: typing.Any) -> typing.Any: import numba cls: typing.Union[typing.Type[AzimuthalObjectXY], typing.Type[AzimuthalObjectRhoPhi], ] try: x_index = recordarraytype.recordlookup.index("x") except ValueError: x_index = None try: y_index = recordarraytype.recordlookup.index("y") except ValueError: y_index = None try: rho_index = recordarraytype.recordlookup.index("rho") except ValueError: rho_index = None try: phi_index = recordarraytype.recordlookup.index("phi") except ValueError: phi_index = None if x_index is not None and y_index is not None: coord1 = recordarraytype.contenttypes[x_index].arraytype.dtype coord2 = recordarraytype.contenttypes[y_index].arraytype.dtype cls = AzimuthalObjectXY elif rho_index is not None and phi_index is not None: coord1 = recordarraytype.contenttypes[rho_index].arraytype.dtype coord2 = recordarraytype.contenttypes[phi_index].arraytype.dtype cls = AzimuthalObjectRhoPhi else: raise numba.TypingError( f"{recordarraytype} is missing azimuthal fields: x/y or rho/phi") return numba.typeof( cls(coord1.cast_python_value(0), coord2.cast_python_value(0)))
def _ltype_of(recordarraytype: typing.Any) -> typing.Any: import numba cls: typing.Union[typing.Type[LongitudinalObjectZ], typing.Type[LongitudinalObjectTheta], typing.Type[LongitudinalObjectEta], ] try: z_index = recordarraytype.recordlookup.index("z") except ValueError: z_index = None try: theta_index = recordarraytype.recordlookup.index("theta") except ValueError: theta_index = None try: eta_index = recordarraytype.recordlookup.index("eta") except ValueError: eta_index = None if z_index is not None: coord1 = recordarraytype.contenttypes[z_index].arraytype.dtype cls = LongitudinalObjectZ elif theta_index is not None: coord1 = recordarraytype.contenttypes[theta_index].arraytype.dtype cls = LongitudinalObjectTheta elif eta_index is not None: coord1 = recordarraytype.contenttypes[eta_index].arraytype.dtype cls = LongitudinalObjectEta else: raise numba.TypingError( f"{recordarraytype} is missing longitudinal fields: z or theta or eta" ) return numba.typeof(cls(coord1.cast_python_value(0)))