diff --git a/docs/eta-pi-p/manual.ipynb b/docs/eta-pi-p/manual.ipynb index 77b087b..f7f6bb7 100644 --- a/docs/eta-pi-p/manual.ipynb +++ b/docs/eta-pi-p/manual.ipynb @@ -54,6 +54,7 @@ "import numpy as np\n", "import pandas as pd\n", "import sympy as sp\n", + "from ampform.io import aslatex\n", "from ampform.kinematics.angles import (\n", " Phi,\n", " Theta,\n", @@ -63,14 +64,22 @@ " ArraySize,\n", " BoostZMatrix,\n", " Energy,\n", + " EuclideanNorm,\n", " FourMomentumSymbol,\n", - " InvariantMass,\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 tensorwaves.data import TFPhaseSpaceGenerator, TFUniformRealNumberGenerator" + "from IPython.display import Latex\n", + "from tensorwaves.data import (\n", + " SympyDataTransformer,\n", + " TFPhaseSpaceGenerator,\n", + " TFUniformRealNumberGenerator,\n", + ")\n", + "from tensorwaves.function.sympy import create_parametrized_function" ] }, { @@ -93,12 +102,12 @@ "metadata": {}, "outputs": [], "source": [ - "s12, m_a2, Gamma_a2 = sp.symbols(r\"s_{12} m_{a_2} \\Gamma_{a_2}\")\n", + "s12, m_a2, gamma_a2 = sp.symbols(r\"s_{12} m_{a_2} \\Gamma_{a_2}\")\n", "theta1, phi1 = sp.symbols(\"theta_1 phi_1\")\n", "a = sp.IndexedBase(\"a\")\n", "m = sp.symbols(\"m\", cls=sp.Idx)\n", "A12 = sp.Sum(a[m] * sp.Ynm(2, m, theta1, phi1), (m, -2, 2)) / (\n", - " s12 - m_a2**2 + sp.I * m_a2 * Gamma_a2\n", + " s12 - m_a2**2 + sp.I * m_a2 * gamma_a2\n", ")\n", "A12" ] @@ -137,7 +146,7 @@ "outputs": [], "source": [ "A12_func = sp.lambdify(\n", - " [s12, a[-2], a[-1], a[0], a[1], a[2], m_a2, Gamma_a2, theta1, phi1],\n", + " [s12, a[-2], a[-1], a[0], a[1], a[2], m_a2, gamma_a2, theta1, phi1],\n", " A12.doit().expand(func=True),\n", ")\n", "A12_func" @@ -210,7 +219,7 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "### $A = A^{12}+A^{23}+A^{31}$" + "### $I = |A|^2 = |A^{12}+A^{23}+A^{31}|^2$" ] }, { @@ -227,9 +236,6 @@ "cell_type": "code", "execution_count": null, "metadata": { - "jupyter": { - "source_hidden": true - }, "tags": [ "scroll-input" ] @@ -245,7 +251,7 @@ " a[1],\n", " a[2],\n", " m_a2,\n", - " Gamma_a2,\n", + " gamma_a2,\n", " theta1,\n", " phi1,\n", " s23,\n", @@ -266,15 +272,43 @@ "intensity_func" ] }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Kinematic variables" + ] + }, { "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", + "@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" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + } + }, + "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", @@ -303,9 +337,9 @@ "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", + "m12_square = SquaredInvariantMass(p12)\n", + "m23_square = SquaredInvariantMass(p23)\n", + "m31_square = SquaredInvariantMass(p31)\n", "\n", "phi12_helicity = Phi(ArrayMultiplication(Bz12, Ry12, Rz12, p1))\n", "phi23_helicity = Phi(ArrayMultiplication(Bz23, Ry23, Rz23, p2))\n", @@ -319,109 +353,27 @@ { "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": {}, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "tags": [ + "hide-input" + ] + }, "outputs": [], "source": [ - "theta31" + "kinematic_variables = {\n", + " theta1: theta12_helicity,\n", + " theta2: theta23_helicity,\n", + " phi1: phi12_helicity,\n", + " phi2: phi23_helicity,\n", + " s12: m12_square,\n", + " s23: m23_square,\n", + " s31: m31_square,\n", + "}\n", + "\n", + "Latex(aslatex(kinematic_variables))" ] }, { @@ -430,25 +382,60 @@ "metadata": {}, "outputs": [], "source": [ - "phi12" + "helicity_transformer = SympyDataTransformer.from_sympy(\n", + " kinematic_variables, backend=\"jax\"\n", + ")" ] }, { - "cell_type": "code", - "execution_count": null, + "cell_type": "markdown", "metadata": {}, - "outputs": [], "source": [ - "phi23" + "### Parameters" ] }, { "cell_type": "code", "execution_count": null, - "metadata": {}, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "tags": [ + "hide-input" + ] + }, "outputs": [], "source": [ - "phi31" + "a_vals = [0, 0.5, 3.5, 4, 2.5]\n", + "b_vals = [-1.5, 4, 0.5]\n", + "c0_val = 2.5\n", + "m_a2_val = 1.32\n", + "m_delta_val = 1.54\n", + "m_nstar_val = 1.87\n", + "gamma_a2_val = 0.1\n", + "gamma_delta_val = 0.1\n", + "gamma_nstar_val = 0.1\n", + "\n", + "parameters_default = {\n", + " a[-2]: a_vals[0],\n", + " a[-1]: a_vals[1],\n", + " a[0]: a_vals[2],\n", + " a[1]: a_vals[3],\n", + " a[2]: a_vals[4],\n", + " b[-1]: b_vals[0],\n", + " b[0]: b_vals[1],\n", + " b[1]: b_vals[2],\n", + " c0: c0_val,\n", + " m_a2: m_a2_val,\n", + " m_delta: m_delta_val,\n", + " m_nstar: m_nstar_val,\n", + " gamma_a2: gamma_a2_val,\n", + " gamma_delta: gamma_delta_val,\n", + " gamma_nstar: gamma_nstar_val,\n", + "}\n", + "\n", + "Latex(aslatex(parameters_default))" ] }, { @@ -490,7 +477,7 @@ " 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)" + "phsp_momenta = phsp_generator.generate(500_000, rng)" ] }, { @@ -514,20 +501,10 @@ ] }, { - "cell_type": "code", - "execution_count": null, + "cell_type": "markdown", "metadata": {}, - "outputs": [], "source": [ - "a_vals = [0, 0.5, 3.5, 4, 2.5]\n", - "b_vals = [-1.5, 4, 0.5]\n", - "c0_val = 2.5\n", - "m_a2_val = 1.32\n", - "m_delta_val = 1.54\n", - "m_nstar_val = 1.87\n", - "Gamma_a2_val = 0.1\n", - "gamma_delta_val = 0.1\n", - "gamma_nstar_val = 0.1" + "### Model components" ] }, { @@ -539,7 +516,7 @@ "phi = np.pi / 4\n", "theta = np.pi / 4\n", "s_values = np.linspace(0, 5, num=500)\n", - "A12_values = A12_func(s_values, *a_vals, m_a2_val, Gamma_a2_val, theta, phi)\n", + "A12_values = A12_func(s_values, *a_vals, m_a2_val, gamma_a2_val, theta, phi)\n", "A23_values = A23_func(s_values, *b_vals, m_delta_val, gamma_delta_val, theta, phi)\n", "A31_values = A31_func(s_values, c0_val, m_nstar_val, gamma_nstar_val)" ] @@ -638,6 +615,71 @@ "fig.tight_layout()\n", "plt.show(fig)" ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Phase space" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "phsp = helicity_transformer(phsp_momenta)\n", + "list(phsp)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "unfolded_expression = intensity_expr.doit().expand(func=True)\n", + "intensity_func = create_parametrized_function(\n", + " expression=unfolded_expression,\n", + " parameters=parameters_default,\n", + " backend=\"jax\",\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "%config InlineBackend.figure_formats = ['png']\n", + "\n", + "fig, ax = plt.subplots(dpi=200)\n", + "hist = ax.hist2d(\n", + " phsp[\"s_{12}\"],\n", + " phsp[\"s_{23}\"],\n", + " bins=200,\n", + " cmin=1e-6,\n", + " density=True,\n", + " cmap=\"jet\",\n", + " vmax=0.15,\n", + " weights=intensity_func(phsp),\n", + ")\n", + "ax.set_title(\"Model-weighted Phase space Dalitz Plot\")\n", + "ax.set_xlabel(R\"$m^2(\\eta \\pi^0)\\;\\left[\\mathrm{GeV}^2\\right]$\")\n", + "ax.set_ylabel(R\"$m^2(\\pi^0 p)\\;\\left[\\mathrm{GeV}^2\\right]$\")\n", + "cbar = fig.colorbar(hist[3], ax=ax)\n", + "fig.tight_layout()\n", + "plt.show()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] } ], "metadata": { diff --git a/pyproject.toml b/pyproject.toml index 27a0560..9ff5bfe 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -106,8 +106,7 @@ builtins-ignorelist = ["display"] "*.ipynb" = [ "ANN", "B018", - "D100", - "D103", + "D1", "E303", "E501", "PLC2401",