Symbolic kinematics#
This report investigates issue compwa.github.io#56. The ideal solution would be to use only SymPy in the existing ampform.kinematics module. This has two benefits:
It allows computing kinematic variables from four-momenta with different computational back-ends.
Expressions for kinematic variable can be inspected through their LaTeX representation.
To simplify things, we investigate 1. by only lambdifying to NumPy. It should be relatively straightforward to lambdify to other back-ends like TensorFlow (as long as they support Einstein summation).
Test sample#
Data sample taken from this test in AmpForm and topology and expected angles taken from here.
topologies = create_isobar_topologies(4)
topology = topologies[1]
angles = _compute_helicity_angles(events, topology)
angles._DataSet__data
{'phi_1+2+3': ScalarSequence([ 2.79758029 2.51292308 -1.07396684],
'theta_1+2+3': ScalarSequence([2.72456853 3.03316287 0.69240082],
'phi_2+3,1+2+3': ScalarSequence([1.0436215 1.8734936 0.16073833],
'theta_2+3,1+2+3': ScalarSequence([2.45361589 1.40639741 0.98079245],
'phi_2,2+3,1+2+3': ScalarSequence([ 0.36955786 -1.68820498 0.63063002],
'theta_2,2+3,1+2+3': ScalarSequence([1.0924374 1.99375767 1.31959621]}
Einstein summation#
First challenge is to express the Einstein summation in the existing implementation in terms of SymPy. The aim is to render the expression resulting nicely as LaTeX while at the same time being able to lambdify the expression to efficient NumPy code. We do this by deriving from UnevaluatedExpression and using the decorator functions provided by the ampform.sympy module.
Define boost and rotation classes#
First, wrap rotations and boosts in a class so with a nice LaTeX printer. Later on, a NumPy printer method will be defined externally for each of them.
class BoostZ(UnevaluatedExpression):
def __new__(cls, beta: sp.Symbol, **hints) -> BoostZ:
return create_expression(cls, beta, **hints)
def as_explicit(self) -> sp.Expr:
gamma = 1 / sp.sqrt(1 - beta**2)
return sp.Matrix([
[gamma, 0, 0, -gamma * beta],
[0, 1, 0, 0],
[0, 0, 1, 0],
[-gamma * beta, 0, 0, gamma],
])
def _latex(self, printer, *args) -> str:
beta, *_ = self.args
beta = printer._print(beta)
return Rf"\boldsymbol{{B_z}}\left({beta}\right)"
class RotationY(UnevaluatedExpression):
def __new__(cls, angle: sp.Symbol, **hints) -> RotationY:
return create_expression(cls, angle, **hints)
def as_explicit(self) -> sp.Expr:
angle = self.args[0]
return sp.Matrix([
[1, 0, 0, 0],
[0, sp.cos(angle), 0, sp.sin(angle)],
[0, 0, 1, 0],
[0, -sp.sin(angle), 0, sp.cos(angle)],
])
def _latex(self, printer, *args) -> str:
angle, *_ = self.args
angle = printer._print(angle)
return Rf"\boldsymbol{{R_y}}\left({angle}\right)"
Define Einstein summation class#
Similarly, we define a ArrayMultiplication class that will eventually be lambdified to numpy.einsum().
Indeed an expression involving these classes looks nice on the top-level:
n_events = 3
momentum = sp.MatrixSymbol("p", m=n_events, n=4)
beta = sp.Symbol("beta")
phi = sp.Symbol("phi")
theta = sp.Symbol("theta")
boosted_momentum = ArrayMultiplication(
BoostZ(beta),
RotationY(-theta),
RotationZ(-phi),
momentum,
)
boosted_momentum
\[\displaystyle \boldsymbol{B_z}\left(\beta\right) \boldsymbol{R_y}\left(- \theta\right) \boldsymbol{R_z}\left(- \phi\right) p\]
Note
It could be that the above can be achieved with SymPy’s ArrayTensorProduct and ArrayContraction. See for instance sympy/sympy#22279.
Define lambdification#
Now we have a problem: lambdification does not work…
print_lambdify([beta, theta, momentum], boosted_momentum)
def _lambdifygenerated(beta, theta, p):
return ( # Not supported in Python with SciPy and NumPy:
# ArrayMultiplication
ArrayMultiplication(
BoostZ(beta), RotationY(-theta), RotationZ(-phi), p
)
)
But lambdification can be defined externally to both the SymPy library and the expression classes. Here’s an implementation for NumPy where we define the lambdification through the expression class (with a _numpycode method):
def print_as_numpy(self, printer: Printer, *args) -> str:
def multiply(matrix, vector):
return (
'einsum("ij...,j...",'
f" transpose({matrix}, axes=(1, 2, 0)),"
f" transpose({vector}))"
)
def recursive_multiply(tensors):
if len(tensors) < 2:
msg = "Need at least two tensors"
raise ValueError(msg)
if len(tensors) == 2:
return multiply(tensors[0], tensors[1])
return multiply(tensors[0], recursive_multiply(tensors[1:]))
printer.module_imports["numpy"].update({"einsum", "transpose"})
tensors = list(map(printer._print, self.args))
if len(tensors) == 0:
return ""
if len(tensors) == 1:
return tensors[0]
return recursive_multiply(tensors)
ArrayMultiplication._numpycode = print_as_numpy
print_lambdify(
symbols=[beta, theta, momentum],
expr=ArrayMultiplication(beta, theta, momentum),
)
def _lambdifygenerated(beta, theta, p):
return einsum(
"ij...,j...",
transpose(beta, axes=(1, 2, 0)),
transpose(
einsum(
"ij...,j...", transpose(theta, axes=(1, 2, 0)), transpose(p)
)
),
)
This also needs to be done for the rotation and boost classes:
print_lambdify([beta, theta, momentum], boosted_momentum)
def _lambdifygenerated(beta, theta, p):
return ( # Not supported in Python with SciPy and NumPy:
# BoostZ
# RotationY
# RotationZ
einsum(
"ij...,j...",
transpose(BoostZ(beta), axes=(1, 2, 0)),
transpose(
einsum(
"ij...,j...",
transpose(RotationY(-theta), axes=(1, 2, 0)),
transpose(
einsum(
"ij...,j...",
transpose(RotationZ(-phi), axes=(1, 2, 0)),
transpose(p),
)
),
)
),
)
)
This time, we define the lambdification through the printer class:
def _print_BoostZ(self: NumPyPrinter, expr: BoostZ) -> str:
self.module_imports["numpy"].update({"array", "ones", "zeros", "sqrt"})
arg = expr.args[0]
beta = self._print(arg)
gamma = f"1 / sqrt(1 - ({beta}) ** 2)"
n_events = f"len({beta})"
zeros = f"zeros({n_events})"
ones = f"ones({n_events})"
return f"""array(
[
[{gamma}, {zeros}, {zeros}, -{gamma} * {beta}],
[{zeros}, {ones}, {zeros}, {zeros}],
[{zeros}, {zeros}, {ones}, {zeros}],
[-{gamma} * {beta}, {zeros}, {zeros}, {gamma}],
]
).transpose(2, 0, 1)"""
NumPyPrinter._print_BoostZ = _print_BoostZ
def _print_RotationY(self: NumPyPrinter, expr: RotationY) -> str:
self.module_imports["numpy"].update({"array", "cos", "ones", "zeros", "sin"})
arg = expr.args[0]
angle = self._print(arg)
n_events = f"len({angle})"
zeros = f"zeros({n_events})"
ones = f"ones({n_events})"
return f"""array(
[
[{ones}, {zeros}, {zeros}, {zeros}],
[{zeros}, cos({angle}), {zeros}, sin({angle})],
[{zeros}, {zeros}, {ones}, {zeros}],
[{zeros}, -sin({angle}), {zeros}, cos({angle})],
]
).transpose(2, 0, 1)"""
NumPyPrinter._print_RotationY = _print_RotationY
print_lambdify([beta, theta, momentum], boosted_momentum)
def _lambdifygenerated(beta, theta, p):
return einsum(
"ij...,j...",
transpose(
array(
[
[
1 / sqrt(1 - (beta) ** 2),
zeros(len(beta)),
zeros(len(beta)),
-1 / sqrt(1 - (beta) ** 2) * beta,
],
[
zeros(len(beta)),
ones(len(beta)),
zeros(len(beta)),
zeros(len(beta)),
],
[
zeros(len(beta)),
zeros(len(beta)),
ones(len(beta)),
zeros(len(beta)),
],
[
-1 / sqrt(1 - (beta) ** 2) * beta,
zeros(len(beta)),
zeros(len(beta)),
1 / sqrt(1 - (beta) ** 2),
],
]
).transpose(2, 0, 1),
axes=(1, 2, 0),
),
transpose(
einsum(
"ij...,j...",
transpose(
array(
[
[
ones(len(-theta)),
zeros(len(-theta)),
zeros(len(-theta)),
zeros(len(-theta)),
],
[
zeros(len(-theta)),
cos(-theta),
zeros(len(-theta)),
sin(-theta),
],
[
zeros(len(-theta)),
zeros(len(-theta)),
ones(len(-theta)),
zeros(len(-theta)),
],
[
zeros(len(-theta)),
-sin(-theta),
zeros(len(-theta)),
cos(-theta),
],
]
).transpose(2, 0, 1),
axes=(1, 2, 0),
),
transpose(
einsum(
"ij...,j...",
transpose(
array(
[
[
ones(len(-phi)),
zeros(len(-phi)),
zeros(len(-phi)),
zeros(len(-phi)),
],
[
zeros(len(-phi)),
cos(-phi),
-sin(-phi),
zeros(len(-phi)),
],
[
zeros(len(-phi)),
sin(-phi),
cos(-phi),
zeros(len(-phi)),
],
[
zeros(len(-phi)),
zeros(len(-phi)),
zeros(len(-phi)),
ones(len(-phi)),
],
]
).transpose(2, 0, 1),
axes=(1, 2, 0),
),
transpose(p),
)
),
)
),
)
Note
The code above contains a lot of duplicate code, such as len(-phi). This could possibly be improved with CodeBlock. See ampform#166.
Angle computation#
Computing phi#
The simplest angle to compute is \(\phi\), because it’s simply \(\phi=\arctan(p_y, p_x)\), with \(p\) a four-momentum (see existing implementation). This means we would need a way to represent \(p_x\) and \(p_y\) and some container for \(\phi\) itself. For convenience, we define \(p_z\) and \(E_p\) as well.
@implement_doit_method()
class FourMomentumX(UnevaluatedExpression):
momentum: ArraySymbol = property(lambda self: self.args[0])
def __new__(cls, momentum: ArraySymbol, **hints) -> FourMomentumX:
return create_expression(cls, momentum, **hints)
def evaluate(self) -> ArraySlice:
return ArraySlice(self.momentum, (slice(None), 1))
def _latex(self, printer, *args) -> str:
momentum = printer._print(self.momentum)
return f"{{{momentum}}}_{{x}}"
@implement_doit_method()
class FourMomentumY(UnevaluatedExpression):
momentum: ArraySymbol = property(lambda self: self.args[0])
def __new__(cls, momentum: ArraySymbol, **hints) -> FourMomentumY:
return create_expression(cls, momentum, **hints)
def evaluate(self) -> ArraySlice:
return ArraySlice(self.momentum, (slice(None), 2))
def _latex(self, printer, *args) -> str:
momentum = printer._print(self.momentum)
return f"{{{momentum}}}_{{y}}"
@implement_doit_method()
class FourMomentumZ(UnevaluatedExpression):
momentum: ArraySymbol = property(lambda self: self.args[0])
def __new__(cls, momentum: ArraySymbol, **hints) -> FourMomentumZ:
return create_expression(cls, momentum, **hints)
def evaluate(self) -> ArraySlice:
return ArraySlice(self.momentum, (slice(None), 3))
def _latex(self, printer, *args) -> str:
momentum = printer._print(self.momentum)
return f"{{{momentum}}}_{{y}}"
@implement_doit_method()
class Energy(UnevaluatedExpression):
momentum: ArraySymbol = property(lambda self: self.args[0])
def __new__(cls, momentum: ArraySymbol, **hints) -> Energy:
return create_expression(cls, momentum, **hints)
def evaluate(self) -> ArraySlice:
return ArraySlice(self.momentum, (slice(None), 0))
def _latex(self, printer, *args) -> str:
momentum = printer._print(self.momentum)
return f"{{E}}_{{{momentum}}}"
@implement_doit_method()
class Phi(UnevaluatedExpression):
momentum: ArraySymbol = property(lambda self: self.args[0])
def __new__(cls, momentum: ArraySymbol, **hints) -> Phi:
return create_expression(cls, momentum, **hints)
def evaluate(self) -> sp.Expr:
p = self.momentum
return sp.atan2(FourMomentumY(p), FourMomentumX(p))
def _latex(self, printer, *args) -> str:
momentum = printer._print(self.momentum)
return Rf"\phi\left({momentum}\right)"
The classes indeed render nicely as LaTeX:
\[\displaystyle p_{1}\]
\[\displaystyle \begin{eqnarray}
\phi\left(p_{1}\right) & = & \operatorname{atan_{2}}{\left({p_{1}}_{y},{p_{1}}_{x} \right)}
\end{eqnarray}\]
\[\displaystyle \begin{eqnarray}
{p_{1}}_{x} & = & p_{1}\left[:, 1\right]
\end{eqnarray}\]
\[\displaystyle \begin{eqnarray}
{p_{1}}_{y} & = & p_{1}\left[:, 2\right]
\end{eqnarray}\]
\[\displaystyle \begin{eqnarray}
{E}_{p_{1}} & = & p_{1}\left[:, 0\right]
\end{eqnarray}\]
Note that the four classes like FourMomentumX expect an ArraySymbol as input. This requires sympy/sympy#22265. This allows lambdifying the above expressions to valid NumPy code. Let’s compare this with the existing implementation using the Test sample:
momentum_sample = events[0]
np.array(momentum_sample)
array([[ 1.35527 , 0.514208 , -0.184219 , 1.23296 ],
[ 0.841933 , 0.0727385, -0.0528868, 0.826163 ],
[ 0.550927 , -0.162529 , 0.29976 , -0.411133 ]])
np_expr = sp.lambdify(p, p_x.doit())
np_expr(momentum_sample)
array([ 0.514208 , 0.0727385, -0.162529 ])
np_expr = sp.lambdify(p, energy.doit())
np_expr(momentum_sample)
array([1.35527 , 0.841933, 0.550927])
Computing theta#
Computing \(\theta\) is more complicated, because requires the norm of the three-momentum of \(p\) (see existing implementation). In other words:
\[
\theta = \arccos\left(\frac{p_z}{\left|\vec{p}\right|}\right)
\quad \mathrm{with} \quad
\left|\vec{p}\right| = \sqrt{p_x^2+p_y^2+p_z^2}
\]
The complication here is that \(\left|\vec{p}\right|\) needs to be computed with np.sum over a slice of the arrays. As of writing, it is not yet possible to write a SymPy expression that can lambdify to this or an equivalent with np.einsum (sympy/sympy#22279). So for now, an intermediate ArrayAxisSum class has to be written for this. See also this remark.
class ArrayAxisSum(sp.Expr):
array: ArraySymbol = property(lambda self: self.args[0])
axis: int | None = property(lambda self: self.args[1])
def __new__(
cls, array: ArraySymbol, axis: int | None = None, **hints
) -> ArrayAxisSum:
if axis is not None and not isinstance(axis, (int, sp.Integer)):
msg = "Only single digits allowed for axis"
raise TypeError(msg)
return create_expression(cls, array, axis, **hints)
def _latex(self, printer, *args) -> str:
A = printer._print(self.array)
if self.axis is None:
return Rf"\sum{{{A}}}"
axis = printer._print(self.axis)
return Rf"\sum_{{\mathrm{{axis{axis}}}}}{{{A}}}"
Looks nice as LaTeX:
A = ArraySymbol("A")
display(
ArrayAxisSum(A, axis=1),
ArrayAxisSum(A),
)
\[\displaystyle \sum_{\mathrm{axis1}}{A}\]
\[\displaystyle \sum{A}\]
Now let’s define a printer method for NumPy:
def _print_ArrayAxisSum(self: NumPyPrinter, expr: ArrayAxisSum) -> str:
self.module_imports["numpy"].add("sum")
array = self._print(expr.array)
axis = self._print(expr.axis)
return f"sum({array}, axis={axis})"
NumPyPrinter._print_ArrayAxisSum = _print_ArrayAxisSum
print_lambdify(A, ArrayAxisSum(A, axis=1))
def _lambdifygenerated(A):
return sum(A, axis=1)
…and let’s check whether it works as expected for a 3-dimensional array:
array([[[ 0, 1],
[ 2, 3],
[ 4, 5]],
[[ 6, 7],
[ 8, 9],
[10, 11]]])
np_expr = sp.lambdify(A, ArrayAxisSum(A))
np_expr(array)
66
np_expr = sp.lambdify(A, ArrayAxisSum(A, axis=0))
np_expr(array)
array([[ 6, 8],
[10, 12],
[14, 16]])
np_expr = sp.lambdify(A, ArrayAxisSum(A, axis=1))
np_expr(array)
array([[ 6, 9],
[24, 27]])
np_expr = sp.lambdify(A, ArrayAxisSum(A, axis=2))
np_expr(array)
array([[ 1, 5, 9],
[13, 17, 21]])
Now we’re ready to define a class that can represent \(\left|\vec{p}\right|\) and that lambdifies to the expressions given by the existing implementation.
@implement_doit_method()
class ThreeMomentumNorm(UnevaluatedExpression):
momentum: ArraySymbol = property(lambda self: self.args[0])
def __new__(cls, momentum: ArraySymbol, **hints) -> ThreeMomentumNorm:
return create_expression(cls, momentum, **hints)
def evaluate(self) -> ArraySlice:
three_momentum = ArraySlice(self.momentum, (slice(None), slice(1, None)))
norm_squared = ArrayAxisSum(three_momentum**2, axis=1)
return sp.sqrt(norm_squared)
def _latex(self, printer, *args) -> str:
three_momentum = printer._print(self.momentum)
return Rf"\left|\vec{{{three_momentum}}}\right|"
def _numpycode(self, printer, *args) -> str:
return printer._print(self.evaluate())
\[\displaystyle \begin{eqnarray}
\left|\vec{p_{1}}\right| & = & \sqrt{\sum_{\mathrm{axis1}}{p_{1}\left[:, 1:\right]^{2}}}
\end{eqnarray}\]
np_expr = sp.lambdify(p, p_norm.doit())
np_expr(momentum_sample)
array([1.34853137, 0.83104344, 0.53413676])
With that, we’re ready to define the Theta class!
@implement_doit_method()
class Theta(UnevaluatedExpression):
momentum: ArraySymbol = property(lambda self: self.args[0])
def __new__(cls, momentum: ArraySymbol, **hints) -> Theta:
return create_expression(cls, momentum, **hints)
def evaluate(self) -> sp.Expr:
p = self.momentum
return sp.acos(FourMomentumZ(p) / ThreeMomentumNorm(p))
def _latex(self, printer, *args) -> str:
momentum = printer._print(self.momentum)
return Rf"\theta\left({momentum}\right)"
The math doesn’t look the best, but this due to the ArrayMultiplication class.
\[\displaystyle \begin{eqnarray}
\theta\left(p_{1}\right) & = & \operatorname{acos}{\left(p_{1}\left[:, 3\right] \frac{1}{\sqrt{\sum_{\mathrm{axis1}}{p_{1}\left[:, 1:\right]^{2}}}} \right)}
\end{eqnarray}\]
At any rate, when we lambdify the whole thing, we have exactly the same result as the original theta() method gave!
Recursive angle computation#
Finally, we are ready to compute kinematic angles recursively from a Topology, just as in the existing implementation, but now with SymPy only.
{0: p0, 1: p1, 2: p2, 3: p3}
Now it’s quite trivial to rewrite the existing implementation with the classes defined above.
def compute_helicity_angles(
events: dict[int, ArraySymbol], topology: Topology
) -> dict[str, sp.Expr]:
if topology.outgoing_edge_ids != set(events):
msg = (
f"Momentum IDs {set(events)} do not match final state edge IDs"
f" {set(topology.outgoing_edge_ids)}"
)
raise ValueError(msg)
def __recursive_helicity_angles(
events: dict[int, ArraySymbol], node_id: int
) -> dict[str, sp.Expr]:
helicity_angles: dict[str, sp.Expr] = {}
child_state_ids = sorted(topology.get_edge_ids_outgoing_from_node(node_id))
if all(topology.edges[i].ending_node_id is None for i in child_state_ids):
state_id = child_state_ids[0]
four_momentum = events[state_id]
phi_label, theta_label = get_helicity_angle_label(topology, state_id)
helicity_angles[phi_label] = Phi(four_momentum)
helicity_angles[theta_label] = Theta(four_momentum)
for state_id in child_state_ids:
edge = topology.edges[state_id]
if edge.ending_node_id is not None:
# recursively determine all momenta ids in the list
sub_momenta_ids = determine_attached_final_state(topology, state_id)
if len(sub_momenta_ids) > 1:
# add all of these momenta together -> defines new subsystem
four_momentum = sum(events[i] for i in sub_momenta_ids)
# boost all of those momenta into this new subsystem
phi = Phi(four_momentum)
theta = Theta(four_momentum)
p3_norm = ThreeMomentumNorm(four_momentum)
beta = p3_norm / Energy(four_momentum)
new_momentum_pool = {
k: ArrayMultiplication(
BoostZ(beta),
RotationY(-theta),
RotationZ(-phi),
p,
)
for k, p in events.items()
if k in sub_momenta_ids
}
# register current angle variables
phi_label, theta_label = get_helicity_angle_label(
topology, state_id
)
helicity_angles[phi_label] = Phi(four_momentum)
helicity_angles[theta_label] = Theta(four_momentum)
# call next recursion
angles = __recursive_helicity_angles(
new_momentum_pool,
edge.ending_node_id,
)
helicity_angles.update(angles)
return helicity_angles
initial_state_id = next(iter(topology.incoming_edge_ids))
initial_state_edge = topology.edges[initial_state_id]
assert initial_state_edge.ending_node_id is not None
return __recursive_helicity_angles(events, initial_state_edge.ending_node_id)
The computation works indeed and can be rendered to both LaTeX and numpy!
symbolic_angles = compute_helicity_angles(momentum_symbols, topology)
list(symbolic_angles)
['phi_1+2+3',
'theta_1+2+3',
'phi_2+3,1+2+3',
'theta_2+3,1+2+3',
'phi_2,2+3,1+2+3',
'theta_2,2+3,1+2+3']
\[\displaystyle \begin{eqnarray}
\phi\left(p_{1} + p_{2} + p_{3}\right) & = & \operatorname{atan_{2}}{\left(\left(p_{1} + p_{2} + p_{3}\right)\left[:, 2\right],\left(p_{1} + p_{2} + p_{3}\right)\left[:, 1\right] \right)}
\end{eqnarray}\]
def _lambdifygenerated(p0, p1, p2, p3):
return arctan2((p1 + p2 + p3)[:, 2], (p1 + p2 + p3)[:, 1])
\[\displaystyle \begin{eqnarray}
\theta\left(p_{1} + p_{2} + p_{3}\right) & = & \operatorname{acos}{\left(\left(p_{1} + p_{2} + p_{3}\right)\left[:, 3\right] \frac{1}{\sqrt{\sum_{\mathrm{axis1}}{\left(p_{1} + p_{2} + p_{3}\right)\left[:, 1:\right]^{2}}}} \right)}
\end{eqnarray}\]
def _lambdifygenerated(p0, p1, p2, p3):
return arccos(
(p1 + p2 + p3)[:, 3]
* sum((p1 + p2 + p3)[:, 1:] ** 2, axis=1) ** (-1 / 2)
)
\[\displaystyle \begin{eqnarray}
\phi\left(\boldsymbol{B_z}\left(\left|\vec{p_{1} + p_{2} + p_{3}}\right| {E}_{p_{1} + p_{2} + p_{3}}^{-1}\right) \boldsymbol{R_y}\left(- \theta\left(p_{1} + p_{2} + p_{3}\right)\right) \boldsymbol{R_z}\left(- \phi\left(p_{1} + p_{2} + p_{3}\right)\right) p_{2} + \boldsymbol{B_z}\left(\left|\vec{p_{1} + p_{2} + p_{3}}\right| {E}_{p_{1} + p_{2} + p_{3}}^{-1}\right) \boldsymbol{R_y}\left(- \theta\left(p_{1} + p_{2} + p_{3}\right)\right) \boldsymbol{R_z}\left(- \phi\left(p_{1} + p_{2} + p_{3}\right)\right) p_{3}\right) & = & \operatorname{atan_{2}}{\left(\left(\boldsymbol{B_z}\left(\sqrt{\sum_{\mathrm{axis1}}{\left(p_{1} + p_{2} + p_{3}\right)\left[:, 1:\right]^{2}}} \left(p_{1} + p_{2} + p_{3}\right)\left[:, 0\right]^{-1}\right) \boldsymbol{R_y}\left(- \operatorname{acos}{\left(\left(p_{1} + p_{2} + p_{3}\right)\left[:, 3\right] \frac{1}{\sqrt{\sum_{\mathrm{axis1}}{\left(p_{1} + p_{2} + p_{3}\right)\left[:, 1:\right]^{2}}}} \right)}\right) \boldsymbol{R_z}\left(- \operatorname{atan_{2}}{\left(\left(p_{1} + p_{2} + p_{3}\right)\left[:, 2\right],\left(p_{1} + p_{2} + p_{3}\right)\left[:, 1\right] \right)}\right) p_{2} + \boldsymbol{B_z}\left(\sqrt{\sum_{\mathrm{axis1}}{\left(p_{1} + p_{2} + p_{3}\right)\left[:, 1:\right]^{2}}} \left(p_{1} + p_{2} + p_{3}\right)\left[:, 0\right]^{-1}\right) \boldsymbol{R_y}\left(- \operatorname{acos}{\left(\left(p_{1} + p_{2} + p_{3}\right)\left[:, 3\right] \frac{1}{\sqrt{\sum_{\mathrm{axis1}}{\left(p_{1} + p_{2} + p_{3}\right)\left[:, 1:\right]^{2}}}} \right)}\right) \boldsymbol{R_z}\left(- \operatorname{atan_{2}}{\left(\left(p_{1} + p_{2} + p_{3}\right)\left[:, 2\right],\left(p_{1} + p_{2} + p_{3}\right)\left[:, 1\right] \right)}\right) p_{3}\right)\left[:, 2\right],\left(\boldsymbol{B_z}\left(\sqrt{\sum_{\mathrm{axis1}}{\left(p_{1} + p_{2} + p_{3}\right)\left[:, 1:\right]^{2}}} \left(p_{1} + p_{2} + p_{3}\right)\left[:, 0\right]^{-1}\right) \boldsymbol{R_y}\left(- \operatorname{acos}{\left(\left(p_{1} + p_{2} + p_{3}\right)\left[:, 3\right] \frac{1}{\sqrt{\sum_{\mathrm{axis1}}{\left(p_{1} + p_{2} + p_{3}\right)\left[:, 1:\right]^{2}}}} \right)}\right) \boldsymbol{R_z}\left(- \operatorname{atan_{2}}{\left(\left(p_{1} + p_{2} + p_{3}\right)\left[:, 2\right],\left(p_{1} + p_{2} + p_{3}\right)\left[:, 1\right] \right)}\right) p_{2} + \boldsymbol{B_z}\left(\sqrt{\sum_{\mathrm{axis1}}{\left(p_{1} + p_{2} + p_{3}\right)\left[:, 1:\right]^{2}}} \left(p_{1} + p_{2} + p_{3}\right)\left[:, 0\right]^{-1}\right) \boldsymbol{R_y}\left(- \operatorname{acos}{\left(\left(p_{1} + p_{2} + p_{3}\right)\left[:, 3\right] \frac{1}{\sqrt{\sum_{\mathrm{axis1}}{\left(p_{1} + p_{2} + p_{3}\right)\left[:, 1:\right]^{2}}}} \right)}\right) \boldsymbol{R_z}\left(- \operatorname{atan_{2}}{\left(\left(p_{1} + p_{2} + p_{3}\right)\left[:, 2\right],\left(p_{1} + p_{2} + p_{3}\right)\left[:, 1\right] \right)}\right) p_{3}\right)\left[:, 1\right] \right)}
\end{eqnarray}\]
def _lambdifygenerated(p0, p1, p2, p3):
return arctan2(
(
einsum(
"ij...,j...",
transpose(
array(
[
[
1
/ sqrt(
1
- (
sqrt(
sum(
(p1 + p2 + p3)[:, 1:] ** 2,
axis=1,
)
)
* (p1 + p2 + p3)[:, 0] ** (-1.0)
)
** 2
),
zeros(
len(
sqrt(
sum(
(p1 + p2 + p3)[:, 1:] ** 2,
axis=1,
)
)
* (p1 + p2 + p3)[:, 0] ** (-1.0)
)
),
zeros(
len(
sqrt(
sum(
(p1 + p2 + p3)[:, 1:] ** 2,
axis=1,
)
)
* (p1 + p2 + p3)[:, 0] ** (-1.0)
)
),
-1
/ sqrt(
1
- (
sqrt(
sum(
(p1 + p2 + p3)[:, 1:] ** 2,
axis=1,
)
)
* (p1 + p2 + p3)[:, 0] ** (-1.0)
)
** 2
)
* sqrt(sum((p1 + p2 + p3)[:, 1:] ** 2, axis=1))
* (p1 + p2 + p3)[:, 0] ** (-1.0),
],
[
zeros(
len(
sqrt(
sum(
(p1 + p2 + p3)[:, 1:] ** 2,
axis=1,
)
)
* (p1 + p2 + p3)[:, 0] ** (-1.0)
)
),
ones(
len(
sqrt(
sum(
(p1 + p2 + p3)[:, 1:] ** 2,
axis=1,
)
)
* (p1 + p2 + p3)[:, 0] ** (-1.0)
)
),
zeros(
len(
sqrt(
sum(
(p1 + p2 + p3)[:, 1:] ** 2,
axis=1,
)
)
* (p1 + p2 + p3)[:, 0] ** (-1.0)
)
),
zeros(
len(
sqrt(
sum(
(p1 + p2 + p3)[:, 1:] ** 2,
axis=1,
)
)
* (p1 + p2 + p3)[:, 0] ** (-1.0)
)
),
],
[
zeros(
len(
sqrt(
sum(
(p1 + p2 + p3)[:, 1:] ** 2,
axis=1,
)
)
* (p1 + p2 + p3)[:, 0] ** (-1.0)
)
),
zeros(
len(
sqrt(
sum(
(p1 + p2 + p3)[:, 1:] ** 2,
axis=1,
)
)
* (p1 + p2 + p3)[:, 0] ** (-1.0)
)
),
ones(
len(
sqrt(
sum(
(p1 + p2 + p3)[:, 1:] ** 2,
axis=1,
)
)
* (p1 + p2 + p3)[:, 0] ** (-1.0)
)
),
zeros(
len(
sqrt(
sum(
(p1 + p2 + p3)[:, 1:] ** 2,
axis=1,
)
)
* (p1 + p2 + p3)[:, 0] ** (-1.0)
)
),
],
[
-1
/ sqrt(
1
- (
sqrt(
sum(
(p1 + p2 + p3)[:, 1:] ** 2,
axis=1,
)
)
* (p1 + p2 + p3)[:, 0] ** (-1.0)
)
** 2
)
* sqrt(sum((p1 + p2 + p3)[:, 1:] ** 2, axis=1))
* (p1 + p2 + p3)[:, 0] ** (-1.0),
zeros(
len(
sqrt(
sum(
(p1 + p2 + p3)[:, 1:] ** 2,
axis=1,
)
)
* (p1 + p2 + p3)[:, 0] ** (-1.0)
)
),
zeros(
len(
sqrt(
sum(
(p1 + p2 + p3)[:, 1:] ** 2,
axis=1,
)
)
* (p1 + p2 + p3)[:, 0] ** (-1.0)
)
),
1
/ sqrt(
1
- (
sqrt(
sum(
(p1 + p2 + p3)[:, 1:] ** 2,
axis=1,
)
)
* (p1 + p2 + p3)[:, 0] ** (-1.0)
)
** 2
),
],
]
).transpose(2, 0, 1),
axes=(1, 2, 0),
),
transpose(
einsum(
"ij...,j...",
transpose(
array(
[
[
ones(
len(
-arccos(
(p1 + p2 + p3)[:, 3]
* sum(
(p1 + p2 + p3)[:, 1:]
** 2,
axis=1,
)
** (-1 / 2)
)
)
),
zeros(
len(
-arccos(
(p1 + p2 + p3)[:, 3]
* sum(
(p1 + p2 + p3)[:, 1:]
** 2,
axis=1,
)
** (-1 / 2)
)
)
),
zeros(
len(
-arccos(
(p1 + p2 + p3)[:, 3]
* sum(
(p1 + p2 + p3)[:, 1:]
** 2,
axis=1,
)
** (-1 / 2)
)
)
),
zeros(
len(
-arccos(
(p1 + p2 + p3)[:, 3]
* sum(
(p1 + p2 + p3)[:, 1:]
** 2,
axis=1,
)
** (-1 / 2)
)
)
),
],
[
zeros(
len(
-arccos(
(p1 + p2 + p3)[:, 3]
* sum(
(p1 + p2 + p3)[:, 1:]
** 2,
axis=1,
)
** (-1 / 2)
)
)
),
cos(
-arccos(
(p1 + p2 + p3)[:, 3]
* sum(
(p1 + p2 + p3)[:, 1:] ** 2,
axis=1,
)
** (-1 / 2)
)
),
zeros(
len(
-arccos(
(p1 + p2 + p3)[:, 3]
* sum(
(p1 + p2 + p3)[:, 1:]
** 2,
axis=1,
)
** (-1 / 2)
)
)
),
sin(
-arccos(
(p1 + p2 + p3)[:, 3]
* sum(
(p1 + p2 + p3)[:, 1:] ** 2,
axis=1,
)
** (-1 / 2)
)
),
],
[
zeros(
len(
-arccos(
(p1 + p2 + p3)[:, 3]
* sum(
(p1 + p2 + p3)[:, 1:]
** 2,
axis=1,
)
** (-1 / 2)
)
)
),
zeros(
len(
-arccos(
(p1 + p2 + p3)[:, 3]
* sum(
(p1 + p2 + p3)[:, 1:]
** 2,
axis=1,
)
** (-1 / 2)
)
)
),
ones(
len(
-arccos(
(p1 + p2 + p3)[:, 3]
* sum(
(p1 + p2 + p3)[:, 1:]
** 2,
axis=1,
)
** (-1 / 2)
)
)
),
zeros(
len(
-arccos(
(p1 + p2 + p3)[:, 3]
* sum(
(p1 + p2 + p3)[:, 1:]
** 2,
axis=1,
)
** (-1 / 2)
)
)
),
],
[
zeros(
len(
-arccos(
(p1 + p2 + p3)[:, 3]
* sum(
(p1 + p2 + p3)[:, 1:]
** 2,
axis=1,
)
** (-1 / 2)
)
)
),
-sin(
-arccos(
(p1 + p2 + p3)[:, 3]
* sum(
(p1 + p2 + p3)[:, 1:] ** 2,
axis=1,
)
** (-1 / 2)
)
),
zeros(
len(
-arccos(
(p1 + p2 + p3)[:, 3]
* sum(
(p1 + p2 + p3)[:, 1:]
** 2,
axis=1,
)
** (-1 / 2)
)
)
),
cos(
-arccos(
(p1 + p2 + p3)[:, 3]
* sum(
(p1 + p2 + p3)[:, 1:] ** 2,
axis=1,
)
** (-1 / 2)
)
),
],
]
).transpose(2, 0, 1),
axes=(1, 2, 0),
),
transpose(
einsum(
"ij...,j...",
transpose(
array(
[
[
ones(
len(
-arctan2(
(p1 + p2 + p3)[
:, 2
],
(p1 + p2 + p3)[
:, 1
],
)
)
),
zeros(
len(
-arctan2(
(p1 + p2 + p3)[
:, 2
],
(p1 + p2 + p3)[
:, 1
],
)
)
),
zeros(
len(
-arctan2(
(p1 + p2 + p3)[
:, 2
],
(p1 + p2 + p3)[
:, 1
],
)
)
),
zeros(
len(
-arctan2(
(p1 + p2 + p3)[
:, 2
],
(p1 + p2 + p3)[
:, 1
],
)
)
),
],
[
zeros(
len(
-arctan2(
(p1 + p2 + p3)[
:, 2
],
(p1 + p2 + p3)[
:, 1
],
)
)
),
cos(
-arctan2(
(p1 + p2 + p3)[:, 2],
(p1 + p2 + p3)[:, 1],
)
),
-sin(
-arctan2(
(p1 + p2 + p3)[:, 2],
(p1 + p2 + p3)[:, 1],
)
),
zeros(
len(
-arctan2(
(p1 + p2 + p3)[
:, 2
],
(p1 + p2 + p3)[
:, 1
],
)
)
),
],
[
zeros(
len(
-arctan2(
(p1 + p2 + p3)[
:, 2
],
(p1 + p2 + p3)[
:, 1
],
)
)
),
sin(
-arctan2(
(p1 + p2 + p3)[:, 2],
(p1 + p2 + p3)[:, 1],
)
),
cos(
-arctan2(
(p1 + p2 + p3)[:, 2],
(p1 + p2 + p3)[:, 1],
)
),
zeros(
len(
-arctan2(
(p1 + p2 + p3)[
:, 2
],
(p1 + p2 + p3)[
:, 1
],
)
)
),
],
[
zeros(
len(
-arctan2(
(p1 + p2 + p3)[
:, 2
],
(p1 + p2 + p3)[
:, 1
],
)
)
),
zeros(
len(
-arctan2(
(p1 + p2 + p3)[
:, 2
],
(p1 + p2 + p3)[
:, 1
],
)
)
),
zeros(
len(
-arctan2(
(p1 + p2 + p3)[
:, 2
],
(p1 + p2 + p3)[
:, 1
],
)
)
),
ones(
len(
-arctan2(
(p1 + p2 + p3)[
:, 2
],
(p1 + p2 + p3)[
:, 1
],
)
)
),
],
]
).transpose(2, 0, 1),
axes=(1, 2, 0),
),
transpose(p2),
)
),
)
),
)
+ einsum(
"ij...,j...",
transpose(
array(
[
[
1
/ sqrt(
1
- (
sqrt(
sum(
(p1 + p2 + p3)[:, 1:] ** 2,
axis=1,
)
)
* (p1 + p2 + p3)[:, 0] ** (-1.0)
)
** 2
),
zeros(
len(
sqrt(
sum(
(p1 + p2 + p3)[:, 1:] ** 2,
axis=1,
)
)
* (p1 + p2 + p3)[:, 0] ** (-1.0)
)
),
zeros(
len(
sqrt(
sum(
(p1 + p2 + p3)[:, 1:] ** 2,
axis=1,
)
)
* (p1 + p2 + p3)[:, 0] ** (-1.0)
)
),
-1
/ sqrt(
1
- (
sqrt(
sum(
(p1 + p2 + p3)[:, 1:] ** 2,
axis=1,
)
)
* (p1 + p2 + p3)[:, 0] ** (-1.0)
)
** 2
)
* sqrt(sum((p1 + p2 + p3)[:, 1:] ** 2, axis=1))
* (p1 + p2 + p3)[:, 0] ** (-1.0),
],
[
zeros(
len(
sqrt(
sum(
(p1 + p2 + p3)[:, 1:] ** 2,
axis=1,
)
)
* (p1 + p2 + p3)[:, 0] ** (-1.0)
)
),
ones(
len(
sqrt(
sum(
(p1 + p2 + p3)[:, 1:] ** 2,
axis=1,
)
)
* (p1 + p2 + p3)[:, 0] ** (-1.0)
)
),
zeros(
len(
sqrt(
sum(
(p1 + p2 + p3)[:, 1:] ** 2,
axis=1,
)
)
* (p1 + p2 + p3)[:, 0] ** (-1.0)
)
),
zeros(
len(
sqrt(
sum(
(p1 + p2 + p3)[:, 1:] ** 2,
axis=1,
)
)
* (p1 + p2 + p3)[:, 0] ** (-1.0)
)
),
],
[
zeros(
len(
sqrt(
sum(
(p1 + p2 + p3)[:, 1:] ** 2,
axis=1,
)
)
* (p1 + p2 + p3)[:, 0] ** (-1.0)
)
),
zeros(
len(
sqrt(
sum(
(p1 + p2 + p3)[:, 1:] ** 2,
axis=1,
)
)
* (p1 + p2 + p3)[:, 0] ** (-1.0)
)
),
ones(
len(
sqrt(
sum(
(p1 + p2 + p3)[:, 1:] ** 2,
axis=1,
)
)
* (p1 + p2 + p3)[:, 0] ** (-1.0)
)
),
zeros(
len(
sqrt(
sum(
(p1 + p2 + p3)[:, 1:] ** 2,
axis=1,
)
)
* (p1 + p2 + p3)[:, 0] ** (-1.0)
)
),
],
[
-1
/ sqrt(
1
- (
sqrt(
sum(
(p1 + p2 + p3)[:, 1:] ** 2,
axis=1,
)
)
* (p1 + p2 + p3)[:, 0] ** (-1.0)
)
** 2
)
* sqrt(sum((p1 + p2 + p3)[:, 1:] ** 2, axis=1))
* (p1 + p2 + p3)[:, 0] ** (-1.0),
zeros(
len(
sqrt(
sum(
(p1 + p2 + p3)[:, 1:] ** 2,
axis=1,
)
)
* (p1 + p2 + p3)[:, 0] ** (-1.0)
)
),
zeros(
len(
sqrt(
sum(
(p1 + p2 + p3)[:, 1:] ** 2,
axis=1,
)
)
* (p1 + p2 + p3)[:, 0] ** (-1.0)
)
),
1
/ sqrt(
1
- (
sqrt(
sum(
(p1 + p2 + p3)[:, 1:] ** 2,
axis=1,
)
)
* (p1 + p2 + p3)[:, 0] ** (-1.0)
)
** 2
),
],
]
).transpose(2, 0, 1),
axes=(1, 2, 0),
),
transpose(
einsum(
"ij...,j...",
transpose(
array(
[
[
ones(
len(
-arccos(
(p1 + p2 + p3)[:, 3]
* sum(
(p1 + p2 + p3)[:, 1:]
** 2,
axis=1,
)
** (-1 / 2)
)
)
),
zeros(
len(
-arccos(
(p1 + p2 + p3)[:, 3]
* sum(
(p1 + p2 + p3)[:, 1:]
** 2,
axis=1,
)
** (-1 / 2)
)
)
),
zeros(
len(
-arccos(
(p1 + p2 + p3)[:, 3]
* sum(
(p1 + p2 + p3)[:, 1:]
** 2,
axis=1,
)
** (-1 / 2)
)
)
),
zeros(
len(
-arccos(
(p1 + p2 + p3)[:, 3]
* sum(
(p1 + p2 + p3)[:, 1:]
** 2,
axis=1,
)
** (-1 / 2)
)
)
),
],
[
zeros(
len(
-arccos(
(p1 + p2 + p3)[:, 3]
* sum(
(p1 + p2 + p3)[:, 1:]
** 2,
axis=1,
)
** (-1 / 2)
)
)
),
cos(
-arccos(
(p1 + p2 + p3)[:, 3]
* sum(
(p1 + p2 + p3)[:, 1:] ** 2,
axis=1,
)
** (-1 / 2)
)
),
zeros(
len(
-arccos(
(p1 + p2 + p3)[:, 3]
* sum(
(p1 + p2 + p3)[:, 1:]
** 2,
axis=1,
)
** (-1 / 2)
)
)
),
sin(
-arccos(
(p1 + p2 + p3)[:, 3]
* sum(
(p1 + p2 + p3)[:, 1:] ** 2,
axis=1,
)
** (-1 / 2)
)
),
],
[
zeros(
len(
-arccos(
(p1 + p2 + p3)[:, 3]
* sum(
(p1 + p2 + p3)[:, 1:]
** 2,
axis=1,
)
** (-1 / 2)
)
)
),
zeros(
len(
-arccos(
(p1 + p2 + p3)[:, 3]
* sum(
(p1 + p2 + p3)[:, 1:]
** 2,
axis=1,
)
** (-1 / 2)
)
)
),
ones(
len(
-arccos(
(p1 + p2 + p3)[:, 3]
* sum(
(p1 + p2 + p3)[:, 1:]
** 2,
axis=1,
)
** (-1 / 2)
)
)
),
zeros(
len(
-arccos(
(p1 + p2 + p3)[:, 3]
* sum(
(p1 + p2 + p3)[:, 1:]
** 2,
axis=1,
)
** (-1 / 2)
)
)
),
],
[
zeros(
len(
-arccos(
(p1 + p2 + p3)[:, 3]
* sum(
(p1 + p2 + p3)[:, 1:]
** 2,
axis=1,
)
** (-1 / 2)
)
)
),
-sin(
-arccos(
(p1 + p2 + p3)[:, 3]
* sum(
(p1 + p2 + p3)[:, 1:] ** 2,
axis=1,
)
** (-1 / 2)
)
),
zeros(
len(
-arccos(
(p1 + p2 + p3)[:, 3]
* sum(
(p1 + p2 + p3)[:, 1:]
** 2,
axis=1,
)
** (-1 / 2)
)
)
),
cos(
-arccos(
(p1 + p2 + p3)[:, 3]
* sum(
(p1 + p2 + p3)[:, 1:] ** 2,
axis=1,
)
** (-1 / 2)
)
),
],
]
).transpose(2, 0, 1),
axes=(1, 2, 0),
),
transpose(
einsum(
"ij...,j...",
transpose(
array(
[
[
ones(
len(
-arctan2(
(p1 + p2 + p3)[
:, 2
],
(p1 + p2 + p3)[
:, 1
],
)
)
),
zeros(
len(
-arctan2(
(p1 + p2 + p3)[
:, 2
],
(p1 + p2 + p3)[
:, 1
],
)
)
),
zeros(
len(
-arctan2(
(p1 + p2 + p3)[
:, 2
],
(p1 + p2 + p3)[
:, 1
],
)
)
),
zeros(
len(
-arctan2(
(p1 + p2 + p3)[
:, 2
],
(p1 + p2 + p3)[
:, 1
],
)
)
),
],
[
zeros(
len(
-arctan2(
(p1 + p2 + p3)[
:, 2
],
(p1 + p2 + p3)[
:, 1
],
)
)
),
cos(
-arctan2(
(p1 + p2 + p3)[:, 2],
(p1 + p2 + p3)[:, 1],
)
),
-sin(
-arctan2(
(p1 + p2 + p3)[:, 2],
(p1 + p2 + p3)[:, 1],
)
),
zeros(
len(
-arctan2(
(p1 + p2 + p3)[
:, 2
],
(p1 + p2 + p3)[
:, 1
],
)
)
),
],
[
zeros(
len(
-arctan2(
(p1 + p2 + p3)[
:, 2
],
(p1 + p2 + p3)[
:, 1
],
)
)
),
sin(
-arctan2(
(p1 + p2 + p3)[:, 2],
(p1 + p2 + p3)[:, 1],
)
),
cos(
-arctan2(
(p1 + p2 + p3)[:, 2],
(p1 + p2 + p3)[:, 1],
)
),
zeros(
len(
-arctan2(
(p1 + p2 + p3)[
:, 2
],
(p1 + p2 + p3)[
:, 1
],
)
)
),
],
[
zeros(
len(
-arctan2(
(p1 + p2 + p3)[
:, 2
],
(p1 + p2 + p3)[
:, 1
],
)
)
),
zeros(
len(
-arctan2(
(p1 + p2 + p3)[
:, 2
],
(p1 + p2 + p3)[
:, 1
],
)
)
),
zeros(
len(
-arctan2(
(p1 + p2 + p3)[
:, 2
],
(p1 + p2 + p3)[
:, 1
],
)
)
),
ones(
len(
-arctan2(
(p1 + p2 + p3)[
:, 2
],
(p1 + p2 + p3)[
:, 1
],
)
)
),
],
]
).transpose(2, 0, 1),
axes=(1, 2, 0),
),
transpose(p3),
)
),
)
),
)
)[:, 2],
(
einsum(
"ij...,j...",
transpose(
array(
[
[
1
/ sqrt(
1
- (
sqrt(
sum(
(p1 + p2 + p3)[:, 1:] ** 2,
axis=1,
)
)
* (p1 + p2 + p3)[:, 0] ** (-1.0)
)
** 2
),
zeros(
len(
sqrt(
sum(
(p1 + p2 + p3)[:, 1:] ** 2,
axis=1,
)
)
* (p1 + p2 + p3)[:, 0] ** (-1.0)
)
),
zeros(
len(
sqrt(
sum(
(p1 + p2 + p3)[:, 1:] ** 2,
axis=1,
)
)
* (p1 + p2 + p3)[:, 0] ** (-1.0)
)
),
-1
/ sqrt(
1
- (
sqrt(
sum(
(p1 + p2 + p3)[:, 1:] ** 2,
axis=1,
)
)
* (p1 + p2 + p3)[:, 0] ** (-1.0)
)
** 2
)
* sqrt(sum((p1 + p2 + p3)[:, 1:] ** 2, axis=1))
* (p1 + p2 + p3)[:, 0] ** (-1.0),
],
[
zeros(
len(
sqrt(
sum(
(p1 + p2 + p3)[:, 1:] ** 2,
axis=1,
)
)
* (p1 + p2 + p3)[:, 0] ** (-1.0)
)
),
ones(
len(
sqrt(
sum(
(p1 + p2 + p3)[:, 1:] ** 2,
axis=1,
)
)
* (p1 + p2 + p3)[:, 0] ** (-1.0)
)
),
zeros(
len(
sqrt(
sum(
(p1 + p2 + p3)[:, 1:] ** 2,
axis=1,
)
)
* (p1 + p2 + p3)[:, 0] ** (-1.0)
)
),
zeros(
len(
sqrt(
sum(
(p1 + p2 + p3)[:, 1:] ** 2,
axis=1,
)
)
* (p1 + p2 + p3)[:, 0] ** (-1.0)
)
),
],
[
zeros(
len(
sqrt(
sum(
(p1 + p2 + p3)[:, 1:] ** 2,
axis=1,
)
)
* (p1 + p2 + p3)[:, 0] ** (-1.0)
)
),
zeros(
len(
sqrt(
sum(
(p1 + p2 + p3)[:, 1:] ** 2,
axis=1,
)
)
* (p1 + p2 + p3)[:, 0] ** (-1.0)
)
),
ones(
len(
sqrt(
sum(
(p1 + p2 + p3)[:, 1:] ** 2,
axis=1,
)
)
* (p1 + p2 + p3)[:, 0] ** (-1.0)
)
),
zeros(
len(
sqrt(
sum(
(p1 + p2 + p3)[:, 1:] ** 2,
axis=1,
)
)
* (p1 + p2 + p3)[:, 0] ** (-1.0)
)
),
],
[
-1
/ sqrt(
1
- (
sqrt(
sum(
(p1 + p2 + p3)[:, 1:] ** 2,
axis=1,
)
)
* (p1 + p2 + p3)[:, 0] ** (-1.0)
)
** 2
)
* sqrt(sum((p1 + p2 + p3)[:, 1:] ** 2, axis=1))
* (p1 + p2 + p3)[:, 0] ** (-1.0),
zeros(
len(
sqrt(
sum(
(p1 + p2 + p3)[:, 1:] ** 2,
axis=1,
)
)
* (p1 + p2 + p3)[:, 0] ** (-1.0)
)
),
zeros(
len(
sqrt(
sum(
(p1 + p2 + p3)[:, 1:] ** 2,
axis=1,
)
)
* (p1 + p2 + p3)[:, 0] ** (-1.0)
)
),
1
/ sqrt(
1
- (
sqrt(
sum(
(p1 + p2 + p3)[:, 1:] ** 2,
axis=1,
)
)
* (p1 + p2 + p3)[:, 0] ** (-1.0)
)
** 2
),
],
]
).transpose(2, 0, 1),
axes=(1, 2, 0),
),
transpose(
einsum(
"ij...,j...",
transpose(
array(
[
[
ones(
len(
-arccos(
(p1 + p2 + p3)[:, 3]
* sum(
(p1 + p2 + p3)[:, 1:]
** 2,
axis=1,
)
** (-1 / 2)
)
)
),
zeros(
len(
-arccos(
(p1 + p2 + p3)[:, 3]
* sum(
(p1 + p2 + p3)[:, 1:]
** 2,
axis=1,
)
** (-1 / 2)
)
)
),
zeros(
len(
-arccos(
(p1 + p2 + p3)[:, 3]
* sum(
(p1 + p2 + p3)[:, 1:]
** 2,
axis=1,
)
** (-1 / 2)
)
)
),
zeros(
len(
-arccos(
(p1 + p2 + p3)[:, 3]
* sum(
(p1 + p2 + p3)[:, 1:]
** 2,
axis=1,
)
** (-1 / 2)
)
)
),
],
[
zeros(
len(
-arccos(
(p1 + p2 + p3)[:, 3]
* sum(
(p1 + p2 + p3)[:, 1:]
** 2,
axis=1,
)
** (-1 / 2)
)
)
),
cos(
-arccos(
(p1 + p2 + p3)[:, 3]
* sum(
(p1 + p2 + p3)[:, 1:] ** 2,
axis=1,
)
** (-1 / 2)
)
),
zeros(
len(
-arccos(
(p1 + p2 + p3)[:, 3]
* sum(
(p1 + p2 + p3)[:, 1:]
** 2,
axis=1,
)
** (-1 / 2)
)
)
),
sin(
-arccos(
(p1 + p2 + p3)[:, 3]
* sum(
(p1 + p2 + p3)[:, 1:] ** 2,
axis=1,
)
** (-1 / 2)
)
),
],
[
zeros(
len(
-arccos(
(p1 + p2 + p3)[:, 3]
* sum(
(p1 + p2 + p3)[:, 1:]
** 2,
axis=1,
)
** (-1 / 2)
)
)
),
zeros(
len(
-arccos(
(p1 + p2 + p3)[:, 3]
* sum(
(p1 + p2 + p3)[:, 1:]
** 2,
axis=1,
)
** (-1 / 2)
)
)
),
ones(
len(
-arccos(
(p1 + p2 + p3)[:, 3]
* sum(
(p1 + p2 + p3)[:, 1:]
** 2,
axis=1,
)
** (-1 / 2)
)
)
),
zeros(
len(
-arccos(
(p1 + p2 + p3)[:, 3]
* sum(
(p1 + p2 + p3)[:, 1:]
** 2,
axis=1,
)
** (-1 / 2)
)
)
),
],
[
zeros(
len(
-arccos(
(p1 + p2 + p3)[:, 3]
* sum(
(p1 + p2 + p3)[:, 1:]
** 2,
axis=1,
)
** (-1 / 2)
)
)
),
-sin(
-arccos(
(p1 + p2 + p3)[:, 3]
* sum(
(p1 + p2 + p3)[:, 1:] ** 2,
axis=1,
)
** (-1 / 2)
)
),
zeros(
len(
-arccos(
(p1 + p2 + p3)[:, 3]
* sum(
(p1 + p2 + p3)[:, 1:]
** 2,
axis=1,
)
** (-1 / 2)
)
)
),
cos(
-arccos(
(p1 + p2 + p3)[:, 3]
* sum(
(p1 + p2 + p3)[:, 1:] ** 2,
axis=1,
)
** (-1 / 2)
)
),
],
]
).transpose(2, 0, 1),
axes=(1, 2, 0),
),
transpose(
einsum(
"ij...,j...",
transpose(
array(
[
[
ones(
len(
-arctan2(
(p1 + p2 + p3)[
:, 2
],
(p1 + p2 + p3)[
:, 1
],
)
)
),
zeros(
len(
-arctan2(
(p1 + p2 + p3)[
:, 2
],
(p1 + p2 + p3)[
:, 1
],
)
)
),
zeros(
len(
-arctan2(
(p1 + p2 + p3)[
:, 2
],
(p1 + p2 + p3)[
:, 1
],
)
)
),
zeros(
len(
-arctan2(
(p1 + p2 + p3)[
:, 2
],
(p1 + p2 + p3)[
:, 1
],
)
)
),
],
[
zeros(
len(
-arctan2(
(p1 + p2 + p3)[
:, 2
],
(p1 + p2 + p3)[
:, 1
],
)
)
),
cos(
-arctan2(
(p1 + p2 + p3)[:, 2],
(p1 + p2 + p3)[:, 1],
)
),
-sin(
-arctan2(
(p1 + p2 + p3)[:, 2],
(p1 + p2 + p3)[:, 1],
)
),
zeros(
len(
-arctan2(
(p1 + p2 + p3)[
:, 2
],
(p1 + p2 + p3)[
:, 1
],
)
)
),
],
[
zeros(
len(
-arctan2(
(p1 + p2 + p3)[
:, 2
],
(p1 + p2 + p3)[
:, 1
],
)
)
),
sin(
-arctan2(
(p1 + p2 + p3)[:, 2],
(p1 + p2 + p3)[:, 1],
)
),
cos(
-arctan2(
(p1 + p2 + p3)[:, 2],
(p1 + p2 + p3)[:, 1],
)
),
zeros(
len(
-arctan2(
(p1 + p2 + p3)[
:, 2
],
(p1 + p2 + p3)[
:, 1
],
)
)
),
],
[
zeros(
len(
-arctan2(
(p1 + p2 + p3)[
:, 2
],
(p1 + p2 + p3)[
:, 1
],
)
)
),
zeros(
len(
-arctan2(
(p1 + p2 + p3)[
:, 2
],
(p1 + p2 + p3)[
:, 1
],
)
)
),
zeros(
len(
-arctan2(
(p1 + p2 + p3)[
:, 2
],
(p1 + p2 + p3)[
:, 1
],
)
)
),
ones(
len(
-arctan2(
(p1 + p2 + p3)[
:, 2
],
(p1 + p2 + p3)[
:, 1
],
)
)
),
],
]
).transpose(2, 0, 1),
axes=(1, 2, 0),
),
transpose(p2),
)
),
)
),
)
+ einsum(
"ij...,j...",
transpose(
array(
[
[
1
/ sqrt(
1
- (
sqrt(
sum(
(p1 + p2 + p3)[:, 1:] ** 2,
axis=1,
)
)
* (p1 + p2 + p3)[:, 0] ** (-1.0)
)
** 2
),
zeros(
len(
sqrt(
sum(
(p1 + p2 + p3)[:, 1:] ** 2,
axis=1,
)
)
* (p1 + p2 + p3)[:, 0] ** (-1.0)
)
),
zeros(
len(
sqrt(
sum(
(p1 + p2 + p3)[:, 1:] ** 2,
axis=1,
)
)
* (p1 + p2 + p3)[:, 0] ** (-1.0)
)
),
-1
/ sqrt(
1
- (
sqrt(
sum(
(p1 + p2 + p3)[:, 1:] ** 2,
axis=1,
)
)
* (p1 + p2 + p3)[:, 0] ** (-1.0)
)
** 2
)
* sqrt(sum((p1 + p2 + p3)[:, 1:] ** 2, axis=1))
* (p1 + p2 + p3)[:, 0] ** (-1.0),
],
[
zeros(
len(
sqrt(
sum(
(p1 + p2 + p3)[:, 1:] ** 2,
axis=1,
)
)
* (p1 + p2 + p3)[:, 0] ** (-1.0)
)
),
ones(
len(
sqrt(
sum(
(p1 + p2 + p3)[:, 1:] ** 2,
axis=1,
)
)
* (p1 + p2 + p3)[:, 0] ** (-1.0)
)
),
zeros(
len(
sqrt(
sum(
(p1 + p2 + p3)[:, 1:] ** 2,
axis=1,
)
)
* (p1 + p2 + p3)[:, 0] ** (-1.0)
)
),
zeros(
len(
sqrt(
sum(
(p1 + p2 + p3)[:, 1:] ** 2,
axis=1,
)
)
* (p1 + p2 + p3)[:, 0] ** (-1.0)
)
),
],
[
zeros(
len(
sqrt(
sum(
(p1 + p2 + p3)[:, 1:] ** 2,
axis=1,
)
)
* (p1 + p2 + p3)[:, 0] ** (-1.0)
)
),
zeros(
len(
sqrt(
sum(
(p1 + p2 + p3)[:, 1:] ** 2,
axis=1,
)
)
* (p1 + p2 + p3)[:, 0] ** (-1.0)
)
),
ones(
len(
sqrt(
sum(
(p1 + p2 + p3)[:, 1:] ** 2,
axis=1,
)
)
* (p1 + p2 + p3)[:, 0] ** (-1.0)
)
),
zeros(
len(
sqrt(
sum(
(p1 + p2 + p3)[:, 1:] ** 2,
axis=1,
)
)
* (p1 + p2 + p3)[:, 0] ** (-1.0)
)
),
],
[
-1
/ sqrt(
1
- (
sqrt(
sum(
(p1 + p2 + p3)[:, 1:] ** 2,
axis=1,
)
)
* (p1 + p2 + p3)[:, 0] ** (-1.0)
)
** 2
)
* sqrt(sum((p1 + p2 + p3)[:, 1:] ** 2, axis=1))
* (p1 + p2 + p3)[:, 0] ** (-1.0),
zeros(
len(
sqrt(
sum(
(p1 + p2 + p3)[:, 1:] ** 2,
axis=1,
)
)
* (p1 + p2 + p3)[:, 0] ** (-1.0)
)
),
zeros(
len(
sqrt(
sum(
(p1 + p2 + p3)[:, 1:] ** 2,
axis=1,
)
)
* (p1 + p2 + p3)[:, 0] ** (-1.0)
)
),
1
/ sqrt(
1
- (
sqrt(
sum(
(p1 + p2 + p3)[:, 1:] ** 2,
axis=1,
)
)
* (p1 + p2 + p3)[:, 0] ** (-1.0)
)
** 2
),
],
]
).transpose(2, 0, 1),
axes=(1, 2, 0),
),
transpose(
einsum(
"ij...,j...",
transpose(
array(
[
[
ones(
len(
-arccos(
(p1 + p2 + p3)[:, 3]
* sum(
(p1 + p2 + p3)[:, 1:]
** 2,
axis=1,
)
** (-1 / 2)
)
)
),
zeros(
len(
-arccos(
(p1 + p2 + p3)[:, 3]
* sum(
(p1 + p2 + p3)[:, 1:]
** 2,
axis=1,
)
** (-1 / 2)
)
)
),
zeros(
len(
-arccos(
(p1 + p2 + p3)[:, 3]
* sum(
(p1 + p2 + p3)[:, 1:]
** 2,
axis=1,
)
** (-1 / 2)
)
)
),
zeros(
len(
-arccos(
(p1 + p2 + p3)[:, 3]
* sum(
(p1 + p2 + p3)[:, 1:]
** 2,
axis=1,
)
** (-1 / 2)
)
)
),
],
[
zeros(
len(
-arccos(
(p1 + p2 + p3)[:, 3]
* sum(
(p1 + p2 + p3)[:, 1:]
** 2,
axis=1,
)
** (-1 / 2)
)
)
),
cos(
-arccos(
(p1 + p2 + p3)[:, 3]
* sum(
(p1 + p2 + p3)[:, 1:] ** 2,
axis=1,
)
** (-1 / 2)
)
),
zeros(
len(
-arccos(
(p1 + p2 + p3)[:, 3]
* sum(
(p1 + p2 + p3)[:, 1:]
** 2,
axis=1,
)
** (-1 / 2)
)
)
),
sin(
-arccos(
(p1 + p2 + p3)[:, 3]
* sum(
(p1 + p2 + p3)[:, 1:] ** 2,
axis=1,
)
** (-1 / 2)
)
),
],
[
zeros(
len(
-arccos(
(p1 + p2 + p3)[:, 3]
* sum(
(p1 + p2 + p3)[:, 1:]
** 2,
axis=1,
)
** (-1 / 2)
)
)
),
zeros(
len(
-arccos(
(p1 + p2 + p3)[:, 3]
* sum(
(p1 + p2 + p3)[:, 1:]
** 2,
axis=1,
)
** (-1 / 2)
)
)
),
ones(
len(
-arccos(
(p1 + p2 + p3)[:, 3]
* sum(
(p1 + p2 + p3)[:, 1:]
** 2,
axis=1,
)
** (-1 / 2)
)
)
),
zeros(
len(
-arccos(
(p1 + p2 + p3)[:, 3]
* sum(
(p1 + p2 + p3)[:, 1:]
** 2,
axis=1,
)
** (-1 / 2)
)
)
),
],
[
zeros(
len(
-arccos(
(p1 + p2 + p3)[:, 3]
* sum(
(p1 + p2 + p3)[:, 1:]
** 2,
axis=1,
)
** (-1 / 2)
)
)
),
-sin(
-arccos(
(p1 + p2 + p3)[:, 3]
* sum(
(p1 + p2 + p3)[:, 1:] ** 2,
axis=1,
)
** (-1 / 2)
)
),
zeros(
len(
-arccos(
(p1 + p2 + p3)[:, 3]
* sum(
(p1 + p2 + p3)[:, 1:]
** 2,
axis=1,
)
** (-1 / 2)
)
)
),
cos(
-arccos(
(p1 + p2 + p3)[:, 3]
* sum(
(p1 + p2 + p3)[:, 1:] ** 2,
axis=1,
)
** (-1 / 2)
)
),
],
]
).transpose(2, 0, 1),
axes=(1, 2, 0),
),
transpose(
einsum(
"ij...,j...",
transpose(
array(
[
[
ones(
len(
-arctan2(
(p1 + p2 + p3)[
:, 2
],
(p1 + p2 + p3)[
:, 1
],
)
)
),
zeros(
len(
-arctan2(
(p1 + p2 + p3)[
:, 2
],
(p1 + p2 + p3)[
:, 1
],
)
)
),
zeros(
len(
-arctan2(
(p1 + p2 + p3)[
:, 2
],
(p1 + p2 + p3)[
:, 1
],
)
)
),
zeros(
len(
-arctan2(
(p1 + p2 + p3)[
:, 2
],
(p1 + p2 + p3)[
:, 1
],
)
)
),
],
[
zeros(
len(
-arctan2(
(p1 + p2 + p3)[
:, 2
],
(p1 + p2 + p3)[
:, 1
],
)
)
),
cos(
-arctan2(
(p1 + p2 + p3)[:, 2],
(p1 + p2 + p3)[:, 1],
)
),
-sin(
-arctan2(
(p1 + p2 + p3)[:, 2],
(p1 + p2 + p3)[:, 1],
)
),
zeros(
len(
-arctan2(
(p1 + p2 + p3)[
:, 2
],
(p1 + p2 + p3)[
:, 1
],
)
)
),
],
[
zeros(
len(
-arctan2(
(p1 + p2 + p3)[
:, 2
],
(p1 + p2 + p3)[
:, 1
],
)
)
),
sin(
-arctan2(
(p1 + p2 + p3)[:, 2],
(p1 + p2 + p3)[:, 1],
)
),
cos(
-arctan2(
(p1 + p2 + p3)[:, 2],
(p1 + p2 + p3)[:, 1],
)
),
zeros(
len(
-arctan2(
(p1 + p2 + p3)[
:, 2
],
(p1 + p2 + p3)[
:, 1
],
)
)
),
],
[
zeros(
len(
-arctan2(
(p1 + p2 + p3)[
:, 2
],
(p1 + p2 + p3)[
:, 1
],
)
)
),
zeros(
len(
-arctan2(
(p1 + p2 + p3)[
:, 2
],
(p1 + p2 + p3)[
:, 1
],
)
)
),
zeros(
len(
-arctan2(
(p1 + p2 + p3)[
:, 2
],
(p1 + p2 + p3)[
:, 1
],
)
)
),
ones(
len(
-arctan2(
(p1 + p2 + p3)[
:, 2
],
(p1 + p2 + p3)[
:, 1
],
)
)
),
],
]
).transpose(2, 0, 1),
axes=(1, 2, 0),
),
transpose(p3),
)
),
)
),
)
)[:, 1],
)
Comparison computed value#
Finally, what’s left is to compare computed data with the Test sample. Here’s a little comparison function that can be called for each angle:
for angle_name in symbolic_angles:
print(angle_name)
%time compare(angle_name)
print()
phi_1+2+3
array([ 2.79758029, 2.51292308, -1.07396684])
array([ 2.79758029, 2.51292308, -1.07396684])
CPU times: user 2.48 ms, sys: 6.56 ms, total: 9.04 ms
Wall time: 7.77 ms
theta_1+2+3
array([2.72456853, 3.03316287, 0.69240082])
array([2.72456853, 3.03316287, 0.69240082])
CPU times: user 5.13 ms, sys: 297 µs, total: 5.43 ms
Wall time: 4.93 ms
phi_2+3,1+2+3
array([1.0436215 , 1.8734936 , 0.16073833])
array([1.0436215 , 1.8734936 , 0.16073833])
CPU times: user 34.1 ms, sys: 0 ns, total: 34.1 ms
Wall time: 33.1 ms
theta_2+3,1+2+3
array([2.45361589, 1.40639741, 0.98079245])
array([2.45361589, 1.40639741, 0.98079245])
CPU times: user 36.1 ms, sys: 451 µs, total: 36.6 ms
Wall time: 35.5 ms
phi_2,2+3,1+2+3
array([ 0.36955786, -1.68820498, 0.63063002])
array([ 0.36955786, -1.68820498, 0.63063002])
CPU times: user 1.25 s, sys: 162 ms, total: 1.42 s
Wall time: 1.41 s
theta_2,2+3,1+2+3
array([1.0924374 , 1.99375767, 1.31959621])
array([1.0924374 , 1.99375767, 1.31959621])
CPU times: user 1.22 s, sys: 166 ms, total: 1.38 s
Wall time: 1.38 s