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)))