diff --git a/docs/lambda-k-pi/manual-symbolic.ipynb b/docs/lambda-k-pi/manual-symbolic.ipynb index b25c642..d9ca8e0 100644 --- a/docs/lambda-k-pi/manual-symbolic.ipynb +++ b/docs/lambda-k-pi/manual-symbolic.ipynb @@ -51,7 +51,30 @@ "import os\n", "import warnings\n", "\n", + "import numpy as np\n", "import sympy as sp\n", + "from ampform.io import aslatex\n", + "from ampform.kinematics.angles import Phi, Theta\n", + "from ampform.kinematics.lorentz import (\n", + " ArrayMultiplication,\n", + " ArraySize,\n", + " BoostZMatrix,\n", + " Energy,\n", + " EuclideanNorm,\n", + " FourMomentumSymbol,\n", + " RotationYMatrix,\n", + " RotationZMatrix,\n", + " ThreeMomentum,\n", + " three_momentum_norm,\n", + ")\n", + "from ampform.sympy import unevaluated\n", + "from ampform.sympy._array_expressions import ArraySum\n", + "from IPython.display import Latex\n", + "from tensorwaves.data import (\n", + " SympyDataTransformer,\n", + " TFPhaseSpaceGenerator,\n", + " TFUniformRealNumberGenerator,\n", + ")\n", "\n", "STATIC_PAGE = \"EXECUTE_NB\" in os.environ\n", "\n", @@ -231,6 +254,212 @@ "intensity_expr = sp.Abs(A12 + A23 + A31) ** 2\n", "intensity_expr" ] + }, + { + "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_lambda = 1.12\n", + "m_k = 0.494\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_k, 2: m_pi, 3: m_lambda},\n", + ")\n", + "phsp_momenta = phsp_generator.generate(500_000, rng)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Kinematic variables" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [ + "hide-input" + ] + }, + "outputs": [], + "source": [ + "@unevaluated\n", + "class SquaredInvariantMass(sp.Expr):\n", + " momentum: sp.Basic\n", + " _latex_repr_ = \"m_{{{momentum}}}^2\"\n", + "\n", + " def evaluate(self) -> sp.Expr:\n", + " p = self.momentum\n", + " p_xyz = ThreeMomentum(p)\n", + " return Energy(p) ** 2 - EuclideanNorm(p_xyz) ** 2\n", + "\n", + "\n", + "def formulate_helicity_angles(\n", + " pi: FourMomentumSymbol, pj: FourMomentumSymbol\n", + ") -> tuple[Theta, Phi]:\n", + " pij = ArraySum(pi, pj)\n", + " beta = three_momentum_norm(pij) / Energy(pij)\n", + " Rz = RotationZMatrix(-Phi(pij), n_events=ArraySize(Phi(pij)))\n", + " Ry = RotationYMatrix(-Theta(pij), n_events=ArraySize(Theta(pij)))\n", + " Bz = BoostZMatrix(beta, n_events=ArraySize(beta))\n", + " pi_boosted = ArrayMultiplication(Bz, Ry, Rz, pi)\n", + " return Theta(pi_boosted), Phi(pi_boosted)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "p1 = FourMomentumSymbol(\"p1\", shape=[])\n", + "p2 = FourMomentumSymbol(\"p2\", shape=[])\n", + "p3 = FourMomentumSymbol(\"p3\", shape=[])\n", + "p12 = ArraySum(p1, p2)\n", + "p23 = ArraySum(p2, p3)\n", + "p31 = ArraySum(p3, p1)\n", + "\n", + "theta1_expr, phi1_expr = formulate_helicity_angles(p1, p2)\n", + "theta2_expr, phi2_expr = formulate_helicity_angles(p2, p3)\n", + "theta3_expr, phi3_expr = formulate_helicity_angles(p3, p1)\n", + "\n", + "s12_expr = SquaredInvariantMass(p12)\n", + "s23_expr = SquaredInvariantMass(p23)\n", + "s31_expr = SquaredInvariantMass(p31)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "kinematic_variables = {\n", + " theta1: theta1_expr,\n", + " theta2: theta2_expr,\n", + " theta3: theta3_expr,\n", + " phi1: phi1_expr,\n", + " phi2: phi2_expr,\n", + " phi3: phi3_expr,\n", + " s12: s12_expr,\n", + " s23: s23_expr,\n", + " s31: s31_expr,\n", + "}\n", + "\n", + "Latex(aslatex(kinematic_variables))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "helicity_transformer = SympyDataTransformer.from_sympy(\n", + " kinematic_variables, backend=\"jax\"\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "phsp = helicity_transformer(phsp_momenta)\n", + "list(phsp)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Parameters" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "tags": [ + "hide-input", + "scroll-input" + ] + }, + "outputs": [], + "source": [ + "a_vals = [0, 1.0, 3.0, 3.5, 2.0] # Slight adjustment to emphasize higher waves\n", + "b_vals = [0, -1.0, 3.5, 0.5, 0] # Adjust for new final state coupling\n", + "c_vals = [0, 0, 3.0, 0, 0] # Adjust if more s-wave or p-wave is expected\n", + "\n", + "m_Kstar2_val = 1.43\n", + "m_Sigma_val = 1.385\n", + "m_Nstar_val = 1.71\n", + "\n", + "Gamma_Kstar2_val = 0.62\n", + "Gamma_Sigma_val = 0.2\n", + "Gamma_Nstar_val = 0.27\n", + "\n", + "l12_val = 2 # I still use 2 assuming K^*_2\n", + "l23_val = 1\n", + "l31_val = 0\n", + "\n", + "parameters_default = {\n", + " m_Kstar2: m_Kstar2_val,\n", + " m_Sigma: m_Sigma_val,\n", + " m_Nstar: m_Nstar_val,\n", + " Gamma_Kstar2: Gamma_Kstar2_val,\n", + " Gamma_Sigma: Gamma_Sigma_val,\n", + " Gamma_Nstar: Gamma_Nstar_val,\n", + " l12: l12_val,\n", + " l23: l23_val,\n", + " l31: l31_val,\n", + "}\n", + "\n", + "a_dict = {a[i]: a_vals[i + l_max] for i in range(-l_max, l_max + 1)}\n", + "b_dict = {b[i]: b_vals[i + l_max] for i in range(-l_max, l_max + 1)}\n", + "c_dict = {c[i]: c_vals[i + l_max] for i in range(-l_max, l_max + 1)}\n", + "parameters_default.update(a_dict)\n", + "parameters_default.update(b_dict)\n", + "parameters_default.update(c_dict)\n", + "\n", + "Latex(aslatex(parameters_default))" + ] } ], "metadata": {