From ab68d833d84b6746665d82520b377da25f307892 Mon Sep 17 00:00:00 2001
From: Vitor Shen <17490173+shenvitor@users.noreply.github.com>
Date: Mon, 19 Aug 2024 15:54:05 +0200
Subject: [PATCH] phsp 4-momenta generation via tensorwave

---
 docs/eta-pi-p/manual.ipynb | 66 +++++++++++++++++++++++++++++++++++---
 1 file changed, 61 insertions(+), 5 deletions(-)

diff --git a/docs/eta-pi-p/manual.ipynb b/docs/eta-pi-p/manual.ipynb
index 082a2d9..1ef5c9f 100644
--- a/docs/eta-pi-p/manual.ipynb
+++ b/docs/eta-pi-p/manual.ipynb
@@ -41,9 +41,6 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {
-    "jupyter": {
-     "source_hidden": true
-    },
     "mystnb": {
      "code_prompt_show": "Import Python libraries"
     },
@@ -55,7 +52,8 @@
    "source": [
     "import matplotlib.pyplot as plt\n",
     "import numpy as np\n",
-    "import sympy as sp"
+    "import sympy as sp\n",
+    "from tensorwaves.data import TFPhaseSpaceGenerator, TFUniformRealNumberGenerator"
    ]
   },
   {
@@ -251,6 +249,64 @@
     "intensity_func"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Phase Space Generation"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Mass for $p\\gamma$ system"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "E_lab_gamma = 8.5\n",
+    "m_proton = 0.938\n",
+    "m_0 = np.sqrt(2 * E_lab_gamma * m_proton + m_proton**2)\n",
+    "m_eta = 0.548\n",
+    "m_pi = 0.135\n",
+    "m_0"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "rng = TFUniformRealNumberGenerator(seed=0)\n",
+    "phsp_generator = TFPhaseSpaceGenerator(\n",
+    "    initial_state_mass=m_0,\n",
+    "    final_state_masses={1: m_eta, 2: m_pi, 3: m_proton},\n",
+    ")\n",
+    "phsp_momenta = phsp_generator.generate(100_000, rng)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "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",
+    "    for i, label in enumerate([\"E\", \"px\", \"py\", \"pz\"])\n",
+    "})"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -401,7 +457,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.12.4"
+   "version": "3.12.5"
   }
  },
  "nbformat": 4,