From 3fe521278d5d7203feb8c8bb8b304218285efe41 Mon Sep 17 00:00:00 2001
From: Vitor Shen <17490173+shenvitor@users.noreply.github.com>
Date: Tue, 20 Aug 2024 14:33:35 +0200
Subject: [PATCH] draft of invariant mass, boost and rotation

---
 docs/eta-pi-p/manual.ipynb | 205 ++++++++++++++++++++++++++++++++++++-
 pyproject.toml             |   1 +
 2 files changed, 203 insertions(+), 3 deletions(-)

diff --git a/docs/eta-pi-p/manual.ipynb b/docs/eta-pi-p/manual.ipynb
index 1ef5c9f..4db7ab7 100644
--- a/docs/eta-pi-p/manual.ipynb
+++ b/docs/eta-pi-p/manual.ipynb
@@ -52,7 +52,24 @@
    "source": [
     "import matplotlib.pyplot as plt\n",
     "import numpy as np\n",
+    "import pandas as pd\n",
     "import sympy as sp\n",
+    "from ampform.kinematics.angles import (\n",
+    "    Phi,\n",
+    "    Theta,\n",
+    ")\n",
+    "from ampform.kinematics.lorentz import (\n",
+    "    ArrayMultiplication,\n",
+    "    ArraySize,\n",
+    "    BoostZMatrix,\n",
+    "    Energy,\n",
+    "    FourMomentumSymbol,\n",
+    "    InvariantMass,\n",
+    "    RotationYMatrix,\n",
+    "    RotationZMatrix,\n",
+    "    three_momentum_norm,\n",
+    ")\n",
+    "from ampform.sympy._array_expressions import ArraySum\n",
     "from tensorwaves.data import TFPhaseSpaceGenerator, TFUniformRealNumberGenerator"
    ]
   },
@@ -249,6 +266,191 @@
     "intensity_func"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "p1 = FourMomentumSymbol(\"p_1\", shape=[])\n",
+    "p2 = FourMomentumSymbol(\"p_2\", shape=[])\n",
+    "p3 = FourMomentumSymbol(\"p_3\", shape=[])\n",
+    "p12 = ArraySum(p1, p2)\n",
+    "p23 = ArraySum(p2, p3)\n",
+    "p31 = ArraySum(p3, p1)\n",
+    "\n",
+    "theta12 = Theta(p12)\n",
+    "theta23 = Theta(p23)\n",
+    "theta31 = Theta(p31)\n",
+    "\n",
+    "phi12 = Phi(p12)\n",
+    "phi23 = Phi(p23)\n",
+    "phi31 = Phi(p31)\n",
+    "\n",
+    "beta12 = three_momentum_norm(p12) / Energy(p12)\n",
+    "beta23 = three_momentum_norm(p23) / Energy(p23)\n",
+    "beta31 = three_momentum_norm(p31) / Energy(p31)\n",
+    "\n",
+    "Bz12 = BoostZMatrix(beta12, n_events=ArraySize(beta12))\n",
+    "Bz23 = BoostZMatrix(beta23, n_events=ArraySize(beta23))\n",
+    "Bz31 = BoostZMatrix(beta31, n_events=ArraySize(beta31))\n",
+    "\n",
+    "Ry12 = RotationYMatrix(-theta12, n_events=ArraySize(theta12))\n",
+    "Ry23 = RotationYMatrix(-theta23, n_events=ArraySize(theta23))\n",
+    "Ry31 = RotationYMatrix(-theta31, n_events=ArraySize(theta31))\n",
+    "\n",
+    "Rz12 = RotationZMatrix(-phi12, n_events=ArraySize(theta12))\n",
+    "Rz23 = RotationZMatrix(-phi23, n_events=ArraySize(theta23))\n",
+    "Rz31 = RotationZMatrix(-phi31, n_events=ArraySize(theta31))\n",
+    "\n",
+    "m12 = InvariantMass(p12)\n",
+    "m23 = InvariantMass(p23)\n",
+    "m31 = InvariantMass(p31)\n",
+    "\n",
+    "phi12_helicity = Phi(ArrayMultiplication(Bz31, Ry31, Rz31, p3))\n",
+    "phi23_helicity = Phi(ArrayMultiplication(Bz23, Ry23, Rz23, p2))\n",
+    "phi31_helicity = Phi(ArrayMultiplication(Bz31, Ry31, Rz31, p3))\n",
+    "\n",
+    "theta12_helicity = Theta(ArrayMultiplication(Bz12, Ry12, Rz12, p1))\n",
+    "theta23_helicity = Theta(ArrayMultiplication(Bz23, Ry23, Rz23, p2))\n",
+    "theta31_helicity = Theta(ArrayMultiplication(Bz31, Ry31, Rz31, p3))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "phi12_helicity"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "theta12_helicity"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "phi23_helicity"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "theta23_helicity"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "phi31_helicity"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "theta31_helicity"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "m12"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "m23"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "m31"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "theta12"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "theta23"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "theta31"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "phi12"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "phi23"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "phi31"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -297,9 +499,6 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import numpy as np\n",
-    "import pandas as pd\n",
-    "\n",
     "pd.DataFrame({\n",
     "    (k, label): np.transpose(v)[i]\n",
     "    for k, v in phsp_momenta.items()\n",
diff --git a/pyproject.toml b/pyproject.toml
index a93930e..27a0560 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -111,6 +111,7 @@ builtins-ignorelist = ["display"]
     "E303",
     "E501",
     "PLC2401",
+    "PLC2701",
     "PLR2004",
     "PLW0603",
     "S101",