docs/report/011.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```{autolink-concat}\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": [
     "physics",
     "sympy"
    ]
   },
   "source": [
    "::::{margin}\n",
    ":::{card} Helicity angles as symbolic expressions\n",
    "TR-011\n",
    "^^^\n",
    "This report has been implemented in and [ampform#177](https://github.com/ComPWA/ampform/issues/177) and  [tensorwaves#345](https://github.com/ComPWA/tensorwaves/issues/345). The report contains some bugs which were also addressed in these PRs.\n",
    "+++\n",
    "✅&nbsp;[ampform#177](https://github.com/ComPWA/ampform/issues/177), [tensorwaves#345](https://github.com/ComPWA/tensorwaves/issues/345)\n",
    ":::\n",
    "::::"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Symbolic kinematics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "%pip install -q ampform==0.11.4 black==21.7b0 graphviz==0.17 numpy==1.19.5 qrules==0.9.2 git+https://github.com/ComPWA/sympy@20570-add-ArraySlice"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [],
   "source": [
    "from __future__ import annotations\n",
    "\n",
    "import inspect\n",
    "from typing import TYPE_CHECKING\n",
    "\n",
    "import black\n",
    "import graphviz\n",
    "import numpy as np\n",
    "import qrules\n",
    "import sympy as sp\n",
    "from ampform.data import EventCollection\n",
    "from ampform.kinematics import (\n",
    "    _compute_helicity_angles,\n",
    "    determine_attached_final_state,\n",
    "    get_helicity_angle_label,\n",
    ")\n",
    "from ampform.sympy import (\n",
    "    UnevaluatedExpression,\n",
    "    create_expression,\n",
    "    implement_doit_method,\n",
    ")\n",
    "from IPython.display import Math, display\n",
    "from qrules.topology import Topology, create_isobar_topologies\n",
    "from sympy.printing.numpy import NumPyPrinter\n",
    "from sympy.tensor.array.expressions.array_expressions import ArraySlice, ArraySymbol\n",
    "\n",
    "if TYPE_CHECKING:\n",
    "    from sympy.printing.printer import Printer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This report investigates issue [compwa.github.io#56](https://github.com/ComPWA/compwa.github.io/issues/56). The ideal solution would be to use only SymPy in the existing [`ampform.kinematics`](https://ampform.readthedocs.io/en/0.11.4/api/ampform.kinematics.html) module. This has two benefits:\n",
    "\n",
    "1. It allows computing kinematic variables from four-momenta with different computational back-ends.\n",
    "2. Expressions for kinematic variable can be inspected through their LaTeX representation.\n",
    "\n",
    "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)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Test sample"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Data sample taken from [this test in AmpForm](https://github.com/ComPWA/ampform/tree/0.11.4/tests/conftest.py#L53-L115) and topology and expected angles taken from [here](https://github.com/ComPWA/ampform/tree/0.11.4/tests/test_kinematics.py#L15-L112)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "topologies = create_isobar_topologies(4)\n",
    "topology = topologies[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "dot = qrules.io.asdot(topology)\n",
    "graphviz.Source(dot)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "graphviz.Source(dot).render(\"011-topology\", format=\"svg\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![](https://user-images.githubusercontent.com/29308176/164990755-ba9dfa9b-ed27-4bcf-b9cf-8fd62a71e3f4.svg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [],
   "source": [
    "events = EventCollection({\n",
    "    0: np.array([  # pi0\n",
    "        (1.35527, 0.514208, -0.184219, 1.23296),\n",
    "        (0.841933, 0.0727385, -0.0528868, 0.826163),\n",
    "        (0.550927, -0.162529, 0.29976, -0.411133),\n",
    "    ]),\n",
    "    1: np.array([  # gamma\n",
    "        (0.755744, -0.305812, 0.284, -0.630057),\n",
    "        (1.02861, 0.784483, 0.614347, -0.255334),\n",
    "        (0.356875, -0.20767, 0.272796, 0.0990739),\n",
    "    ]),\n",
    "    2: np.array([  # pi0\n",
    "        (0.208274, -0.061663, -0.0211864, 0.144596),\n",
    "        (0.461193, -0.243319, -0.283044, -0.234866),\n",
    "        (1.03294, 0.82872, -0.0465425, -0.599834),\n",
    "    ]),\n",
    "    3: np.array([  # pi0\n",
    "        (0.777613, -0.146733, -0.0785946, -0.747499),\n",
    "        (0.765168, -0.613903, -0.278416, -0.335962),\n",
    "        (1.15616, -0.458522, -0.526014, 0.911894),\n",
    "    ]),\n",
    "})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "keep_output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'phi_1+2+3': ScalarSequence([ 2.79758029  2.51292308 -1.07396684],\n",
       " 'theta_1+2+3': ScalarSequence([2.72456853 3.03316287 0.69240082],\n",
       " 'phi_2+3,1+2+3': ScalarSequence([1.0436215  1.8734936  0.16073833],\n",
       " 'theta_2+3,1+2+3': ScalarSequence([2.45361589 1.40639741 0.98079245],\n",
       " 'phi_2,2+3,1+2+3': ScalarSequence([ 0.36955786 -1.68820498  0.63063002],\n",
       " 'theta_2,2+3,1+2+3': ScalarSequence([1.0924374  1.99375767 1.31959621]}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "angles = _compute_helicity_angles(events, topology)\n",
    "angles._DataSet__data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Einstein summation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First challenge is to express the Einstein summation in [the existing implementation](https://github.com/ComPWA/ampform/tree/0.11.4/src/ampform/data.py#L207-L214) 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`](https://ampform.readthedocs.io/en/0.11.4/api/ampform.sympy.html#ampform.sympy.UnevaluatedExpression) and using the decorator functions provided by the [`ampform.sympy`](https://ampform.readthedocs.io/en/0.11.4/api/ampform.sympy.html) module."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define boost and rotation classes"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, wrap rotations and boosts in a class so with a nice LaTeX printer. {ref}`Later on <report/011:Define lambdification>`, a NumPy printer method will be defined externally for each of them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "class BoostZ(UnevaluatedExpression):\n",
    "    def __new__(cls, beta: sp.Symbol, **hints) -> BoostZ:\n",
    "        return create_expression(cls, beta, **hints)\n",
    "\n",
    "    def as_explicit(self) -> sp.Expr:\n",
    "        gamma = 1 / sp.sqrt(1 - beta**2)\n",
    "        return sp.Matrix([\n",
    "            [gamma, 0, 0, -gamma * beta],\n",
    "            [0, 1, 0, 0],\n",
    "            [0, 0, 1, 0],\n",
    "            [-gamma * beta, 0, 0, gamma],\n",
    "        ])\n",
    "\n",
    "    def _latex(self, printer, *args) -> str:\n",
    "        beta, *_ = self.args\n",
    "        beta = printer._print(beta)\n",
    "        return Rf\"\\boldsymbol{{B_z}}\\left({beta}\\right)\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "class RotationY(UnevaluatedExpression):\n",
    "    def __new__(cls, angle: sp.Symbol, **hints) -> RotationY:\n",
    "        return create_expression(cls, angle, **hints)\n",
    "\n",
    "    def as_explicit(self) -> sp.Expr:\n",
    "        angle = self.args[0]\n",
    "        return sp.Matrix([\n",
    "            [1, 0, 0, 0],\n",
    "            [0, sp.cos(angle), 0, sp.sin(angle)],\n",
    "            [0, 0, 1, 0],\n",
    "            [0, -sp.sin(angle), 0, sp.cos(angle)],\n",
    "        ])\n",
    "\n",
    "    def _latex(self, printer, *args) -> str:\n",
    "        angle, *_ = self.args\n",
    "        angle = printer._print(angle)\n",
    "        return Rf\"\\boldsymbol{{R_y}}\\left({angle}\\right)\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [],
   "source": [
    "class RotationZ(UnevaluatedExpression):\n",
    "    def __new__(cls, angle: sp.Symbol, **hints) -> RotationZ:\n",
    "        return create_expression(cls, angle, **hints)\n",
    "\n",
    "    def as_explicit(self) -> sp.Expr:\n",
    "        angle = self.args[0]\n",
    "        return sp.Matrix([\n",
    "            [1, 0, 0, 0],\n",
    "            [0, sp.cos(angle), -sp.sin(angle), 0],\n",
    "            [0, sp.sin(angle), sp.cos(angle), 0],\n",
    "            [0, 0, 0, 1],\n",
    "        ])\n",
    "\n",
    "    def _latex(self, printer, *args) -> str:\n",
    "        angle, *_ = self.args\n",
    "        angle = printer._print(angle)\n",
    "        return Rf\"\\boldsymbol{{R_z}}\\left({angle}\\right)\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define Einstein summation class"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Similarly, we define a `ArrayMultiplication` class that will eventually be lambdified to {func}`numpy.einsum`.\n",
    "\n",
    "(011-array-multiplication)="
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "class ArrayMultiplication(sp.Expr):\n",
    "    def __new__(cls, *tensors: sp.Symbol, **hints):\n",
    "        return create_expression(cls, *tensors, **hints)\n",
    "\n",
    "    def _latex(self, printer, *args) -> str:\n",
    "        tensors_latex = map(printer._print, self.args)\n",
    "        return \" \".join(tensors_latex)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Indeed an expression involving these classes looks nice on the top-level:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "keep_output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\boldsymbol{B_z}\\left(\\beta\\right) \\boldsymbol{R_y}\\left(- \\theta\\right) \\boldsymbol{R_z}\\left(- \\phi\\right) p$"
      ],
      "text/plain": [
       "ArrayMultiplication(BoostZ(beta), RotationY(-theta), RotationZ(-phi), p)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "n_events = 3\n",
    "momentum = sp.MatrixSymbol(\"p\", m=n_events, n=4)\n",
    "beta = sp.Symbol(\"beta\")\n",
    "phi = sp.Symbol(\"phi\")\n",
    "theta = sp.Symbol(\"theta\")\n",
    "\n",
    "boosted_momentum = ArrayMultiplication(\n",
    "    BoostZ(beta),\n",
    "    RotationY(-theta),\n",
    "    RotationZ(-phi),\n",
    "    momentum,\n",
    ")\n",
    "boosted_momentum"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ":::{note}\n",
    "\n",
    "It could be that the above can be achieved with SymPy's `ArrayTensorProduct` and `ArrayContraction`. See for instance [sympy/sympy#22279](https://github.com/sympy/sympy/issues/22279).\n",
    "\n",
    ":::"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define lambdification"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [],
   "source": [
    "# small helper function\n",
    "\n",
    "\n",
    "def print_lambdify(symbols, expr):\n",
    "    np_expr = sp.lambdify(symbols, expr)\n",
    "    src = inspect.getsource(np_expr)\n",
    "    src = black.format_str(src, mode=black.Mode(line_length=79))\n",
    "    print(src)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we have a problem: lambdification does not work..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "keep_output"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "def _lambdifygenerated(beta, theta, p):\n",
      "    return (  # Not supported in Python with SciPy and NumPy:\n",
      "        # ArrayMultiplication\n",
      "        ArrayMultiplication(\n",
      "            BoostZ(beta), RotationY(-theta), RotationZ(-phi), p\n",
      "        )\n",
      "    )\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print_lambdify([beta, theta, momentum], boosted_momentum)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "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):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def print_as_numpy(self, printer: Printer, *args) -> str:\n",
    "    def multiply(matrix, vector):\n",
    "        return (\n",
    "            'einsum(\"ij...,j...\",'\n",
    "            f\" transpose({matrix}, axes=(1, 2, 0)),\"\n",
    "            f\" transpose({vector}))\"\n",
    "        )\n",
    "\n",
    "    def recursive_multiply(tensors):\n",
    "        if len(tensors) < 2:\n",
    "            msg = \"Need at least two tensors\"\n",
    "            raise ValueError(msg)\n",
    "        if len(tensors) == 2:\n",
    "            return multiply(tensors[0], tensors[1])\n",
    "        return multiply(tensors[0], recursive_multiply(tensors[1:]))\n",
    "\n",
    "    printer.module_imports[\"numpy\"].update({\"einsum\", \"transpose\"})\n",
    "    tensors = list(map(printer._print, self.args))\n",
    "    if len(tensors) == 0:\n",
    "        return \"\"\n",
    "    if len(tensors) == 1:\n",
    "        return tensors[0]\n",
    "    return recursive_multiply(tensors)\n",
    "\n",
    "\n",
    "ArrayMultiplication._numpycode = print_as_numpy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "keep_output"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "def _lambdifygenerated(beta, theta, p):\n",
      "    return einsum(\n",
      "        \"ij...,j...\",\n",
      "        transpose(beta, axes=(1, 2, 0)),\n",
      "        transpose(\n",
      "            einsum(\n",
      "                \"ij...,j...\", transpose(theta, axes=(1, 2, 0)), transpose(p)\n",
      "            )\n",
      "        ),\n",
      "    )\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print_lambdify(\n",
    "    symbols=[beta, theta, momentum],\n",
    "    expr=ArrayMultiplication(beta, theta, momentum),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This also needs to be done for the rotation and boost classes:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "keep_output"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "def _lambdifygenerated(beta, theta, p):\n",
      "    return (  # Not supported in Python with SciPy and NumPy:\n",
      "        # BoostZ\n",
      "        # RotationY\n",
      "        # RotationZ\n",
      "        einsum(\n",
      "            \"ij...,j...\",\n",
      "            transpose(BoostZ(beta), axes=(1, 2, 0)),\n",
      "            transpose(\n",
      "                einsum(\n",
      "                    \"ij...,j...\",\n",
      "                    transpose(RotationY(-theta), axes=(1, 2, 0)),\n",
      "                    transpose(\n",
      "                        einsum(\n",
      "                            \"ij...,j...\",\n",
      "                            transpose(RotationZ(-phi), axes=(1, 2, 0)),\n",
      "                            transpose(p),\n",
      "                        )\n",
      "                    ),\n",
      "                )\n",
      "            ),\n",
      "        )\n",
      "    )\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print_lambdify([beta, theta, momentum], boosted_momentum)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This time, we define the lambdification **through the printer class**:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def _print_BoostZ(self: NumPyPrinter, expr: BoostZ) -> str:\n",
    "    self.module_imports[\"numpy\"].update({\"array\", \"ones\", \"zeros\", \"sqrt\"})\n",
    "    arg = expr.args[0]\n",
    "    beta = self._print(arg)\n",
    "    gamma = f\"1 / sqrt(1 - ({beta}) ** 2)\"\n",
    "    n_events = f\"len({beta})\"\n",
    "    zeros = f\"zeros({n_events})\"\n",
    "    ones = f\"ones({n_events})\"\n",
    "    return f\"\"\"array(\n",
    "        [\n",
    "            [{gamma}, {zeros}, {zeros}, -{gamma} * {beta}],\n",
    "            [{zeros}, {ones}, {zeros}, {zeros}],\n",
    "            [{zeros}, {zeros}, {ones}, {zeros}],\n",
    "            [-{gamma} * {beta}, {zeros}, {zeros}, {gamma}],\n",
    "        ]\n",
    "    ).transpose(2, 0, 1)\"\"\"\n",
    "\n",
    "\n",
    "NumPyPrinter._print_BoostZ = _print_BoostZ"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def _print_RotationY(self: NumPyPrinter, expr: RotationY) -> str:\n",
    "    self.module_imports[\"numpy\"].update({\"array\", \"cos\", \"ones\", \"zeros\", \"sin\"})\n",
    "    arg = expr.args[0]\n",
    "    angle = self._print(arg)\n",
    "    n_events = f\"len({angle})\"\n",
    "    zeros = f\"zeros({n_events})\"\n",
    "    ones = f\"ones({n_events})\"\n",
    "    return f\"\"\"array(\n",
    "        [\n",
    "            [{ones}, {zeros}, {zeros}, {zeros}],\n",
    "            [{zeros}, cos({angle}), {zeros}, sin({angle})],\n",
    "            [{zeros}, {zeros}, {ones}, {zeros}],\n",
    "            [{zeros}, -sin({angle}), {zeros}, cos({angle})],\n",
    "        ]\n",
    "    ).transpose(2, 0, 1)\"\"\"\n",
    "\n",
    "\n",
    "NumPyPrinter._print_RotationY = _print_RotationY"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [],
   "source": [
    "def _print_RotationZ(self: NumPyPrinter, expr: RotationZ) -> str:\n",
    "    self.module_imports[\"numpy\"].update({\"array\", \"cos\", \"ones\", \"zeros\", \"sin\"})\n",
    "    arg = expr.args[0]\n",
    "    angle = self._print(arg)\n",
    "    n_events = f\"len({angle})\"\n",
    "    zeros = f\"zeros({n_events})\"\n",
    "    ones = f\"ones({n_events})\"\n",
    "    return f\"\"\"array(\n",
    "        [\n",
    "            [{ones}, {zeros}, {zeros}, {zeros}],\n",
    "            [{zeros}, cos({angle}), -sin({angle}), {zeros}],\n",
    "            [{zeros}, sin({angle}), cos({angle}), {zeros}],\n",
    "            [{zeros}, {zeros}, {zeros}, {ones}],\n",
    "        ]\n",
    "    ).transpose(2, 0, 1)\"\"\"\n",
    "\n",
    "\n",
    "NumPyPrinter._print_RotationZ = _print_RotationZ"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "scroll-output",
     "keep_output"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "def _lambdifygenerated(beta, theta, p):\n",
      "    return einsum(\n",
      "        \"ij...,j...\",\n",
      "        transpose(\n",
      "            array(\n",
      "                [\n",
      "                    [\n",
      "                        1 / sqrt(1 - (beta) ** 2),\n",
      "                        zeros(len(beta)),\n",
      "                        zeros(len(beta)),\n",
      "                        -1 / sqrt(1 - (beta) ** 2) * beta,\n",
      "                    ],\n",
      "                    [\n",
      "                        zeros(len(beta)),\n",
      "                        ones(len(beta)),\n",
      "                        zeros(len(beta)),\n",
      "                        zeros(len(beta)),\n",
      "                    ],\n",
      "                    [\n",
      "                        zeros(len(beta)),\n",
      "                        zeros(len(beta)),\n",
      "                        ones(len(beta)),\n",
      "                        zeros(len(beta)),\n",
      "                    ],\n",
      "                    [\n",
      "                        -1 / sqrt(1 - (beta) ** 2) * beta,\n",
      "                        zeros(len(beta)),\n",
      "                        zeros(len(beta)),\n",
      "                        1 / sqrt(1 - (beta) ** 2),\n",
      "                    ],\n",
      "                ]\n",
      "            ).transpose(2, 0, 1),\n",
      "            axes=(1, 2, 0),\n",
      "        ),\n",
      "        transpose(\n",
      "            einsum(\n",
      "                \"ij...,j...\",\n",
      "                transpose(\n",
      "                    array(\n",
      "                        [\n",
      "                            [\n",
      "                                ones(len(-theta)),\n",
      "                                zeros(len(-theta)),\n",
      "                                zeros(len(-theta)),\n",
      "                                zeros(len(-theta)),\n",
      "                            ],\n",
      "                            [\n",
      "                                zeros(len(-theta)),\n",
      "                                cos(-theta),\n",
      "                                zeros(len(-theta)),\n",
      "                                sin(-theta),\n",
      "                            ],\n",
      "                            [\n",
      "                                zeros(len(-theta)),\n",
      "                                zeros(len(-theta)),\n",
      "                                ones(len(-theta)),\n",
      "                                zeros(len(-theta)),\n",
      "                            ],\n",
      "                            [\n",
      "                                zeros(len(-theta)),\n",
      "                                -sin(-theta),\n",
      "                                zeros(len(-theta)),\n",
      "                                cos(-theta),\n",
      "                            ],\n",
      "                        ]\n",
      "                    ).transpose(2, 0, 1),\n",
      "                    axes=(1, 2, 0),\n",
      "                ),\n",
      "                transpose(\n",
      "                    einsum(\n",
      "                        \"ij...,j...\",\n",
      "                        transpose(\n",
      "                            array(\n",
      "                                [\n",
      "                                    [\n",
      "                                        ones(len(-phi)),\n",
      "                                        zeros(len(-phi)),\n",
      "                                        zeros(len(-phi)),\n",
      "                                        zeros(len(-phi)),\n",
      "                                    ],\n",
      "                                    [\n",
      "                                        zeros(len(-phi)),\n",
      "                                        cos(-phi),\n",
      "                                        -sin(-phi),\n",
      "                                        zeros(len(-phi)),\n",
      "                                    ],\n",
      "                                    [\n",
      "                                        zeros(len(-phi)),\n",
      "                                        sin(-phi),\n",
      "                                        cos(-phi),\n",
      "                                        zeros(len(-phi)),\n",
      "                                    ],\n",
      "                                    [\n",
      "                                        zeros(len(-phi)),\n",
      "                                        zeros(len(-phi)),\n",
      "                                        zeros(len(-phi)),\n",
      "                                        ones(len(-phi)),\n",
      "                                    ],\n",
      "                                ]\n",
      "                            ).transpose(2, 0, 1),\n",
      "                            axes=(1, 2, 0),\n",
      "                        ),\n",
      "                        transpose(p),\n",
      "                    )\n",
      "                ),\n",
      "            )\n",
      "        ),\n",
      "    )\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print_lambdify([beta, theta, momentum], boosted_momentum)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ":::{note}\n",
    "\n",
    "The code above contains a lot of duplicate code, such as `len(-phi)`. This could possibly be improved with {class}`~sympy.codegen.ast.CodeBlock`. See [ampform#166](https://github.com/ComPWA/ampform/issues/166).\n",
    "\n",
    ":::"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Angle computation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Computing phi"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "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](https://github.com/ComPWA/ampform/blob/f38179483ac20664584014d9eb659866966681c2/src/ampform/data.py#L161-L162)). 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."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "scroll-input"
    ]
   },
   "outputs": [],
   "source": [
    "@implement_doit_method()\n",
    "class FourMomentumX(UnevaluatedExpression):\n",
    "    momentum: ArraySymbol = property(lambda self: self.args[0])\n",
    "\n",
    "    def __new__(cls, momentum: ArraySymbol, **hints) -> FourMomentumX:\n",
    "        return create_expression(cls, momentum, **hints)\n",
    "\n",
    "    def evaluate(self) -> ArraySlice:\n",
    "        return ArraySlice(self.momentum, (slice(None), 1))\n",
    "\n",
    "    def _latex(self, printer, *args) -> str:\n",
    "        momentum = printer._print(self.momentum)\n",
    "        return f\"{{{momentum}}}_{{x}}\"\n",
    "\n",
    "\n",
    "@implement_doit_method()\n",
    "class FourMomentumY(UnevaluatedExpression):\n",
    "    momentum: ArraySymbol = property(lambda self: self.args[0])\n",
    "\n",
    "    def __new__(cls, momentum: ArraySymbol, **hints) -> FourMomentumY:\n",
    "        return create_expression(cls, momentum, **hints)\n",
    "\n",
    "    def evaluate(self) -> ArraySlice:\n",
    "        return ArraySlice(self.momentum, (slice(None), 2))\n",
    "\n",
    "    def _latex(self, printer, *args) -> str:\n",
    "        momentum = printer._print(self.momentum)\n",
    "        return f\"{{{momentum}}}_{{y}}\"\n",
    "\n",
    "\n",
    "@implement_doit_method()\n",
    "class FourMomentumZ(UnevaluatedExpression):\n",
    "    momentum: ArraySymbol = property(lambda self: self.args[0])\n",
    "\n",
    "    def __new__(cls, momentum: ArraySymbol, **hints) -> FourMomentumZ:\n",
    "        return create_expression(cls, momentum, **hints)\n",
    "\n",
    "    def evaluate(self) -> ArraySlice:\n",
    "        return ArraySlice(self.momentum, (slice(None), 3))\n",
    "\n",
    "    def _latex(self, printer, *args) -> str:\n",
    "        momentum = printer._print(self.momentum)\n",
    "        return f\"{{{momentum}}}_{{y}}\"\n",
    "\n",
    "\n",
    "@implement_doit_method()\n",
    "class Energy(UnevaluatedExpression):\n",
    "    momentum: ArraySymbol = property(lambda self: self.args[0])\n",
    "\n",
    "    def __new__(cls, momentum: ArraySymbol, **hints) -> Energy:\n",
    "        return create_expression(cls, momentum, **hints)\n",
    "\n",
    "    def evaluate(self) -> ArraySlice:\n",
    "        return ArraySlice(self.momentum, (slice(None), 0))\n",
    "\n",
    "    def _latex(self, printer, *args) -> str:\n",
    "        momentum = printer._print(self.momentum)\n",
    "        return f\"{{E}}_{{{momentum}}}\"\n",
    "\n",
    "\n",
    "@implement_doit_method()\n",
    "class Phi(UnevaluatedExpression):\n",
    "    momentum: ArraySymbol = property(lambda self: self.args[0])\n",
    "\n",
    "    def __new__(cls, momentum: ArraySymbol, **hints) -> Phi:\n",
    "        return create_expression(cls, momentum, **hints)\n",
    "\n",
    "    def evaluate(self) -> sp.Expr:\n",
    "        p = self.momentum\n",
    "        return sp.atan2(FourMomentumY(p), FourMomentumX(p))\n",
    "\n",
    "    def _latex(self, printer, *args) -> str:\n",
    "        momentum = printer._print(self.momentum)\n",
    "        return Rf\"\\phi\\left({momentum}\\right)\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The classes indeed render nicely as LaTeX:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input",
     "keep_output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle p_{1}$"
      ],
      "text/plain": [
       "p_1"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\begin{eqnarray}\n",
       "\\phi\\left(p_{1}\\right) & = & \\operatorname{atan_{2}}{\\left({p_{1}}_{y},{p_{1}}_{x} \\right)} \n",
       "\\end{eqnarray}$"
      ],
      "text/plain": [
       "<IPython.core.display.Math object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\begin{eqnarray}\n",
       "{p_{1}}_{x} & = & p_{1}\\left[:, 1\\right] \n",
       "\\end{eqnarray}$"
      ],
      "text/plain": [
       "<IPython.core.display.Math object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\begin{eqnarray}\n",
       "{p_{1}}_{y} & = & p_{1}\\left[:, 2\\right] \n",
       "\\end{eqnarray}$"
      ],
      "text/plain": [
       "<IPython.core.display.Math object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\begin{eqnarray}\n",
       "{E}_{p_{1}} & = & p_{1}\\left[:, 0\\right] \n",
       "\\end{eqnarray}$"
      ],
      "text/plain": [
       "<IPython.core.display.Math object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "p = ArraySymbol(\"p_1\")\n",
    "phi = Phi(p)\n",
    "p_x = FourMomentumX(p)\n",
    "p_y = FourMomentumY(p)\n",
    "energy = Energy(p)\n",
    "\n",
    "math_style = {\"environment\": \"eqnarray\"}\n",
    "display(\n",
    "    p,\n",
    "    Math(sp.multiline_latex(Phi(p), phi.evaluate(), **math_style)),\n",
    "    Math(sp.multiline_latex(p_x, p_x.doit(), **math_style)),\n",
    "    Math(sp.multiline_latex(p_y, p_y.doit(), **math_style)),\n",
    "    Math(sp.multiline_latex(energy, energy.doit(), **math_style)),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that the four classes like `FourMomentumX` expect an `ArraySymbol` as input. This requires [sympy/sympy#22265](https://github.com/sympy/sympy/issues/22265). This allows lambdifying the above expressions to valid NumPy code. Let's compare this with the existing implementation using the {ref}`report/011:Test sample`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "keep_output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[ 1.35527  ,  0.514208 , -0.184219 ,  1.23296  ],\n",
       "       [ 0.841933 ,  0.0727385, -0.0528868,  0.826163 ],\n",
       "       [ 0.550927 , -0.162529 ,  0.29976  , -0.411133 ]])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "momentum_sample = events[0]\n",
    "np.array(momentum_sample)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "keep_output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([ 0.514208 ,  0.0727385, -0.162529 ])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np_expr = sp.lambdify(p, p_x.doit())\n",
    "np_expr(momentum_sample)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "assert np.array_equal(_, momentum_sample.p_x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "keep_output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([1.35527 , 0.841933, 0.550927])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np_expr = sp.lambdify(p, energy.doit())\n",
    "np_expr(momentum_sample)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "assert np.array_equal(_, momentum_sample.energy)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "keep_output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([-0.34401236, -0.62867104,  2.06762909])"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "array([-0.34401236, -0.62867104,  2.06762909])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np_expr = sp.lambdify(p, phi.doit())\n",
    "display(np.array(momentum_sample.phi()))\n",
    "np_expr(momentum_sample)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "assert np.array_equal(_, momentum_sample.phi())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Computing theta"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Computing $\\theta$ is more complicated, because requires the norm of the three-momentum of $p$ (see [existing implementation](https://github.com/ComPWA/ampform/blob/f38179483ac20664584014d9eb659866966681c2/src/ampform/data.py#L164-L165)). In other words:\n",
    "\n",
    "$$\n",
    "\\theta = \\arccos\\left(\\frac{p_z}{\\left|\\vec{p}\\right|}\\right)\n",
    "\\quad \\mathrm{with} \\quad\n",
    "\\left|\\vec{p}\\right| = \\sqrt{p_x^2+p_y^2+p_z^2}\n",
    "$$\n",
    "\n",
    "The complication here is that $\\left|\\vec{p}\\right|$ [needs to be computed with `np.sum`](https://github.com/ComPWA/ampform/blob/f381794/src/ampform/data.py#L157-L159) 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](https://github.com/sympy/sympy/issues/22279)). So for now, an intermediate `ArrayAxisSum` class has to be written for this. See also [this remark](https://github.com/sympy/sympy/discussions/22269#discussioncomment-1473005)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "class ArrayAxisSum(sp.Expr):\n",
    "    array: ArraySymbol = property(lambda self: self.args[0])\n",
    "    axis: int | None = property(lambda self: self.args[1])\n",
    "\n",
    "    def __new__(\n",
    "        cls, array: ArraySymbol, axis: int | None = None, **hints\n",
    "    ) -> ArrayAxisSum:\n",
    "        if axis is not None and not isinstance(axis, (int, sp.Integer)):\n",
    "            msg = \"Only single digits allowed for axis\"\n",
    "            raise TypeError(msg)\n",
    "        return create_expression(cls, array, axis, **hints)\n",
    "\n",
    "    def _latex(self, printer, *args) -> str:\n",
    "        A = printer._print(self.array)\n",
    "        if self.axis is None:\n",
    "            return Rf\"\\sum{{{A}}}\"\n",
    "        axis = printer._print(self.axis)\n",
    "        return Rf\"\\sum_{{\\mathrm{{axis{axis}}}}}{{{A}}}\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Looks nice as LaTeX:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "keep_output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\sum_{\\mathrm{axis1}}{A}$"
      ],
      "text/plain": [
       "ArrayAxisSum(A, 1)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\sum{A}$"
      ],
      "text/plain": [
       "ArrayAxisSum(A, None)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "A = ArraySymbol(\"A\")\n",
    "display(\n",
    "    ArrayAxisSum(A, axis=1),\n",
    "    ArrayAxisSum(A),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's define a printer method for NumPy:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def _print_ArrayAxisSum(self: NumPyPrinter, expr: ArrayAxisSum) -> str:\n",
    "    self.module_imports[\"numpy\"].add(\"sum\")\n",
    "    array = self._print(expr.array)\n",
    "    axis = self._print(expr.axis)\n",
    "    return f\"sum({array}, axis={axis})\"\n",
    "\n",
    "\n",
    "NumPyPrinter._print_ArrayAxisSum = _print_ArrayAxisSum"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "keep_output"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "def _lambdifygenerated(A):\n",
      "    return sum(A, axis=1)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print_lambdify(A, ArrayAxisSum(A, axis=1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "...and let's check whether it works as expected for a 3-dimensional array:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "keep_output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[[ 0,  1],\n",
       "        [ 2,  3],\n",
       "        [ 4,  5]],\n",
       "\n",
       "       [[ 6,  7],\n",
       "        [ 8,  9],\n",
       "        [10, 11]]])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "array = np.array(range(12)).reshape(2, 3, 2)\n",
    "array"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "keep_output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "66"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np_expr = sp.lambdify(A, ArrayAxisSum(A))\n",
    "np_expr(array)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "assert _ == np.sum(array)\n",
    "assert _ == sum(range(12))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "keep_output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[ 6,  8],\n",
       "       [10, 12],\n",
       "       [14, 16]])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np_expr = sp.lambdify(A, ArrayAxisSum(A, axis=0))\n",
    "np_expr(array)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "assert np.array_equal(_, np.sum(array, axis=0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "keep_output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[ 6,  9],\n",
       "       [24, 27]])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np_expr = sp.lambdify(A, ArrayAxisSum(A, axis=1))\n",
    "np_expr(array)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "assert np.array_equal(_, np.sum(array, axis=1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "keep_output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[ 1,  5,  9],\n",
       "       [13, 17, 21]])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np_expr = sp.lambdify(A, ArrayAxisSum(A, axis=2))\n",
    "np_expr(array)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "assert np.array_equal(_, np.sum(array, axis=2))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "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](https://github.com/ComPWA/ampform/blob/f38179483ac20664584014d9eb659866966681c2/src/ampform/data.py#L153-L159)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@implement_doit_method()\n",
    "class ThreeMomentumNorm(UnevaluatedExpression):\n",
    "    momentum: ArraySymbol = property(lambda self: self.args[0])\n",
    "\n",
    "    def __new__(cls, momentum: ArraySymbol, **hints) -> ThreeMomentumNorm:\n",
    "        return create_expression(cls, momentum, **hints)\n",
    "\n",
    "    def evaluate(self) -> ArraySlice:\n",
    "        three_momentum = ArraySlice(self.momentum, (slice(None), slice(1, None)))\n",
    "        norm_squared = ArrayAxisSum(three_momentum**2, axis=1)\n",
    "        return sp.sqrt(norm_squared)\n",
    "\n",
    "    def _latex(self, printer, *args) -> str:\n",
    "        three_momentum = printer._print(self.momentum)\n",
    "        return Rf\"\\left|\\vec{{{three_momentum}}}\\right|\"\n",
    "\n",
    "    def _numpycode(self, printer, *args) -> str:\n",
    "        return printer._print(self.evaluate())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input",
     "keep_output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\begin{eqnarray}\n",
       "\\left|\\vec{p_{1}}\\right| & = & \\sqrt{\\sum_{\\mathrm{axis1}}{p_{1}\\left[:, 1:\\right]^{2}}} \n",
       "\\end{eqnarray}$"
      ],
      "text/plain": [
       "<IPython.core.display.Math object>"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "p_norm = ThreeMomentumNorm(p)\n",
    "Math(sp.multiline_latex(p_norm, p_norm.doit(), **math_style))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "keep_output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([1.34853137, 0.83104344, 0.53413676])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np_expr = sp.lambdify(p, p_norm.doit())\n",
    "np_expr(momentum_sample)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "assert np.array_equal(_, momentum_sample.p_norm())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With that, we're ready to define the `Theta` class!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@implement_doit_method()\n",
    "class Theta(UnevaluatedExpression):\n",
    "    momentum: ArraySymbol = property(lambda self: self.args[0])\n",
    "\n",
    "    def __new__(cls, momentum: ArraySymbol, **hints) -> Theta:\n",
    "        return create_expression(cls, momentum, **hints)\n",
    "\n",
    "    def evaluate(self) -> sp.Expr:\n",
    "        p = self.momentum\n",
    "        return sp.acos(FourMomentumZ(p) / ThreeMomentumNorm(p))\n",
    "\n",
    "    def _latex(self, printer, *args) -> str:\n",
    "        momentum = printer._print(self.momentum)\n",
    "        return Rf\"\\theta\\left({momentum}\\right)\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The math doesn't look the best, but this due to the {ref}`ArrayMultiplication <011-array-multiplication>` class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input",
     "keep_output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\begin{eqnarray}\n",
       "\\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)} \n",
       "\\end{eqnarray}$"
      ],
      "text/plain": [
       "<IPython.core.display.Math object>"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Math(sp.multiline_latex(Theta(p), Theta(p).doit(), **math_style))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "At any rate, when we lambdify the whole thing, we have exactly the same result as the [original `theta()` method](https://github.com/ComPWA/ampform/blob/f38179483ac20664584014d9eb659866966681c2/src/ampform/data.py#L164-L165) gave!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "keep_output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([0.41702412, 0.10842903, 2.4491907 ])"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "array([0.41702412, 0.10842903, 2.4491907 ])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np_expr = sp.lambdify(p, Theta(p).doit())\n",
    "display(np.array(momentum_sample.theta()))\n",
    "np_expr(momentum_sample)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "assert np.allclose(_, momentum_sample.theta())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Recursive angle computation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we are ready to compute kinematic angles recursively from a {class}`~qrules.topology.Topology`, just as in the [existing implementation](https://github.com/ComPWA/ampform/blob/307bbb1a09e75a320fe812945d7f0c7e0cdcc9f8/src/ampform/kinematics.py#L204-L280), but now with SymPy only."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "keep_output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{0: p0, 1: p1, 2: p2, 3: p3}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def create_event_collection_from_topology(\n",
    "    topology: Topology,\n",
    ") -> dict[int, ArraySymbol]:\n",
    "    n_final_states = len(topology.outgoing_edge_ids)\n",
    "    return {i: ArraySymbol(f\"p{i}\") for i in range(n_final_states)}\n",
    "\n",
    "\n",
    "momentum_symbols = create_event_collection_from_topology(topology)\n",
    "momentum_symbols"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now it's quite trivial to rewrite the [existing implementation](https://github.com/ComPWA/ampform/blob/307bbb1a09e75a320fe812945d7f0c7e0cdcc9f8/src/ampform/kinematics.py#L204-L280) with the classes defined above.\n",
    "\n",
    "```{autolink-skip}\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "scroll-input"
    ]
   },
   "outputs": [],
   "source": [
    "def compute_helicity_angles(\n",
    "    events: dict[int, ArraySymbol], topology: Topology\n",
    ") -> dict[str, sp.Expr]:\n",
    "    if topology.outgoing_edge_ids != set(events):\n",
    "        msg = (\n",
    "            f\"Momentum IDs {set(events)} do not match final state edge IDs\"\n",
    "            f\" {set(topology.outgoing_edge_ids)}\"\n",
    "        )\n",
    "        raise ValueError(msg)\n",
    "\n",
    "    def __recursive_helicity_angles(\n",
    "        events: dict[int, ArraySymbol], node_id: int\n",
    "    ) -> dict[str, sp.Expr]:\n",
    "        helicity_angles: dict[str, sp.Expr] = {}\n",
    "        child_state_ids = sorted(topology.get_edge_ids_outgoing_from_node(node_id))\n",
    "        if all(topology.edges[i].ending_node_id is None for i in child_state_ids):\n",
    "            state_id = child_state_ids[0]\n",
    "            four_momentum = events[state_id]\n",
    "            phi_label, theta_label = get_helicity_angle_label(topology, state_id)\n",
    "            helicity_angles[phi_label] = Phi(four_momentum)\n",
    "            helicity_angles[theta_label] = Theta(four_momentum)\n",
    "        for state_id in child_state_ids:\n",
    "            edge = topology.edges[state_id]\n",
    "            if edge.ending_node_id is not None:\n",
    "                # recursively determine all momenta ids in the list\n",
    "                sub_momenta_ids = determine_attached_final_state(topology, state_id)\n",
    "                if len(sub_momenta_ids) > 1:\n",
    "                    # add all of these momenta together -> defines new subsystem\n",
    "                    four_momentum = sum(events[i] for i in sub_momenta_ids)\n",
    "\n",
    "                    # boost all of those momenta into this new subsystem\n",
    "                    phi = Phi(four_momentum)\n",
    "                    theta = Theta(four_momentum)\n",
    "                    p3_norm = ThreeMomentumNorm(four_momentum)\n",
    "                    beta = p3_norm / Energy(four_momentum)\n",
    "                    new_momentum_pool = {\n",
    "                        k: ArrayMultiplication(\n",
    "                            BoostZ(beta),\n",
    "                            RotationY(-theta),\n",
    "                            RotationZ(-phi),\n",
    "                            p,\n",
    "                        )\n",
    "                        for k, p in events.items()\n",
    "                        if k in sub_momenta_ids\n",
    "                    }\n",
    "\n",
    "                    # register current angle variables\n",
    "                    phi_label, theta_label = get_helicity_angle_label(\n",
    "                        topology, state_id\n",
    "                    )\n",
    "                    helicity_angles[phi_label] = Phi(four_momentum)\n",
    "                    helicity_angles[theta_label] = Theta(four_momentum)\n",
    "\n",
    "                    # call next recursion\n",
    "                    angles = __recursive_helicity_angles(\n",
    "                        new_momentum_pool,\n",
    "                        edge.ending_node_id,\n",
    "                    )\n",
    "                    helicity_angles.update(angles)\n",
    "\n",
    "        return helicity_angles\n",
    "\n",
    "    initial_state_id = next(iter(topology.incoming_edge_ids))\n",
    "    initial_state_edge = topology.edges[initial_state_id]\n",
    "    assert initial_state_edge.ending_node_id is not None\n",
    "    return __recursive_helicity_angles(events, initial_state_edge.ending_node_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The computation works indeed and can be rendered to both LaTeX and {mod}`numpy`!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "keep_output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['phi_1+2+3',\n",
       " 'theta_1+2+3',\n",
       " 'phi_2+3,1+2+3',\n",
       " 'theta_2+3,1+2+3',\n",
       " 'phi_2,2+3,1+2+3',\n",
       " 'theta_2,2+3,1+2+3']"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "symbolic_angles = compute_helicity_angles(momentum_symbols, topology)\n",
    "list(symbolic_angles)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input",
     "scroll-output",
     "full-width",
     "keep_output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\begin{eqnarray}\n",
       "\\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)} \n",
       "\\end{eqnarray}$"
      ],
      "text/plain": [
       "<IPython.core.display.Math object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "def _lambdifygenerated(p0, p1, p2, p3):\n",
      "    return arctan2((p1 + p2 + p3)[:, 2], (p1 + p2 + p3)[:, 1])\n",
      "\n",
      "\n"
     ]
    },
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\begin{eqnarray}\n",
       "\\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)} \n",
       "\\end{eqnarray}$"
      ],
      "text/plain": [
       "<IPython.core.display.Math object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "def _lambdifygenerated(p0, p1, p2, p3):\n",
      "    return arccos(\n",
      "        (p1 + p2 + p3)[:, 3]\n",
      "        * sum((p1 + p2 + p3)[:, 1:] ** 2, axis=1) ** (-1 / 2)\n",
      "    )\n",
      "\n",
      "\n"
     ]
    },
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\begin{eqnarray}\n",
       "\\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)} \n",
       "\\end{eqnarray}$"
      ],
      "text/plain": [
       "<IPython.core.display.Math object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "def _lambdifygenerated(p0, p1, p2, p3):\n",
      "    return arctan2(\n",
      "        (\n",
      "            einsum(\n",
      "                \"ij...,j...\",\n",
      "                transpose(\n",
      "                    array(\n",
      "                        [\n",
      "                            [\n",
      "                                1\n",
      "                                / sqrt(\n",
      "                                    1\n",
      "                                    - (\n",
      "                                        sqrt(\n",
      "                                            sum(\n",
      "                                                (p1 + p2 + p3)[:, 1:] ** 2,\n",
      "                                                axis=1,\n",
      "                                            )\n",
      "                                        )\n",
      "                                        * (p1 + p2 + p3)[:, 0] ** (-1.0)\n",
      "                                    )\n",
      "                                    ** 2\n",
      "                                ),\n",
      "                                zeros(\n",
      "                                    len(\n",
      "                                        sqrt(\n",
      "                                            sum(\n",
      "                                                (p1 + p2 + p3)[:, 1:] ** 2,\n",
      "                                                axis=1,\n",
      "                                            )\n",
      "                                        )\n",
      "                                        * (p1 + p2 + p3)[:, 0] ** (-1.0)\n",
      "                                    )\n",
      "                                ),\n",
      "                                zeros(\n",
      "                                    len(\n",
      "                                        sqrt(\n",
      "                                            sum(\n",
      "                                                (p1 + p2 + p3)[:, 1:] ** 2,\n",
      "                                                axis=1,\n",
      "                                            )\n",
      "                                        )\n",
      "                                        * (p1 + p2 + p3)[:, 0] ** (-1.0)\n",
      "                                    )\n",
      "                                ),\n",
      "                                -1\n",
      "                                / sqrt(\n",
      "                                    1\n",
      "                                    - (\n",
      "                                        sqrt(\n",
      "                                            sum(\n",
      "                                                (p1 + p2 + p3)[:, 1:] ** 2,\n",
      "                                                axis=1,\n",
      "                                            )\n",
      "                                        )\n",
      "                                        * (p1 + p2 + p3)[:, 0] ** (-1.0)\n",
      "                                    )\n",
      "                                    ** 2\n",
      "                                )\n",
      "                                * sqrt(sum((p1 + p2 + p3)[:, 1:] ** 2, axis=1))\n",
      "                                * (p1 + p2 + p3)[:, 0] ** (-1.0),\n",
      "                            ],\n",
      "                            [\n",
      "                                zeros(\n",
      "                                    len(\n",
      "                                        sqrt(\n",
      "                                            sum(\n",
      "                                                (p1 + p2 + p3)[:, 1:] ** 2,\n",
      "                                                axis=1,\n",
      "                                            )\n",
      "                                        )\n",
      "                                        * (p1 + p2 + p3)[:, 0] ** (-1.0)\n",
      "                                    )\n",
      "                                ),\n",
      "                                ones(\n",
      "                                    len(\n",
      "                                        sqrt(\n",
      "                                            sum(\n",
      "                                                (p1 + p2 + p3)[:, 1:] ** 2,\n",
      "                                                axis=1,\n",
      "                                            )\n",
      "                                        )\n",
      "                                        * (p1 + p2 + p3)[:, 0] ** (-1.0)\n",
      "                                    )\n",
      "                                ),\n",
      "                                zeros(\n",
      "                                    len(\n",
      "                                        sqrt(\n",
      "                                            sum(\n",
      "                                                (p1 + p2 + p3)[:, 1:] ** 2,\n",
      "                                                axis=1,\n",
      "                                            )\n",
      "                                        )\n",
      "                                        * (p1 + p2 + p3)[:, 0] ** (-1.0)\n",
      "                                    )\n",
      "                                ),\n",
      "                                zeros(\n",
      "                                    len(\n",
      "                                        sqrt(\n",
      "                                            sum(\n",
      "                                                (p1 + p2 + p3)[:, 1:] ** 2,\n",
      "                                                axis=1,\n",
      "                                            )\n",
      "                                        )\n",
      "                                        * (p1 + p2 + p3)[:, 0] ** (-1.0)\n",
      "                                    )\n",
      "                                ),\n",
      "                            ],\n",
      "                            [\n",
      "                                zeros(\n",
      "                                    len(\n",
      "                                        sqrt(\n",
      "                                            sum(\n",
      "                                                (p1 + p2 + p3)[:, 1:] ** 2,\n",
      "                                                axis=1,\n",
      "                                            )\n",
      "                                        )\n",
      "                                        * (p1 + p2 + p3)[:, 0] ** (-1.0)\n",
      "                                    )\n",
      "                                ),\n",
      "                                zeros(\n",
      "                                    len(\n",
      "                                        sqrt(\n",
      "                                            sum(\n",
      "                                                (p1 + p2 + p3)[:, 1:] ** 2,\n",
      "                                                axis=1,\n",
      "                                            )\n",
      "                                        )\n",
      "                                        * (p1 + p2 + p3)[:, 0] ** (-1.0)\n",
      "                                    )\n",
      "                                ),\n",
      "                                ones(\n",
      "                                    len(\n",
      "                                        sqrt(\n",
      "                                            sum(\n",
      "                                                (p1 + p2 + p3)[:, 1:] ** 2,\n",
      "                                                axis=1,\n",
      "                                            )\n",
      "                                        )\n",
      "                                        * (p1 + p2 + p3)[:, 0] ** (-1.0)\n",
      "                                    )\n",
      "                                ),\n",
      "                                zeros(\n",
      "                                    len(\n",
      "                                        sqrt(\n",
      "                                            sum(\n",
      "                                                (p1 + p2 + p3)[:, 1:] ** 2,\n",
      "                                                axis=1,\n",
      "                                            )\n",
      "                                        )\n",
      "                                        * (p1 + p2 + p3)[:, 0] ** (-1.0)\n",
      "                                    )\n",
      "                                ),\n",
      "                            ],\n",
      "                            [\n",
      "                                -1\n",
      "                                / sqrt(\n",
      "                                    1\n",
      "                                    - (\n",
      "                                        sqrt(\n",
      "                                            sum(\n",
      "                                                (p1 + p2 + p3)[:, 1:] ** 2,\n",
      "                                                axis=1,\n",
      "                                            )\n",
      "                                        )\n",
      "                                        * (p1 + p2 + p3)[:, 0] ** (-1.0)\n",
      "                                    )\n",
      "                                    ** 2\n",
      "                                )\n",
      "                                * sqrt(sum((p1 + p2 + p3)[:, 1:] ** 2, axis=1))\n",
      "                                * (p1 + p2 + p3)[:, 0] ** (-1.0),\n",
      "                                zeros(\n",
      "                                    len(\n",
      "                                        sqrt(\n",
      "                                            sum(\n",
      "                                                (p1 + p2 + p3)[:, 1:] ** 2,\n",
      "                                                axis=1,\n",
      "                                            )\n",
      "                                        )\n",
      "                                        * (p1 + p2 + p3)[:, 0] ** (-1.0)\n",
      "                                    )\n",
      "                                ),\n",
      "                                zeros(\n",
      "                                    len(\n",
      "                                        sqrt(\n",
      "                                            sum(\n",
      "                                                (p1 + p2 + p3)[:, 1:] ** 2,\n",
      "                                                axis=1,\n",
      "                                            )\n",
      "                                        )\n",
      "                                        * (p1 + p2 + p3)[:, 0] ** (-1.0)\n",
      "                                    )\n",
      "                                ),\n",
      "                                1\n",
      "                                / sqrt(\n",
      "                                    1\n",
      "                                    - (\n",
      "                                        sqrt(\n",
      "                                            sum(\n",
      "                                                (p1 + p2 + p3)[:, 1:] ** 2,\n",
      "                                                axis=1,\n",
      "                                            )\n",
      "                                        )\n",
      "                                        * (p1 + p2 + p3)[:, 0] ** (-1.0)\n",
      "                                    )\n",
      "                                    ** 2\n",
      "                                ),\n",
      "                            ],\n",
      "                        ]\n",
      "                    ).transpose(2, 0, 1),\n",
      "                    axes=(1, 2, 0),\n",
      "                ),\n",
      "                transpose(\n",
      "                    einsum(\n",
      "                        \"ij...,j...\",\n",
      "                        transpose(\n",
      "                            array(\n",
      "                                [\n",
      "                                    [\n",
      "                                        ones(\n",
      "                                            len(\n",
      "                                                -arccos(\n",
      "                                                    (p1 + p2 + p3)[:, 3]\n",
      "                                                    * sum(\n",
      "                                                        (p1 + p2 + p3)[:, 1:]\n",
      "                                                        ** 2,\n",
      "                                                        axis=1,\n",
      "                                                    )\n",
      "                                                    ** (-1 / 2)\n",
      "                                                )\n",
      "                                            )\n",
      "                                        ),\n",
      "                                        zeros(\n",
      "                                            len(\n",
      "                                                -arccos(\n",
      "                                                    (p1 + p2 + p3)[:, 3]\n",
      "                                                    * sum(\n",
      "                                                        (p1 + p2 + p3)[:, 1:]\n",
      "                                                        ** 2,\n",
      "                                                        axis=1,\n",
      "                                                    )\n",
      "                                                    ** (-1 / 2)\n",
      "                                                )\n",
      "                                            )\n",
      "                                        ),\n",
      "                                        zeros(\n",
      "                                            len(\n",
      "                                                -arccos(\n",
      "                                                    (p1 + p2 + p3)[:, 3]\n",
      "                                                    * sum(\n",
      "                                                        (p1 + p2 + p3)[:, 1:]\n",
      "                                                        ** 2,\n",
      "                                                        axis=1,\n",
      "                                                    )\n",
      "                                                    ** (-1 / 2)\n",
      "                                                )\n",
      "                                            )\n",
      "                                        ),\n",
      "                                        zeros(\n",
      "                                            len(\n",
      "                                                -arccos(\n",
      "                                                    (p1 + p2 + p3)[:, 3]\n",
      "                                                    * sum(\n",
      "                                                        (p1 + p2 + p3)[:, 1:]\n",
      "                                                        ** 2,\n",
      "                                                        axis=1,\n",
      "                                                    )\n",
      "                                                    ** (-1 / 2)\n",
      "                                                )\n",
      "                                            )\n",
      "                                        ),\n",
      "                                    ],\n",
      "                                    [\n",
      "                                        zeros(\n",
      "                                            len(\n",
      "                                                -arccos(\n",
      "                                                    (p1 + p2 + p3)[:, 3]\n",
      "                                                    * sum(\n",
      "                                                        (p1 + p2 + p3)[:, 1:]\n",
      "                                                        ** 2,\n",
      "                                                        axis=1,\n",
      "                                                    )\n",
      "                                                    ** (-1 / 2)\n",
      "                                                )\n",
      "                                            )\n",
      "                                        ),\n",
      "                                        cos(\n",
      "                                            -arccos(\n",
      "                                                (p1 + p2 + p3)[:, 3]\n",
      "                                                * sum(\n",
      "                                                    (p1 + p2 + p3)[:, 1:] ** 2,\n",
      "                                                    axis=1,\n",
      "                                                )\n",
      "                                                ** (-1 / 2)\n",
      "                                            )\n",
      "                                        ),\n",
      "                                        zeros(\n",
      "                                            len(\n",
      "                                                -arccos(\n",
      "                                                    (p1 + p2 + p3)[:, 3]\n",
      "                                                    * sum(\n",
      "                                                        (p1 + p2 + p3)[:, 1:]\n",
      "                                                        ** 2,\n",
      "                                                        axis=1,\n",
      "                                                    )\n",
      "                                                    ** (-1 / 2)\n",
      "                                                )\n",
      "                                            )\n",
      "                                        ),\n",
      "                                        sin(\n",
      "                                            -arccos(\n",
      "                                                (p1 + p2 + p3)[:, 3]\n",
      "                                                * sum(\n",
      "                                                    (p1 + p2 + p3)[:, 1:] ** 2,\n",
      "                                                    axis=1,\n",
      "                                                )\n",
      "                                                ** (-1 / 2)\n",
      "                                            )\n",
      "                                        ),\n",
      "                                    ],\n",
      "                                    [\n",
      "                                        zeros(\n",
      "                                            len(\n",
      "                                                -arccos(\n",
      "                                                    (p1 + p2 + p3)[:, 3]\n",
      "                                                    * sum(\n",
      "                                                        (p1 + p2 + p3)[:, 1:]\n",
      "                                                        ** 2,\n",
      "                                                        axis=1,\n",
      "                                                    )\n",
      "                                                    ** (-1 / 2)\n",
      "                                                )\n",
      "                                            )\n",
      "                                        ),\n",
      "                                        zeros(\n",
      "                                            len(\n",
      "                                                -arccos(\n",
      "                                                    (p1 + p2 + p3)[:, 3]\n",
      "                                                    * sum(\n",
      "                                                        (p1 + p2 + p3)[:, 1:]\n",
      "                                                        ** 2,\n",
      "                                                        axis=1,\n",
      "                                                    )\n",
      "                                                    ** (-1 / 2)\n",
      "                                                )\n",
      "                                            )\n",
      "                                        ),\n",
      "                                        ones(\n",
      "                                            len(\n",
      "                                                -arccos(\n",
      "                                                    (p1 + p2 + p3)[:, 3]\n",
      "                                                    * sum(\n",
      "                                                        (p1 + p2 + p3)[:, 1:]\n",
      "                                                        ** 2,\n",
      "                                                        axis=1,\n",
      "                                                    )\n",
      "                                                    ** (-1 / 2)\n",
      "                                                )\n",
      "                                            )\n",
      "                                        ),\n",
      "                                        zeros(\n",
      "                                            len(\n",
      "                                                -arccos(\n",
      "                                                    (p1 + p2 + p3)[:, 3]\n",
      "                                                    * sum(\n",
      "                                                        (p1 + p2 + p3)[:, 1:]\n",
      "                                                        ** 2,\n",
      "                                                        axis=1,\n",
      "                                                    )\n",
      "                                                    ** (-1 / 2)\n",
      "                                                )\n",
      "                                            )\n",
      "                                        ),\n",
      "                                    ],\n",
      "                                    [\n",
      "                                        zeros(\n",
      "                                            len(\n",
      "                                                -arccos(\n",
      "                                                    (p1 + p2 + p3)[:, 3]\n",
      "                                                    * sum(\n",
      "                                                        (p1 + p2 + p3)[:, 1:]\n",
      "                                                        ** 2,\n",
      "                                                        axis=1,\n",
      "                                                    )\n",
      "                                                    ** (-1 / 2)\n",
      "                                                )\n",
      "                                            )\n",
      "                                        ),\n",
      "                                        -sin(\n",
      "                                            -arccos(\n",
      "                                                (p1 + p2 + p3)[:, 3]\n",
      "                                                * sum(\n",
      "                                                    (p1 + p2 + p3)[:, 1:] ** 2,\n",
      "                                                    axis=1,\n",
      "                                                )\n",
      "                                                ** (-1 / 2)\n",
      "                                            )\n",
      "                                        ),\n",
      "                                        zeros(\n",
      "                                            len(\n",
      "                                                -arccos(\n",
      "                                                    (p1 + p2 + p3)[:, 3]\n",
      "                                                    * sum(\n",
      "                                                        (p1 + p2 + p3)[:, 1:]\n",
      "                                                        ** 2,\n",
      "                                                        axis=1,\n",
      "                                                    )\n",
      "                                                    ** (-1 / 2)\n",
      "                                                )\n",
      "                                            )\n",
      "                                        ),\n",
      "                                        cos(\n",
      "                                            -arccos(\n",
      "                                                (p1 + p2 + p3)[:, 3]\n",
      "                                                * sum(\n",
      "                                                    (p1 + p2 + p3)[:, 1:] ** 2,\n",
      "                                                    axis=1,\n",
      "                                                )\n",
      "                                                ** (-1 / 2)\n",
      "                                            )\n",
      "                                        ),\n",
      "                                    ],\n",
      "                                ]\n",
      "                            ).transpose(2, 0, 1),\n",
      "                            axes=(1, 2, 0),\n",
      "                        ),\n",
      "                        transpose(\n",
      "                            einsum(\n",
      "                                \"ij...,j...\",\n",
      "                                transpose(\n",
      "                                    array(\n",
      "                                        [\n",
      "                                            [\n",
      "                                                ones(\n",
      "                                                    len(\n",
      "                                                        -arctan2(\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 2\n",
      "                                                            ],\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 1\n",
      "                                                            ],\n",
      "                                                        )\n",
      "                                                    )\n",
      "                                                ),\n",
      "                                                zeros(\n",
      "                                                    len(\n",
      "                                                        -arctan2(\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 2\n",
      "                                                            ],\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 1\n",
      "                                                            ],\n",
      "                                                        )\n",
      "                                                    )\n",
      "                                                ),\n",
      "                                                zeros(\n",
      "                                                    len(\n",
      "                                                        -arctan2(\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 2\n",
      "                                                            ],\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 1\n",
      "                                                            ],\n",
      "                                                        )\n",
      "                                                    )\n",
      "                                                ),\n",
      "                                                zeros(\n",
      "                                                    len(\n",
      "                                                        -arctan2(\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 2\n",
      "                                                            ],\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 1\n",
      "                                                            ],\n",
      "                                                        )\n",
      "                                                    )\n",
      "                                                ),\n",
      "                                            ],\n",
      "                                            [\n",
      "                                                zeros(\n",
      "                                                    len(\n",
      "                                                        -arctan2(\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 2\n",
      "                                                            ],\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 1\n",
      "                                                            ],\n",
      "                                                        )\n",
      "                                                    )\n",
      "                                                ),\n",
      "                                                cos(\n",
      "                                                    -arctan2(\n",
      "                                                        (p1 + p2 + p3)[:, 2],\n",
      "                                                        (p1 + p2 + p3)[:, 1],\n",
      "                                                    )\n",
      "                                                ),\n",
      "                                                -sin(\n",
      "                                                    -arctan2(\n",
      "                                                        (p1 + p2 + p3)[:, 2],\n",
      "                                                        (p1 + p2 + p3)[:, 1],\n",
      "                                                    )\n",
      "                                                ),\n",
      "                                                zeros(\n",
      "                                                    len(\n",
      "                                                        -arctan2(\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 2\n",
      "                                                            ],\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 1\n",
      "                                                            ],\n",
      "                                                        )\n",
      "                                                    )\n",
      "                                                ),\n",
      "                                            ],\n",
      "                                            [\n",
      "                                                zeros(\n",
      "                                                    len(\n",
      "                                                        -arctan2(\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 2\n",
      "                                                            ],\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 1\n",
      "                                                            ],\n",
      "                                                        )\n",
      "                                                    )\n",
      "                                                ),\n",
      "                                                sin(\n",
      "                                                    -arctan2(\n",
      "                                                        (p1 + p2 + p3)[:, 2],\n",
      "                                                        (p1 + p2 + p3)[:, 1],\n",
      "                                                    )\n",
      "                                                ),\n",
      "                                                cos(\n",
      "                                                    -arctan2(\n",
      "                                                        (p1 + p2 + p3)[:, 2],\n",
      "                                                        (p1 + p2 + p3)[:, 1],\n",
      "                                                    )\n",
      "                                                ),\n",
      "                                                zeros(\n",
      "                                                    len(\n",
      "                                                        -arctan2(\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 2\n",
      "                                                            ],\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 1\n",
      "                                                            ],\n",
      "                                                        )\n",
      "                                                    )\n",
      "                                                ),\n",
      "                                            ],\n",
      "                                            [\n",
      "                                                zeros(\n",
      "                                                    len(\n",
      "                                                        -arctan2(\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 2\n",
      "                                                            ],\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 1\n",
      "                                                            ],\n",
      "                                                        )\n",
      "                                                    )\n",
      "                                                ),\n",
      "                                                zeros(\n",
      "                                                    len(\n",
      "                                                        -arctan2(\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 2\n",
      "                                                            ],\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 1\n",
      "                                                            ],\n",
      "                                                        )\n",
      "                                                    )\n",
      "                                                ),\n",
      "                                                zeros(\n",
      "                                                    len(\n",
      "                                                        -arctan2(\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 2\n",
      "                                                            ],\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 1\n",
      "                                                            ],\n",
      "                                                        )\n",
      "                                                    )\n",
      "                                                ),\n",
      "                                                ones(\n",
      "                                                    len(\n",
      "                                                        -arctan2(\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 2\n",
      "                                                            ],\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 1\n",
      "                                                            ],\n",
      "                                                        )\n",
      "                                                    )\n",
      "                                                ),\n",
      "                                            ],\n",
      "                                        ]\n",
      "                                    ).transpose(2, 0, 1),\n",
      "                                    axes=(1, 2, 0),\n",
      "                                ),\n",
      "                                transpose(p2),\n",
      "                            )\n",
      "                        ),\n",
      "                    )\n",
      "                ),\n",
      "            )\n",
      "            + einsum(\n",
      "                \"ij...,j...\",\n",
      "                transpose(\n",
      "                    array(\n",
      "                        [\n",
      "                            [\n",
      "                                1\n",
      "                                / sqrt(\n",
      "                                    1\n",
      "                                    - (\n",
      "                                        sqrt(\n",
      "                                            sum(\n",
      "                                                (p1 + p2 + p3)[:, 1:] ** 2,\n",
      "                                                axis=1,\n",
      "                                            )\n",
      "                                        )\n",
      "                                        * (p1 + p2 + p3)[:, 0] ** (-1.0)\n",
      "                                    )\n",
      "                                    ** 2\n",
      "                                ),\n",
      "                                zeros(\n",
      "                                    len(\n",
      "                                        sqrt(\n",
      "                                            sum(\n",
      "                                                (p1 + p2 + p3)[:, 1:] ** 2,\n",
      "                                                axis=1,\n",
      "                                            )\n",
      "                                        )\n",
      "                                        * (p1 + p2 + p3)[:, 0] ** (-1.0)\n",
      "                                    )\n",
      "                                ),\n",
      "                                zeros(\n",
      "                                    len(\n",
      "                                        sqrt(\n",
      "                                            sum(\n",
      "                                                (p1 + p2 + p3)[:, 1:] ** 2,\n",
      "                                                axis=1,\n",
      "                                            )\n",
      "                                        )\n",
      "                                        * (p1 + p2 + p3)[:, 0] ** (-1.0)\n",
      "                                    )\n",
      "                                ),\n",
      "                                -1\n",
      "                                / sqrt(\n",
      "                                    1\n",
      "                                    - (\n",
      "                                        sqrt(\n",
      "                                            sum(\n",
      "                                                (p1 + p2 + p3)[:, 1:] ** 2,\n",
      "                                                axis=1,\n",
      "                                            )\n",
      "                                        )\n",
      "                                        * (p1 + p2 + p3)[:, 0] ** (-1.0)\n",
      "                                    )\n",
      "                                    ** 2\n",
      "                                )\n",
      "                                * sqrt(sum((p1 + p2 + p3)[:, 1:] ** 2, axis=1))\n",
      "                                * (p1 + p2 + p3)[:, 0] ** (-1.0),\n",
      "                            ],\n",
      "                            [\n",
      "                                zeros(\n",
      "                                    len(\n",
      "                                        sqrt(\n",
      "                                            sum(\n",
      "                                                (p1 + p2 + p3)[:, 1:] ** 2,\n",
      "                                                axis=1,\n",
      "                                            )\n",
      "                                        )\n",
      "                                        * (p1 + p2 + p3)[:, 0] ** (-1.0)\n",
      "                                    )\n",
      "                                ),\n",
      "                                ones(\n",
      "                                    len(\n",
      "                                        sqrt(\n",
      "                                            sum(\n",
      "                                                (p1 + p2 + p3)[:, 1:] ** 2,\n",
      "                                                axis=1,\n",
      "                                            )\n",
      "                                        )\n",
      "                                        * (p1 + p2 + p3)[:, 0] ** (-1.0)\n",
      "                                    )\n",
      "                                ),\n",
      "                                zeros(\n",
      "                                    len(\n",
      "                                        sqrt(\n",
      "                                            sum(\n",
      "                                                (p1 + p2 + p3)[:, 1:] ** 2,\n",
      "                                                axis=1,\n",
      "                                            )\n",
      "                                        )\n",
      "                                        * (p1 + p2 + p3)[:, 0] ** (-1.0)\n",
      "                                    )\n",
      "                                ),\n",
      "                                zeros(\n",
      "                                    len(\n",
      "                                        sqrt(\n",
      "                                            sum(\n",
      "                                                (p1 + p2 + p3)[:, 1:] ** 2,\n",
      "                                                axis=1,\n",
      "                                            )\n",
      "                                        )\n",
      "                                        * (p1 + p2 + p3)[:, 0] ** (-1.0)\n",
      "                                    )\n",
      "                                ),\n",
      "                            ],\n",
      "                            [\n",
      "                                zeros(\n",
      "                                    len(\n",
      "                                        sqrt(\n",
      "                                            sum(\n",
      "                                                (p1 + p2 + p3)[:, 1:] ** 2,\n",
      "                                                axis=1,\n",
      "                                            )\n",
      "                                        )\n",
      "                                        * (p1 + p2 + p3)[:, 0] ** (-1.0)\n",
      "                                    )\n",
      "                                ),\n",
      "                                zeros(\n",
      "                                    len(\n",
      "                                        sqrt(\n",
      "                                            sum(\n",
      "                                                (p1 + p2 + p3)[:, 1:] ** 2,\n",
      "                                                axis=1,\n",
      "                                            )\n",
      "                                        )\n",
      "                                        * (p1 + p2 + p3)[:, 0] ** (-1.0)\n",
      "                                    )\n",
      "                                ),\n",
      "                                ones(\n",
      "                                    len(\n",
      "                                        sqrt(\n",
      "                                            sum(\n",
      "                                                (p1 + p2 + p3)[:, 1:] ** 2,\n",
      "                                                axis=1,\n",
      "                                            )\n",
      "                                        )\n",
      "                                        * (p1 + p2 + p3)[:, 0] ** (-1.0)\n",
      "                                    )\n",
      "                                ),\n",
      "                                zeros(\n",
      "                                    len(\n",
      "                                        sqrt(\n",
      "                                            sum(\n",
      "                                                (p1 + p2 + p3)[:, 1:] ** 2,\n",
      "                                                axis=1,\n",
      "                                            )\n",
      "                                        )\n",
      "                                        * (p1 + p2 + p3)[:, 0] ** (-1.0)\n",
      "                                    )\n",
      "                                ),\n",
      "                            ],\n",
      "                            [\n",
      "                                -1\n",
      "                                / sqrt(\n",
      "                                    1\n",
      "                                    - (\n",
      "                                        sqrt(\n",
      "                                            sum(\n",
      "                                                (p1 + p2 + p3)[:, 1:] ** 2,\n",
      "                                                axis=1,\n",
      "                                            )\n",
      "                                        )\n",
      "                                        * (p1 + p2 + p3)[:, 0] ** (-1.0)\n",
      "                                    )\n",
      "                                    ** 2\n",
      "                                )\n",
      "                                * sqrt(sum((p1 + p2 + p3)[:, 1:] ** 2, axis=1))\n",
      "                                * (p1 + p2 + p3)[:, 0] ** (-1.0),\n",
      "                                zeros(\n",
      "                                    len(\n",
      "                                        sqrt(\n",
      "                                            sum(\n",
      "                                                (p1 + p2 + p3)[:, 1:] ** 2,\n",
      "                                                axis=1,\n",
      "                                            )\n",
      "                                        )\n",
      "                                        * (p1 + p2 + p3)[:, 0] ** (-1.0)\n",
      "                                    )\n",
      "                                ),\n",
      "                                zeros(\n",
      "                                    len(\n",
      "                                        sqrt(\n",
      "                                            sum(\n",
      "                                                (p1 + p2 + p3)[:, 1:] ** 2,\n",
      "                                                axis=1,\n",
      "                                            )\n",
      "                                        )\n",
      "                                        * (p1 + p2 + p3)[:, 0] ** (-1.0)\n",
      "                                    )\n",
      "                                ),\n",
      "                                1\n",
      "                                / sqrt(\n",
      "                                    1\n",
      "                                    - (\n",
      "                                        sqrt(\n",
      "                                            sum(\n",
      "                                                (p1 + p2 + p3)[:, 1:] ** 2,\n",
      "                                                axis=1,\n",
      "                                            )\n",
      "                                        )\n",
      "                                        * (p1 + p2 + p3)[:, 0] ** (-1.0)\n",
      "                                    )\n",
      "                                    ** 2\n",
      "                                ),\n",
      "                            ],\n",
      "                        ]\n",
      "                    ).transpose(2, 0, 1),\n",
      "                    axes=(1, 2, 0),\n",
      "                ),\n",
      "                transpose(\n",
      "                    einsum(\n",
      "                        \"ij...,j...\",\n",
      "                        transpose(\n",
      "                            array(\n",
      "                                [\n",
      "                                    [\n",
      "                                        ones(\n",
      "                                            len(\n",
      "                                                -arccos(\n",
      "                                                    (p1 + p2 + p3)[:, 3]\n",
      "                                                    * sum(\n",
      "                                                        (p1 + p2 + p3)[:, 1:]\n",
      "                                                        ** 2,\n",
      "                                                        axis=1,\n",
      "                                                    )\n",
      "                                                    ** (-1 / 2)\n",
      "                                                )\n",
      "                                            )\n",
      "                                        ),\n",
      "                                        zeros(\n",
      "                                            len(\n",
      "                                                -arccos(\n",
      "                                                    (p1 + p2 + p3)[:, 3]\n",
      "                                                    * sum(\n",
      "                                                        (p1 + p2 + p3)[:, 1:]\n",
      "                                                        ** 2,\n",
      "                                                        axis=1,\n",
      "                                                    )\n",
      "                                                    ** (-1 / 2)\n",
      "                                                )\n",
      "                                            )\n",
      "                                        ),\n",
      "                                        zeros(\n",
      "                                            len(\n",
      "                                                -arccos(\n",
      "                                                    (p1 + p2 + p3)[:, 3]\n",
      "                                                    * sum(\n",
      "                                                        (p1 + p2 + p3)[:, 1:]\n",
      "                                                        ** 2,\n",
      "                                                        axis=1,\n",
      "                                                    )\n",
      "                                                    ** (-1 / 2)\n",
      "                                                )\n",
      "                                            )\n",
      "                                        ),\n",
      "                                        zeros(\n",
      "                                            len(\n",
      "                                                -arccos(\n",
      "                                                    (p1 + p2 + p3)[:, 3]\n",
      "                                                    * sum(\n",
      "                                                        (p1 + p2 + p3)[:, 1:]\n",
      "                                                        ** 2,\n",
      "                                                        axis=1,\n",
      "                                                    )\n",
      "                                                    ** (-1 / 2)\n",
      "                                                )\n",
      "                                            )\n",
      "                                        ),\n",
      "                                    ],\n",
      "                                    [\n",
      "                                        zeros(\n",
      "                                            len(\n",
      "                                                -arccos(\n",
      "                                                    (p1 + p2 + p3)[:, 3]\n",
      "                                                    * sum(\n",
      "                                                        (p1 + p2 + p3)[:, 1:]\n",
      "                                                        ** 2,\n",
      "                                                        axis=1,\n",
      "                                                    )\n",
      "                                                    ** (-1 / 2)\n",
      "                                                )\n",
      "                                            )\n",
      "                                        ),\n",
      "                                        cos(\n",
      "                                            -arccos(\n",
      "                                                (p1 + p2 + p3)[:, 3]\n",
      "                                                * sum(\n",
      "                                                    (p1 + p2 + p3)[:, 1:] ** 2,\n",
      "                                                    axis=1,\n",
      "                                                )\n",
      "                                                ** (-1 / 2)\n",
      "                                            )\n",
      "                                        ),\n",
      "                                        zeros(\n",
      "                                            len(\n",
      "                                                -arccos(\n",
      "                                                    (p1 + p2 + p3)[:, 3]\n",
      "                                                    * sum(\n",
      "                                                        (p1 + p2 + p3)[:, 1:]\n",
      "                                                        ** 2,\n",
      "                                                        axis=1,\n",
      "                                                    )\n",
      "                                                    ** (-1 / 2)\n",
      "                                                )\n",
      "                                            )\n",
      "                                        ),\n",
      "                                        sin(\n",
      "                                            -arccos(\n",
      "                                                (p1 + p2 + p3)[:, 3]\n",
      "                                                * sum(\n",
      "                                                    (p1 + p2 + p3)[:, 1:] ** 2,\n",
      "                                                    axis=1,\n",
      "                                                )\n",
      "                                                ** (-1 / 2)\n",
      "                                            )\n",
      "                                        ),\n",
      "                                    ],\n",
      "                                    [\n",
      "                                        zeros(\n",
      "                                            len(\n",
      "                                                -arccos(\n",
      "                                                    (p1 + p2 + p3)[:, 3]\n",
      "                                                    * sum(\n",
      "                                                        (p1 + p2 + p3)[:, 1:]\n",
      "                                                        ** 2,\n",
      "                                                        axis=1,\n",
      "                                                    )\n",
      "                                                    ** (-1 / 2)\n",
      "                                                )\n",
      "                                            )\n",
      "                                        ),\n",
      "                                        zeros(\n",
      "                                            len(\n",
      "                                                -arccos(\n",
      "                                                    (p1 + p2 + p3)[:, 3]\n",
      "                                                    * sum(\n",
      "                                                        (p1 + p2 + p3)[:, 1:]\n",
      "                                                        ** 2,\n",
      "                                                        axis=1,\n",
      "                                                    )\n",
      "                                                    ** (-1 / 2)\n",
      "                                                )\n",
      "                                            )\n",
      "                                        ),\n",
      "                                        ones(\n",
      "                                            len(\n",
      "                                                -arccos(\n",
      "                                                    (p1 + p2 + p3)[:, 3]\n",
      "                                                    * sum(\n",
      "                                                        (p1 + p2 + p3)[:, 1:]\n",
      "                                                        ** 2,\n",
      "                                                        axis=1,\n",
      "                                                    )\n",
      "                                                    ** (-1 / 2)\n",
      "                                                )\n",
      "                                            )\n",
      "                                        ),\n",
      "                                        zeros(\n",
      "                                            len(\n",
      "                                                -arccos(\n",
      "                                                    (p1 + p2 + p3)[:, 3]\n",
      "                                                    * sum(\n",
      "                                                        (p1 + p2 + p3)[:, 1:]\n",
      "                                                        ** 2,\n",
      "                                                        axis=1,\n",
      "                                                    )\n",
      "                                                    ** (-1 / 2)\n",
      "                                                )\n",
      "                                            )\n",
      "                                        ),\n",
      "                                    ],\n",
      "                                    [\n",
      "                                        zeros(\n",
      "                                            len(\n",
      "                                                -arccos(\n",
      "                                                    (p1 + p2 + p3)[:, 3]\n",
      "                                                    * sum(\n",
      "                                                        (p1 + p2 + p3)[:, 1:]\n",
      "                                                        ** 2,\n",
      "                                                        axis=1,\n",
      "                                                    )\n",
      "                                                    ** (-1 / 2)\n",
      "                                                )\n",
      "                                            )\n",
      "                                        ),\n",
      "                                        -sin(\n",
      "                                            -arccos(\n",
      "                                                (p1 + p2 + p3)[:, 3]\n",
      "                                                * sum(\n",
      "                                                    (p1 + p2 + p3)[:, 1:] ** 2,\n",
      "                                                    axis=1,\n",
      "                                                )\n",
      "                                                ** (-1 / 2)\n",
      "                                            )\n",
      "                                        ),\n",
      "                                        zeros(\n",
      "                                            len(\n",
      "                                                -arccos(\n",
      "                                                    (p1 + p2 + p3)[:, 3]\n",
      "                                                    * sum(\n",
      "                                                        (p1 + p2 + p3)[:, 1:]\n",
      "                                                        ** 2,\n",
      "                                                        axis=1,\n",
      "                                                    )\n",
      "                                                    ** (-1 / 2)\n",
      "                                                )\n",
      "                                            )\n",
      "                                        ),\n",
      "                                        cos(\n",
      "                                            -arccos(\n",
      "                                                (p1 + p2 + p3)[:, 3]\n",
      "                                                * sum(\n",
      "                                                    (p1 + p2 + p3)[:, 1:] ** 2,\n",
      "                                                    axis=1,\n",
      "                                                )\n",
      "                                                ** (-1 / 2)\n",
      "                                            )\n",
      "                                        ),\n",
      "                                    ],\n",
      "                                ]\n",
      "                            ).transpose(2, 0, 1),\n",
      "                            axes=(1, 2, 0),\n",
      "                        ),\n",
      "                        transpose(\n",
      "                            einsum(\n",
      "                                \"ij...,j...\",\n",
      "                                transpose(\n",
      "                                    array(\n",
      "                                        [\n",
      "                                            [\n",
      "                                                ones(\n",
      "                                                    len(\n",
      "                                                        -arctan2(\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 2\n",
      "                                                            ],\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 1\n",
      "                                                            ],\n",
      "                                                        )\n",
      "                                                    )\n",
      "                                                ),\n",
      "                                                zeros(\n",
      "                                                    len(\n",
      "                                                        -arctan2(\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 2\n",
      "                                                            ],\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 1\n",
      "                                                            ],\n",
      "                                                        )\n",
      "                                                    )\n",
      "                                                ),\n",
      "                                                zeros(\n",
      "                                                    len(\n",
      "                                                        -arctan2(\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 2\n",
      "                                                            ],\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 1\n",
      "                                                            ],\n",
      "                                                        )\n",
      "                                                    )\n",
      "                                                ),\n",
      "                                                zeros(\n",
      "                                                    len(\n",
      "                                                        -arctan2(\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 2\n",
      "                                                            ],\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 1\n",
      "                                                            ],\n",
      "                                                        )\n",
      "                                                    )\n",
      "                                                ),\n",
      "                                            ],\n",
      "                                            [\n",
      "                                                zeros(\n",
      "                                                    len(\n",
      "                                                        -arctan2(\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 2\n",
      "                                                            ],\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 1\n",
      "                                                            ],\n",
      "                                                        )\n",
      "                                                    )\n",
      "                                                ),\n",
      "                                                cos(\n",
      "                                                    -arctan2(\n",
      "                                                        (p1 + p2 + p3)[:, 2],\n",
      "                                                        (p1 + p2 + p3)[:, 1],\n",
      "                                                    )\n",
      "                                                ),\n",
      "                                                -sin(\n",
      "                                                    -arctan2(\n",
      "                                                        (p1 + p2 + p3)[:, 2],\n",
      "                                                        (p1 + p2 + p3)[:, 1],\n",
      "                                                    )\n",
      "                                                ),\n",
      "                                                zeros(\n",
      "                                                    len(\n",
      "                                                        -arctan2(\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 2\n",
      "                                                            ],\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 1\n",
      "                                                            ],\n",
      "                                                        )\n",
      "                                                    )\n",
      "                                                ),\n",
      "                                            ],\n",
      "                                            [\n",
      "                                                zeros(\n",
      "                                                    len(\n",
      "                                                        -arctan2(\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 2\n",
      "                                                            ],\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 1\n",
      "                                                            ],\n",
      "                                                        )\n",
      "                                                    )\n",
      "                                                ),\n",
      "                                                sin(\n",
      "                                                    -arctan2(\n",
      "                                                        (p1 + p2 + p3)[:, 2],\n",
      "                                                        (p1 + p2 + p3)[:, 1],\n",
      "                                                    )\n",
      "                                                ),\n",
      "                                                cos(\n",
      "                                                    -arctan2(\n",
      "                                                        (p1 + p2 + p3)[:, 2],\n",
      "                                                        (p1 + p2 + p3)[:, 1],\n",
      "                                                    )\n",
      "                                                ),\n",
      "                                                zeros(\n",
      "                                                    len(\n",
      "                                                        -arctan2(\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 2\n",
      "                                                            ],\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 1\n",
      "                                                            ],\n",
      "                                                        )\n",
      "                                                    )\n",
      "                                                ),\n",
      "                                            ],\n",
      "                                            [\n",
      "                                                zeros(\n",
      "                                                    len(\n",
      "                                                        -arctan2(\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 2\n",
      "                                                            ],\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 1\n",
      "                                                            ],\n",
      "                                                        )\n",
      "                                                    )\n",
      "                                                ),\n",
      "                                                zeros(\n",
      "                                                    len(\n",
      "                                                        -arctan2(\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 2\n",
      "                                                            ],\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 1\n",
      "                                                            ],\n",
      "                                                        )\n",
      "                                                    )\n",
      "                                                ),\n",
      "                                                zeros(\n",
      "                                                    len(\n",
      "                                                        -arctan2(\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 2\n",
      "                                                            ],\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 1\n",
      "                                                            ],\n",
      "                                                        )\n",
      "                                                    )\n",
      "                                                ),\n",
      "                                                ones(\n",
      "                                                    len(\n",
      "                                                        -arctan2(\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 2\n",
      "                                                            ],\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 1\n",
      "                                                            ],\n",
      "                                                        )\n",
      "                                                    )\n",
      "                                                ),\n",
      "                                            ],\n",
      "                                        ]\n",
      "                                    ).transpose(2, 0, 1),\n",
      "                                    axes=(1, 2, 0),\n",
      "                                ),\n",
      "                                transpose(p3),\n",
      "                            )\n",
      "                        ),\n",
      "                    )\n",
      "                ),\n",
      "            )\n",
      "        )[:, 2],\n",
      "        (\n",
      "            einsum(\n",
      "                \"ij...,j...\",\n",
      "                transpose(\n",
      "                    array(\n",
      "                        [\n",
      "                            [\n",
      "                                1\n",
      "                                / sqrt(\n",
      "                                    1\n",
      "                                    - (\n",
      "                                        sqrt(\n",
      "                                            sum(\n",
      "                                                (p1 + p2 + p3)[:, 1:] ** 2,\n",
      "                                                axis=1,\n",
      "                                            )\n",
      "                                        )\n",
      "                                        * (p1 + p2 + p3)[:, 0] ** (-1.0)\n",
      "                                    )\n",
      "                                    ** 2\n",
      "                                ),\n",
      "                                zeros(\n",
      "                                    len(\n",
      "                                        sqrt(\n",
      "                                            sum(\n",
      "                                                (p1 + p2 + p3)[:, 1:] ** 2,\n",
      "                                                axis=1,\n",
      "                                            )\n",
      "                                        )\n",
      "                                        * (p1 + p2 + p3)[:, 0] ** (-1.0)\n",
      "                                    )\n",
      "                                ),\n",
      "                                zeros(\n",
      "                                    len(\n",
      "                                        sqrt(\n",
      "                                            sum(\n",
      "                                                (p1 + p2 + p3)[:, 1:] ** 2,\n",
      "                                                axis=1,\n",
      "                                            )\n",
      "                                        )\n",
      "                                        * (p1 + p2 + p3)[:, 0] ** (-1.0)\n",
      "                                    )\n",
      "                                ),\n",
      "                                -1\n",
      "                                / sqrt(\n",
      "                                    1\n",
      "                                    - (\n",
      "                                        sqrt(\n",
      "                                            sum(\n",
      "                                                (p1 + p2 + p3)[:, 1:] ** 2,\n",
      "                                                axis=1,\n",
      "                                            )\n",
      "                                        )\n",
      "                                        * (p1 + p2 + p3)[:, 0] ** (-1.0)\n",
      "                                    )\n",
      "                                    ** 2\n",
      "                                )\n",
      "                                * sqrt(sum((p1 + p2 + p3)[:, 1:] ** 2, axis=1))\n",
      "                                * (p1 + p2 + p3)[:, 0] ** (-1.0),\n",
      "                            ],\n",
      "                            [\n",
      "                                zeros(\n",
      "                                    len(\n",
      "                                        sqrt(\n",
      "                                            sum(\n",
      "                                                (p1 + p2 + p3)[:, 1:] ** 2,\n",
      "                                                axis=1,\n",
      "                                            )\n",
      "                                        )\n",
      "                                        * (p1 + p2 + p3)[:, 0] ** (-1.0)\n",
      "                                    )\n",
      "                                ),\n",
      "                                ones(\n",
      "                                    len(\n",
      "                                        sqrt(\n",
      "                                            sum(\n",
      "                                                (p1 + p2 + p3)[:, 1:] ** 2,\n",
      "                                                axis=1,\n",
      "                                            )\n",
      "                                        )\n",
      "                                        * (p1 + p2 + p3)[:, 0] ** (-1.0)\n",
      "                                    )\n",
      "                                ),\n",
      "                                zeros(\n",
      "                                    len(\n",
      "                                        sqrt(\n",
      "                                            sum(\n",
      "                                                (p1 + p2 + p3)[:, 1:] ** 2,\n",
      "                                                axis=1,\n",
      "                                            )\n",
      "                                        )\n",
      "                                        * (p1 + p2 + p3)[:, 0] ** (-1.0)\n",
      "                                    )\n",
      "                                ),\n",
      "                                zeros(\n",
      "                                    len(\n",
      "                                        sqrt(\n",
      "                                            sum(\n",
      "                                                (p1 + p2 + p3)[:, 1:] ** 2,\n",
      "                                                axis=1,\n",
      "                                            )\n",
      "                                        )\n",
      "                                        * (p1 + p2 + p3)[:, 0] ** (-1.0)\n",
      "                                    )\n",
      "                                ),\n",
      "                            ],\n",
      "                            [\n",
      "                                zeros(\n",
      "                                    len(\n",
      "                                        sqrt(\n",
      "                                            sum(\n",
      "                                                (p1 + p2 + p3)[:, 1:] ** 2,\n",
      "                                                axis=1,\n",
      "                                            )\n",
      "                                        )\n",
      "                                        * (p1 + p2 + p3)[:, 0] ** (-1.0)\n",
      "                                    )\n",
      "                                ),\n",
      "                                zeros(\n",
      "                                    len(\n",
      "                                        sqrt(\n",
      "                                            sum(\n",
      "                                                (p1 + p2 + p3)[:, 1:] ** 2,\n",
      "                                                axis=1,\n",
      "                                            )\n",
      "                                        )\n",
      "                                        * (p1 + p2 + p3)[:, 0] ** (-1.0)\n",
      "                                    )\n",
      "                                ),\n",
      "                                ones(\n",
      "                                    len(\n",
      "                                        sqrt(\n",
      "                                            sum(\n",
      "                                                (p1 + p2 + p3)[:, 1:] ** 2,\n",
      "                                                axis=1,\n",
      "                                            )\n",
      "                                        )\n",
      "                                        * (p1 + p2 + p3)[:, 0] ** (-1.0)\n",
      "                                    )\n",
      "                                ),\n",
      "                                zeros(\n",
      "                                    len(\n",
      "                                        sqrt(\n",
      "                                            sum(\n",
      "                                                (p1 + p2 + p3)[:, 1:] ** 2,\n",
      "                                                axis=1,\n",
      "                                            )\n",
      "                                        )\n",
      "                                        * (p1 + p2 + p3)[:, 0] ** (-1.0)\n",
      "                                    )\n",
      "                                ),\n",
      "                            ],\n",
      "                            [\n",
      "                                -1\n",
      "                                / sqrt(\n",
      "                                    1\n",
      "                                    - (\n",
      "                                        sqrt(\n",
      "                                            sum(\n",
      "                                                (p1 + p2 + p3)[:, 1:] ** 2,\n",
      "                                                axis=1,\n",
      "                                            )\n",
      "                                        )\n",
      "                                        * (p1 + p2 + p3)[:, 0] ** (-1.0)\n",
      "                                    )\n",
      "                                    ** 2\n",
      "                                )\n",
      "                                * sqrt(sum((p1 + p2 + p3)[:, 1:] ** 2, axis=1))\n",
      "                                * (p1 + p2 + p3)[:, 0] ** (-1.0),\n",
      "                                zeros(\n",
      "                                    len(\n",
      "                                        sqrt(\n",
      "                                            sum(\n",
      "                                                (p1 + p2 + p3)[:, 1:] ** 2,\n",
      "                                                axis=1,\n",
      "                                            )\n",
      "                                        )\n",
      "                                        * (p1 + p2 + p3)[:, 0] ** (-1.0)\n",
      "                                    )\n",
      "                                ),\n",
      "                                zeros(\n",
      "                                    len(\n",
      "                                        sqrt(\n",
      "                                            sum(\n",
      "                                                (p1 + p2 + p3)[:, 1:] ** 2,\n",
      "                                                axis=1,\n",
      "                                            )\n",
      "                                        )\n",
      "                                        * (p1 + p2 + p3)[:, 0] ** (-1.0)\n",
      "                                    )\n",
      "                                ),\n",
      "                                1\n",
      "                                / sqrt(\n",
      "                                    1\n",
      "                                    - (\n",
      "                                        sqrt(\n",
      "                                            sum(\n",
      "                                                (p1 + p2 + p3)[:, 1:] ** 2,\n",
      "                                                axis=1,\n",
      "                                            )\n",
      "                                        )\n",
      "                                        * (p1 + p2 + p3)[:, 0] ** (-1.0)\n",
      "                                    )\n",
      "                                    ** 2\n",
      "                                ),\n",
      "                            ],\n",
      "                        ]\n",
      "                    ).transpose(2, 0, 1),\n",
      "                    axes=(1, 2, 0),\n",
      "                ),\n",
      "                transpose(\n",
      "                    einsum(\n",
      "                        \"ij...,j...\",\n",
      "                        transpose(\n",
      "                            array(\n",
      "                                [\n",
      "                                    [\n",
      "                                        ones(\n",
      "                                            len(\n",
      "                                                -arccos(\n",
      "                                                    (p1 + p2 + p3)[:, 3]\n",
      "                                                    * sum(\n",
      "                                                        (p1 + p2 + p3)[:, 1:]\n",
      "                                                        ** 2,\n",
      "                                                        axis=1,\n",
      "                                                    )\n",
      "                                                    ** (-1 / 2)\n",
      "                                                )\n",
      "                                            )\n",
      "                                        ),\n",
      "                                        zeros(\n",
      "                                            len(\n",
      "                                                -arccos(\n",
      "                                                    (p1 + p2 + p3)[:, 3]\n",
      "                                                    * sum(\n",
      "                                                        (p1 + p2 + p3)[:, 1:]\n",
      "                                                        ** 2,\n",
      "                                                        axis=1,\n",
      "                                                    )\n",
      "                                                    ** (-1 / 2)\n",
      "                                                )\n",
      "                                            )\n",
      "                                        ),\n",
      "                                        zeros(\n",
      "                                            len(\n",
      "                                                -arccos(\n",
      "                                                    (p1 + p2 + p3)[:, 3]\n",
      "                                                    * sum(\n",
      "                                                        (p1 + p2 + p3)[:, 1:]\n",
      "                                                        ** 2,\n",
      "                                                        axis=1,\n",
      "                                                    )\n",
      "                                                    ** (-1 / 2)\n",
      "                                                )\n",
      "                                            )\n",
      "                                        ),\n",
      "                                        zeros(\n",
      "                                            len(\n",
      "                                                -arccos(\n",
      "                                                    (p1 + p2 + p3)[:, 3]\n",
      "                                                    * sum(\n",
      "                                                        (p1 + p2 + p3)[:, 1:]\n",
      "                                                        ** 2,\n",
      "                                                        axis=1,\n",
      "                                                    )\n",
      "                                                    ** (-1 / 2)\n",
      "                                                )\n",
      "                                            )\n",
      "                                        ),\n",
      "                                    ],\n",
      "                                    [\n",
      "                                        zeros(\n",
      "                                            len(\n",
      "                                                -arccos(\n",
      "                                                    (p1 + p2 + p3)[:, 3]\n",
      "                                                    * sum(\n",
      "                                                        (p1 + p2 + p3)[:, 1:]\n",
      "                                                        ** 2,\n",
      "                                                        axis=1,\n",
      "                                                    )\n",
      "                                                    ** (-1 / 2)\n",
      "                                                )\n",
      "                                            )\n",
      "                                        ),\n",
      "                                        cos(\n",
      "                                            -arccos(\n",
      "                                                (p1 + p2 + p3)[:, 3]\n",
      "                                                * sum(\n",
      "                                                    (p1 + p2 + p3)[:, 1:] ** 2,\n",
      "                                                    axis=1,\n",
      "                                                )\n",
      "                                                ** (-1 / 2)\n",
      "                                            )\n",
      "                                        ),\n",
      "                                        zeros(\n",
      "                                            len(\n",
      "                                                -arccos(\n",
      "                                                    (p1 + p2 + p3)[:, 3]\n",
      "                                                    * sum(\n",
      "                                                        (p1 + p2 + p3)[:, 1:]\n",
      "                                                        ** 2,\n",
      "                                                        axis=1,\n",
      "                                                    )\n",
      "                                                    ** (-1 / 2)\n",
      "                                                )\n",
      "                                            )\n",
      "                                        ),\n",
      "                                        sin(\n",
      "                                            -arccos(\n",
      "                                                (p1 + p2 + p3)[:, 3]\n",
      "                                                * sum(\n",
      "                                                    (p1 + p2 + p3)[:, 1:] ** 2,\n",
      "                                                    axis=1,\n",
      "                                                )\n",
      "                                                ** (-1 / 2)\n",
      "                                            )\n",
      "                                        ),\n",
      "                                    ],\n",
      "                                    [\n",
      "                                        zeros(\n",
      "                                            len(\n",
      "                                                -arccos(\n",
      "                                                    (p1 + p2 + p3)[:, 3]\n",
      "                                                    * sum(\n",
      "                                                        (p1 + p2 + p3)[:, 1:]\n",
      "                                                        ** 2,\n",
      "                                                        axis=1,\n",
      "                                                    )\n",
      "                                                    ** (-1 / 2)\n",
      "                                                )\n",
      "                                            )\n",
      "                                        ),\n",
      "                                        zeros(\n",
      "                                            len(\n",
      "                                                -arccos(\n",
      "                                                    (p1 + p2 + p3)[:, 3]\n",
      "                                                    * sum(\n",
      "                                                        (p1 + p2 + p3)[:, 1:]\n",
      "                                                        ** 2,\n",
      "                                                        axis=1,\n",
      "                                                    )\n",
      "                                                    ** (-1 / 2)\n",
      "                                                )\n",
      "                                            )\n",
      "                                        ),\n",
      "                                        ones(\n",
      "                                            len(\n",
      "                                                -arccos(\n",
      "                                                    (p1 + p2 + p3)[:, 3]\n",
      "                                                    * sum(\n",
      "                                                        (p1 + p2 + p3)[:, 1:]\n",
      "                                                        ** 2,\n",
      "                                                        axis=1,\n",
      "                                                    )\n",
      "                                                    ** (-1 / 2)\n",
      "                                                )\n",
      "                                            )\n",
      "                                        ),\n",
      "                                        zeros(\n",
      "                                            len(\n",
      "                                                -arccos(\n",
      "                                                    (p1 + p2 + p3)[:, 3]\n",
      "                                                    * sum(\n",
      "                                                        (p1 + p2 + p3)[:, 1:]\n",
      "                                                        ** 2,\n",
      "                                                        axis=1,\n",
      "                                                    )\n",
      "                                                    ** (-1 / 2)\n",
      "                                                )\n",
      "                                            )\n",
      "                                        ),\n",
      "                                    ],\n",
      "                                    [\n",
      "                                        zeros(\n",
      "                                            len(\n",
      "                                                -arccos(\n",
      "                                                    (p1 + p2 + p3)[:, 3]\n",
      "                                                    * sum(\n",
      "                                                        (p1 + p2 + p3)[:, 1:]\n",
      "                                                        ** 2,\n",
      "                                                        axis=1,\n",
      "                                                    )\n",
      "                                                    ** (-1 / 2)\n",
      "                                                )\n",
      "                                            )\n",
      "                                        ),\n",
      "                                        -sin(\n",
      "                                            -arccos(\n",
      "                                                (p1 + p2 + p3)[:, 3]\n",
      "                                                * sum(\n",
      "                                                    (p1 + p2 + p3)[:, 1:] ** 2,\n",
      "                                                    axis=1,\n",
      "                                                )\n",
      "                                                ** (-1 / 2)\n",
      "                                            )\n",
      "                                        ),\n",
      "                                        zeros(\n",
      "                                            len(\n",
      "                                                -arccos(\n",
      "                                                    (p1 + p2 + p3)[:, 3]\n",
      "                                                    * sum(\n",
      "                                                        (p1 + p2 + p3)[:, 1:]\n",
      "                                                        ** 2,\n",
      "                                                        axis=1,\n",
      "                                                    )\n",
      "                                                    ** (-1 / 2)\n",
      "                                                )\n",
      "                                            )\n",
      "                                        ),\n",
      "                                        cos(\n",
      "                                            -arccos(\n",
      "                                                (p1 + p2 + p3)[:, 3]\n",
      "                                                * sum(\n",
      "                                                    (p1 + p2 + p3)[:, 1:] ** 2,\n",
      "                                                    axis=1,\n",
      "                                                )\n",
      "                                                ** (-1 / 2)\n",
      "                                            )\n",
      "                                        ),\n",
      "                                    ],\n",
      "                                ]\n",
      "                            ).transpose(2, 0, 1),\n",
      "                            axes=(1, 2, 0),\n",
      "                        ),\n",
      "                        transpose(\n",
      "                            einsum(\n",
      "                                \"ij...,j...\",\n",
      "                                transpose(\n",
      "                                    array(\n",
      "                                        [\n",
      "                                            [\n",
      "                                                ones(\n",
      "                                                    len(\n",
      "                                                        -arctan2(\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 2\n",
      "                                                            ],\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 1\n",
      "                                                            ],\n",
      "                                                        )\n",
      "                                                    )\n",
      "                                                ),\n",
      "                                                zeros(\n",
      "                                                    len(\n",
      "                                                        -arctan2(\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 2\n",
      "                                                            ],\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 1\n",
      "                                                            ],\n",
      "                                                        )\n",
      "                                                    )\n",
      "                                                ),\n",
      "                                                zeros(\n",
      "                                                    len(\n",
      "                                                        -arctan2(\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 2\n",
      "                                                            ],\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 1\n",
      "                                                            ],\n",
      "                                                        )\n",
      "                                                    )\n",
      "                                                ),\n",
      "                                                zeros(\n",
      "                                                    len(\n",
      "                                                        -arctan2(\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 2\n",
      "                                                            ],\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 1\n",
      "                                                            ],\n",
      "                                                        )\n",
      "                                                    )\n",
      "                                                ),\n",
      "                                            ],\n",
      "                                            [\n",
      "                                                zeros(\n",
      "                                                    len(\n",
      "                                                        -arctan2(\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 2\n",
      "                                                            ],\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 1\n",
      "                                                            ],\n",
      "                                                        )\n",
      "                                                    )\n",
      "                                                ),\n",
      "                                                cos(\n",
      "                                                    -arctan2(\n",
      "                                                        (p1 + p2 + p3)[:, 2],\n",
      "                                                        (p1 + p2 + p3)[:, 1],\n",
      "                                                    )\n",
      "                                                ),\n",
      "                                                -sin(\n",
      "                                                    -arctan2(\n",
      "                                                        (p1 + p2 + p3)[:, 2],\n",
      "                                                        (p1 + p2 + p3)[:, 1],\n",
      "                                                    )\n",
      "                                                ),\n",
      "                                                zeros(\n",
      "                                                    len(\n",
      "                                                        -arctan2(\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 2\n",
      "                                                            ],\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 1\n",
      "                                                            ],\n",
      "                                                        )\n",
      "                                                    )\n",
      "                                                ),\n",
      "                                            ],\n",
      "                                            [\n",
      "                                                zeros(\n",
      "                                                    len(\n",
      "                                                        -arctan2(\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 2\n",
      "                                                            ],\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 1\n",
      "                                                            ],\n",
      "                                                        )\n",
      "                                                    )\n",
      "                                                ),\n",
      "                                                sin(\n",
      "                                                    -arctan2(\n",
      "                                                        (p1 + p2 + p3)[:, 2],\n",
      "                                                        (p1 + p2 + p3)[:, 1],\n",
      "                                                    )\n",
      "                                                ),\n",
      "                                                cos(\n",
      "                                                    -arctan2(\n",
      "                                                        (p1 + p2 + p3)[:, 2],\n",
      "                                                        (p1 + p2 + p3)[:, 1],\n",
      "                                                    )\n",
      "                                                ),\n",
      "                                                zeros(\n",
      "                                                    len(\n",
      "                                                        -arctan2(\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 2\n",
      "                                                            ],\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 1\n",
      "                                                            ],\n",
      "                                                        )\n",
      "                                                    )\n",
      "                                                ),\n",
      "                                            ],\n",
      "                                            [\n",
      "                                                zeros(\n",
      "                                                    len(\n",
      "                                                        -arctan2(\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 2\n",
      "                                                            ],\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 1\n",
      "                                                            ],\n",
      "                                                        )\n",
      "                                                    )\n",
      "                                                ),\n",
      "                                                zeros(\n",
      "                                                    len(\n",
      "                                                        -arctan2(\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 2\n",
      "                                                            ],\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 1\n",
      "                                                            ],\n",
      "                                                        )\n",
      "                                                    )\n",
      "                                                ),\n",
      "                                                zeros(\n",
      "                                                    len(\n",
      "                                                        -arctan2(\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 2\n",
      "                                                            ],\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 1\n",
      "                                                            ],\n",
      "                                                        )\n",
      "                                                    )\n",
      "                                                ),\n",
      "                                                ones(\n",
      "                                                    len(\n",
      "                                                        -arctan2(\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 2\n",
      "                                                            ],\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 1\n",
      "                                                            ],\n",
      "                                                        )\n",
      "                                                    )\n",
      "                                                ),\n",
      "                                            ],\n",
      "                                        ]\n",
      "                                    ).transpose(2, 0, 1),\n",
      "                                    axes=(1, 2, 0),\n",
      "                                ),\n",
      "                                transpose(p2),\n",
      "                            )\n",
      "                        ),\n",
      "                    )\n",
      "                ),\n",
      "            )\n",
      "            + einsum(\n",
      "                \"ij...,j...\",\n",
      "                transpose(\n",
      "                    array(\n",
      "                        [\n",
      "                            [\n",
      "                                1\n",
      "                                / sqrt(\n",
      "                                    1\n",
      "                                    - (\n",
      "                                        sqrt(\n",
      "                                            sum(\n",
      "                                                (p1 + p2 + p3)[:, 1:] ** 2,\n",
      "                                                axis=1,\n",
      "                                            )\n",
      "                                        )\n",
      "                                        * (p1 + p2 + p3)[:, 0] ** (-1.0)\n",
      "                                    )\n",
      "                                    ** 2\n",
      "                                ),\n",
      "                                zeros(\n",
      "                                    len(\n",
      "                                        sqrt(\n",
      "                                            sum(\n",
      "                                                (p1 + p2 + p3)[:, 1:] ** 2,\n",
      "                                                axis=1,\n",
      "                                            )\n",
      "                                        )\n",
      "                                        * (p1 + p2 + p3)[:, 0] ** (-1.0)\n",
      "                                    )\n",
      "                                ),\n",
      "                                zeros(\n",
      "                                    len(\n",
      "                                        sqrt(\n",
      "                                            sum(\n",
      "                                                (p1 + p2 + p3)[:, 1:] ** 2,\n",
      "                                                axis=1,\n",
      "                                            )\n",
      "                                        )\n",
      "                                        * (p1 + p2 + p3)[:, 0] ** (-1.0)\n",
      "                                    )\n",
      "                                ),\n",
      "                                -1\n",
      "                                / sqrt(\n",
      "                                    1\n",
      "                                    - (\n",
      "                                        sqrt(\n",
      "                                            sum(\n",
      "                                                (p1 + p2 + p3)[:, 1:] ** 2,\n",
      "                                                axis=1,\n",
      "                                            )\n",
      "                                        )\n",
      "                                        * (p1 + p2 + p3)[:, 0] ** (-1.0)\n",
      "                                    )\n",
      "                                    ** 2\n",
      "                                )\n",
      "                                * sqrt(sum((p1 + p2 + p3)[:, 1:] ** 2, axis=1))\n",
      "                                * (p1 + p2 + p3)[:, 0] ** (-1.0),\n",
      "                            ],\n",
      "                            [\n",
      "                                zeros(\n",
      "                                    len(\n",
      "                                        sqrt(\n",
      "                                            sum(\n",
      "                                                (p1 + p2 + p3)[:, 1:] ** 2,\n",
      "                                                axis=1,\n",
      "                                            )\n",
      "                                        )\n",
      "                                        * (p1 + p2 + p3)[:, 0] ** (-1.0)\n",
      "                                    )\n",
      "                                ),\n",
      "                                ones(\n",
      "                                    len(\n",
      "                                        sqrt(\n",
      "                                            sum(\n",
      "                                                (p1 + p2 + p3)[:, 1:] ** 2,\n",
      "                                                axis=1,\n",
      "                                            )\n",
      "                                        )\n",
      "                                        * (p1 + p2 + p3)[:, 0] ** (-1.0)\n",
      "                                    )\n",
      "                                ),\n",
      "                                zeros(\n",
      "                                    len(\n",
      "                                        sqrt(\n",
      "                                            sum(\n",
      "                                                (p1 + p2 + p3)[:, 1:] ** 2,\n",
      "                                                axis=1,\n",
      "                                            )\n",
      "                                        )\n",
      "                                        * (p1 + p2 + p3)[:, 0] ** (-1.0)\n",
      "                                    )\n",
      "                                ),\n",
      "                                zeros(\n",
      "                                    len(\n",
      "                                        sqrt(\n",
      "                                            sum(\n",
      "                                                (p1 + p2 + p3)[:, 1:] ** 2,\n",
      "                                                axis=1,\n",
      "                                            )\n",
      "                                        )\n",
      "                                        * (p1 + p2 + p3)[:, 0] ** (-1.0)\n",
      "                                    )\n",
      "                                ),\n",
      "                            ],\n",
      "                            [\n",
      "                                zeros(\n",
      "                                    len(\n",
      "                                        sqrt(\n",
      "                                            sum(\n",
      "                                                (p1 + p2 + p3)[:, 1:] ** 2,\n",
      "                                                axis=1,\n",
      "                                            )\n",
      "                                        )\n",
      "                                        * (p1 + p2 + p3)[:, 0] ** (-1.0)\n",
      "                                    )\n",
      "                                ),\n",
      "                                zeros(\n",
      "                                    len(\n",
      "                                        sqrt(\n",
      "                                            sum(\n",
      "                                                (p1 + p2 + p3)[:, 1:] ** 2,\n",
      "                                                axis=1,\n",
      "                                            )\n",
      "                                        )\n",
      "                                        * (p1 + p2 + p3)[:, 0] ** (-1.0)\n",
      "                                    )\n",
      "                                ),\n",
      "                                ones(\n",
      "                                    len(\n",
      "                                        sqrt(\n",
      "                                            sum(\n",
      "                                                (p1 + p2 + p3)[:, 1:] ** 2,\n",
      "                                                axis=1,\n",
      "                                            )\n",
      "                                        )\n",
      "                                        * (p1 + p2 + p3)[:, 0] ** (-1.0)\n",
      "                                    )\n",
      "                                ),\n",
      "                                zeros(\n",
      "                                    len(\n",
      "                                        sqrt(\n",
      "                                            sum(\n",
      "                                                (p1 + p2 + p3)[:, 1:] ** 2,\n",
      "                                                axis=1,\n",
      "                                            )\n",
      "                                        )\n",
      "                                        * (p1 + p2 + p3)[:, 0] ** (-1.0)\n",
      "                                    )\n",
      "                                ),\n",
      "                            ],\n",
      "                            [\n",
      "                                -1\n",
      "                                / sqrt(\n",
      "                                    1\n",
      "                                    - (\n",
      "                                        sqrt(\n",
      "                                            sum(\n",
      "                                                (p1 + p2 + p3)[:, 1:] ** 2,\n",
      "                                                axis=1,\n",
      "                                            )\n",
      "                                        )\n",
      "                                        * (p1 + p2 + p3)[:, 0] ** (-1.0)\n",
      "                                    )\n",
      "                                    ** 2\n",
      "                                )\n",
      "                                * sqrt(sum((p1 + p2 + p3)[:, 1:] ** 2, axis=1))\n",
      "                                * (p1 + p2 + p3)[:, 0] ** (-1.0),\n",
      "                                zeros(\n",
      "                                    len(\n",
      "                                        sqrt(\n",
      "                                            sum(\n",
      "                                                (p1 + p2 + p3)[:, 1:] ** 2,\n",
      "                                                axis=1,\n",
      "                                            )\n",
      "                                        )\n",
      "                                        * (p1 + p2 + p3)[:, 0] ** (-1.0)\n",
      "                                    )\n",
      "                                ),\n",
      "                                zeros(\n",
      "                                    len(\n",
      "                                        sqrt(\n",
      "                                            sum(\n",
      "                                                (p1 + p2 + p3)[:, 1:] ** 2,\n",
      "                                                axis=1,\n",
      "                                            )\n",
      "                                        )\n",
      "                                        * (p1 + p2 + p3)[:, 0] ** (-1.0)\n",
      "                                    )\n",
      "                                ),\n",
      "                                1\n",
      "                                / sqrt(\n",
      "                                    1\n",
      "                                    - (\n",
      "                                        sqrt(\n",
      "                                            sum(\n",
      "                                                (p1 + p2 + p3)[:, 1:] ** 2,\n",
      "                                                axis=1,\n",
      "                                            )\n",
      "                                        )\n",
      "                                        * (p1 + p2 + p3)[:, 0] ** (-1.0)\n",
      "                                    )\n",
      "                                    ** 2\n",
      "                                ),\n",
      "                            ],\n",
      "                        ]\n",
      "                    ).transpose(2, 0, 1),\n",
      "                    axes=(1, 2, 0),\n",
      "                ),\n",
      "                transpose(\n",
      "                    einsum(\n",
      "                        \"ij...,j...\",\n",
      "                        transpose(\n",
      "                            array(\n",
      "                                [\n",
      "                                    [\n",
      "                                        ones(\n",
      "                                            len(\n",
      "                                                -arccos(\n",
      "                                                    (p1 + p2 + p3)[:, 3]\n",
      "                                                    * sum(\n",
      "                                                        (p1 + p2 + p3)[:, 1:]\n",
      "                                                        ** 2,\n",
      "                                                        axis=1,\n",
      "                                                    )\n",
      "                                                    ** (-1 / 2)\n",
      "                                                )\n",
      "                                            )\n",
      "                                        ),\n",
      "                                        zeros(\n",
      "                                            len(\n",
      "                                                -arccos(\n",
      "                                                    (p1 + p2 + p3)[:, 3]\n",
      "                                                    * sum(\n",
      "                                                        (p1 + p2 + p3)[:, 1:]\n",
      "                                                        ** 2,\n",
      "                                                        axis=1,\n",
      "                                                    )\n",
      "                                                    ** (-1 / 2)\n",
      "                                                )\n",
      "                                            )\n",
      "                                        ),\n",
      "                                        zeros(\n",
      "                                            len(\n",
      "                                                -arccos(\n",
      "                                                    (p1 + p2 + p3)[:, 3]\n",
      "                                                    * sum(\n",
      "                                                        (p1 + p2 + p3)[:, 1:]\n",
      "                                                        ** 2,\n",
      "                                                        axis=1,\n",
      "                                                    )\n",
      "                                                    ** (-1 / 2)\n",
      "                                                )\n",
      "                                            )\n",
      "                                        ),\n",
      "                                        zeros(\n",
      "                                            len(\n",
      "                                                -arccos(\n",
      "                                                    (p1 + p2 + p3)[:, 3]\n",
      "                                                    * sum(\n",
      "                                                        (p1 + p2 + p3)[:, 1:]\n",
      "                                                        ** 2,\n",
      "                                                        axis=1,\n",
      "                                                    )\n",
      "                                                    ** (-1 / 2)\n",
      "                                                )\n",
      "                                            )\n",
      "                                        ),\n",
      "                                    ],\n",
      "                                    [\n",
      "                                        zeros(\n",
      "                                            len(\n",
      "                                                -arccos(\n",
      "                                                    (p1 + p2 + p3)[:, 3]\n",
      "                                                    * sum(\n",
      "                                                        (p1 + p2 + p3)[:, 1:]\n",
      "                                                        ** 2,\n",
      "                                                        axis=1,\n",
      "                                                    )\n",
      "                                                    ** (-1 / 2)\n",
      "                                                )\n",
      "                                            )\n",
      "                                        ),\n",
      "                                        cos(\n",
      "                                            -arccos(\n",
      "                                                (p1 + p2 + p3)[:, 3]\n",
      "                                                * sum(\n",
      "                                                    (p1 + p2 + p3)[:, 1:] ** 2,\n",
      "                                                    axis=1,\n",
      "                                                )\n",
      "                                                ** (-1 / 2)\n",
      "                                            )\n",
      "                                        ),\n",
      "                                        zeros(\n",
      "                                            len(\n",
      "                                                -arccos(\n",
      "                                                    (p1 + p2 + p3)[:, 3]\n",
      "                                                    * sum(\n",
      "                                                        (p1 + p2 + p3)[:, 1:]\n",
      "                                                        ** 2,\n",
      "                                                        axis=1,\n",
      "                                                    )\n",
      "                                                    ** (-1 / 2)\n",
      "                                                )\n",
      "                                            )\n",
      "                                        ),\n",
      "                                        sin(\n",
      "                                            -arccos(\n",
      "                                                (p1 + p2 + p3)[:, 3]\n",
      "                                                * sum(\n",
      "                                                    (p1 + p2 + p3)[:, 1:] ** 2,\n",
      "                                                    axis=1,\n",
      "                                                )\n",
      "                                                ** (-1 / 2)\n",
      "                                            )\n",
      "                                        ),\n",
      "                                    ],\n",
      "                                    [\n",
      "                                        zeros(\n",
      "                                            len(\n",
      "                                                -arccos(\n",
      "                                                    (p1 + p2 + p3)[:, 3]\n",
      "                                                    * sum(\n",
      "                                                        (p1 + p2 + p3)[:, 1:]\n",
      "                                                        ** 2,\n",
      "                                                        axis=1,\n",
      "                                                    )\n",
      "                                                    ** (-1 / 2)\n",
      "                                                )\n",
      "                                            )\n",
      "                                        ),\n",
      "                                        zeros(\n",
      "                                            len(\n",
      "                                                -arccos(\n",
      "                                                    (p1 + p2 + p3)[:, 3]\n",
      "                                                    * sum(\n",
      "                                                        (p1 + p2 + p3)[:, 1:]\n",
      "                                                        ** 2,\n",
      "                                                        axis=1,\n",
      "                                                    )\n",
      "                                                    ** (-1 / 2)\n",
      "                                                )\n",
      "                                            )\n",
      "                                        ),\n",
      "                                        ones(\n",
      "                                            len(\n",
      "                                                -arccos(\n",
      "                                                    (p1 + p2 + p3)[:, 3]\n",
      "                                                    * sum(\n",
      "                                                        (p1 + p2 + p3)[:, 1:]\n",
      "                                                        ** 2,\n",
      "                                                        axis=1,\n",
      "                                                    )\n",
      "                                                    ** (-1 / 2)\n",
      "                                                )\n",
      "                                            )\n",
      "                                        ),\n",
      "                                        zeros(\n",
      "                                            len(\n",
      "                                                -arccos(\n",
      "                                                    (p1 + p2 + p3)[:, 3]\n",
      "                                                    * sum(\n",
      "                                                        (p1 + p2 + p3)[:, 1:]\n",
      "                                                        ** 2,\n",
      "                                                        axis=1,\n",
      "                                                    )\n",
      "                                                    ** (-1 / 2)\n",
      "                                                )\n",
      "                                            )\n",
      "                                        ),\n",
      "                                    ],\n",
      "                                    [\n",
      "                                        zeros(\n",
      "                                            len(\n",
      "                                                -arccos(\n",
      "                                                    (p1 + p2 + p3)[:, 3]\n",
      "                                                    * sum(\n",
      "                                                        (p1 + p2 + p3)[:, 1:]\n",
      "                                                        ** 2,\n",
      "                                                        axis=1,\n",
      "                                                    )\n",
      "                                                    ** (-1 / 2)\n",
      "                                                )\n",
      "                                            )\n",
      "                                        ),\n",
      "                                        -sin(\n",
      "                                            -arccos(\n",
      "                                                (p1 + p2 + p3)[:, 3]\n",
      "                                                * sum(\n",
      "                                                    (p1 + p2 + p3)[:, 1:] ** 2,\n",
      "                                                    axis=1,\n",
      "                                                )\n",
      "                                                ** (-1 / 2)\n",
      "                                            )\n",
      "                                        ),\n",
      "                                        zeros(\n",
      "                                            len(\n",
      "                                                -arccos(\n",
      "                                                    (p1 + p2 + p3)[:, 3]\n",
      "                                                    * sum(\n",
      "                                                        (p1 + p2 + p3)[:, 1:]\n",
      "                                                        ** 2,\n",
      "                                                        axis=1,\n",
      "                                                    )\n",
      "                                                    ** (-1 / 2)\n",
      "                                                )\n",
      "                                            )\n",
      "                                        ),\n",
      "                                        cos(\n",
      "                                            -arccos(\n",
      "                                                (p1 + p2 + p3)[:, 3]\n",
      "                                                * sum(\n",
      "                                                    (p1 + p2 + p3)[:, 1:] ** 2,\n",
      "                                                    axis=1,\n",
      "                                                )\n",
      "                                                ** (-1 / 2)\n",
      "                                            )\n",
      "                                        ),\n",
      "                                    ],\n",
      "                                ]\n",
      "                            ).transpose(2, 0, 1),\n",
      "                            axes=(1, 2, 0),\n",
      "                        ),\n",
      "                        transpose(\n",
      "                            einsum(\n",
      "                                \"ij...,j...\",\n",
      "                                transpose(\n",
      "                                    array(\n",
      "                                        [\n",
      "                                            [\n",
      "                                                ones(\n",
      "                                                    len(\n",
      "                                                        -arctan2(\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 2\n",
      "                                                            ],\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 1\n",
      "                                                            ],\n",
      "                                                        )\n",
      "                                                    )\n",
      "                                                ),\n",
      "                                                zeros(\n",
      "                                                    len(\n",
      "                                                        -arctan2(\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 2\n",
      "                                                            ],\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 1\n",
      "                                                            ],\n",
      "                                                        )\n",
      "                                                    )\n",
      "                                                ),\n",
      "                                                zeros(\n",
      "                                                    len(\n",
      "                                                        -arctan2(\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 2\n",
      "                                                            ],\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 1\n",
      "                                                            ],\n",
      "                                                        )\n",
      "                                                    )\n",
      "                                                ),\n",
      "                                                zeros(\n",
      "                                                    len(\n",
      "                                                        -arctan2(\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 2\n",
      "                                                            ],\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 1\n",
      "                                                            ],\n",
      "                                                        )\n",
      "                                                    )\n",
      "                                                ),\n",
      "                                            ],\n",
      "                                            [\n",
      "                                                zeros(\n",
      "                                                    len(\n",
      "                                                        -arctan2(\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 2\n",
      "                                                            ],\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 1\n",
      "                                                            ],\n",
      "                                                        )\n",
      "                                                    )\n",
      "                                                ),\n",
      "                                                cos(\n",
      "                                                    -arctan2(\n",
      "                                                        (p1 + p2 + p3)[:, 2],\n",
      "                                                        (p1 + p2 + p3)[:, 1],\n",
      "                                                    )\n",
      "                                                ),\n",
      "                                                -sin(\n",
      "                                                    -arctan2(\n",
      "                                                        (p1 + p2 + p3)[:, 2],\n",
      "                                                        (p1 + p2 + p3)[:, 1],\n",
      "                                                    )\n",
      "                                                ),\n",
      "                                                zeros(\n",
      "                                                    len(\n",
      "                                                        -arctan2(\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 2\n",
      "                                                            ],\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 1\n",
      "                                                            ],\n",
      "                                                        )\n",
      "                                                    )\n",
      "                                                ),\n",
      "                                            ],\n",
      "                                            [\n",
      "                                                zeros(\n",
      "                                                    len(\n",
      "                                                        -arctan2(\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 2\n",
      "                                                            ],\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 1\n",
      "                                                            ],\n",
      "                                                        )\n",
      "                                                    )\n",
      "                                                ),\n",
      "                                                sin(\n",
      "                                                    -arctan2(\n",
      "                                                        (p1 + p2 + p3)[:, 2],\n",
      "                                                        (p1 + p2 + p3)[:, 1],\n",
      "                                                    )\n",
      "                                                ),\n",
      "                                                cos(\n",
      "                                                    -arctan2(\n",
      "                                                        (p1 + p2 + p3)[:, 2],\n",
      "                                                        (p1 + p2 + p3)[:, 1],\n",
      "                                                    )\n",
      "                                                ),\n",
      "                                                zeros(\n",
      "                                                    len(\n",
      "                                                        -arctan2(\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 2\n",
      "                                                            ],\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 1\n",
      "                                                            ],\n",
      "                                                        )\n",
      "                                                    )\n",
      "                                                ),\n",
      "                                            ],\n",
      "                                            [\n",
      "                                                zeros(\n",
      "                                                    len(\n",
      "                                                        -arctan2(\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 2\n",
      "                                                            ],\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 1\n",
      "                                                            ],\n",
      "                                                        )\n",
      "                                                    )\n",
      "                                                ),\n",
      "                                                zeros(\n",
      "                                                    len(\n",
      "                                                        -arctan2(\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 2\n",
      "                                                            ],\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 1\n",
      "                                                            ],\n",
      "                                                        )\n",
      "                                                    )\n",
      "                                                ),\n",
      "                                                zeros(\n",
      "                                                    len(\n",
      "                                                        -arctan2(\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 2\n",
      "                                                            ],\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 1\n",
      "                                                            ],\n",
      "                                                        )\n",
      "                                                    )\n",
      "                                                ),\n",
      "                                                ones(\n",
      "                                                    len(\n",
      "                                                        -arctan2(\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 2\n",
      "                                                            ],\n",
      "                                                            (p1 + p2 + p3)[\n",
      "                                                                :, 1\n",
      "                                                            ],\n",
      "                                                        )\n",
      "                                                    )\n",
      "                                                ),\n",
      "                                            ],\n",
      "                                        ]\n",
      "                                    ).transpose(2, 0, 1),\n",
      "                                    axes=(1, 2, 0),\n",
      "                                ),\n",
      "                                transpose(p3),\n",
      "                            )\n",
      "                        ),\n",
      "                    )\n",
      "                ),\n",
      "            )\n",
      "        )[:, 1],\n",
      "    )\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "def display_kinematic_variable(name: str) -> None:\n",
    "    expr = symbolic_angles[name]\n",
    "    display(Math(sp.multiline_latex(expr, expr.doit(), **math_style)))\n",
    "    print_lambdify(momentum_symbols.values(), expr.doit())\n",
    "\n",
    "\n",
    "for angle_name in list(symbolic_angles)[:3]:\n",
    "    display_kinematic_variable(angle_name)\n",
    "    print()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Comparison computed value"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, what's left is to compare computed data with the {ref}`report/011:Test sample`. Here's a little comparison function that can be called for each angle:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "def compare(angle_name: str) -> None:\n",
    "    np_expr = sp.lambdify(\n",
    "        args=momentum_symbols.values(),\n",
    "        expr=symbolic_angles[angle_name].doit(),\n",
    "        modules=\"numpy\",\n",
    "    )\n",
    "    computed = np_expr(*events.values())\n",
    "    expected = np.array(angles[angle_name])\n",
    "    display(computed, expected)\n",
    "    np.testing.assert_allclose(computed, expected)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```{autolink-skip}\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "keep_output"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "phi_1+2+3\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "array([ 2.79758029,  2.51292308, -1.07396684])"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "array([ 2.79758029,  2.51292308, -1.07396684])"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 2.48 ms, sys: 6.56 ms, total: 9.04 ms\n",
      "Wall time: 7.77 ms\n",
      "\n",
      "theta_1+2+3\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "array([2.72456853, 3.03316287, 0.69240082])"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "array([2.72456853, 3.03316287, 0.69240082])"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 5.13 ms, sys: 297 µs, total: 5.43 ms\n",
      "Wall time: 4.93 ms\n",
      "\n",
      "phi_2+3,1+2+3\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "array([1.0436215 , 1.8734936 , 0.16073833])"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "array([1.0436215 , 1.8734936 , 0.16073833])"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 34.1 ms, sys: 0 ns, total: 34.1 ms\n",
      "Wall time: 33.1 ms\n",
      "\n",
      "theta_2+3,1+2+3\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "array([2.45361589, 1.40639741, 0.98079245])"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "array([2.45361589, 1.40639741, 0.98079245])"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 36.1 ms, sys: 451 µs, total: 36.6 ms\n",
      "Wall time: 35.5 ms\n",
      "\n",
      "phi_2,2+3,1+2+3\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "array([ 0.36955786, -1.68820498,  0.63063002])"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "array([ 0.36955786, -1.68820498,  0.63063002])"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 1.25 s, sys: 162 ms, total: 1.42 s\n",
      "Wall time: 1.41 s\n",
      "\n",
      "theta_2,2+3,1+2+3\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "array([1.0924374 , 1.99375767, 1.31959621])"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "array([1.0924374 , 1.99375767, 1.31959621])"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 1.22 s, sys: 166 ms, total: 1.38 s\n",
      "Wall time: 1.38 s\n",
      "\n"
     ]
    }
   ],
   "source": [
    "for angle_name in symbolic_angles:\n",
    "    print(angle_name)\n",
    "    %time compare(angle_name)\n",
    "    print()"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}