From b9dc49caf0d667d44cd1064515edde14f48bff52 Mon Sep 17 00:00:00 2001 From: Remco de Boer <29308176+redeboer@users.noreply.github.com> Date: Fri, 26 Apr 2024 21:55:53 +0200 Subject: [PATCH] FEAT: define `RelativisticBreitWigner` and `FormFactor` --- docs/jpsi2ksp.ipynb | 126 ++++++--------------------- src/ampform_dpd/dynamics/__init__.py | 52 +++++++++++ 2 files changed, 81 insertions(+), 97 deletions(-) diff --git a/docs/jpsi2ksp.ipynb b/docs/jpsi2ksp.ipynb index bd6891d9..a62dffb3 100644 --- a/docs/jpsi2ksp.ipynb +++ b/docs/jpsi2ksp.ipynb @@ -31,16 +31,16 @@ "import logging\n", "import os\n", "import warnings\n", - "from typing import TYPE_CHECKING, Any\n", + "from typing import TYPE_CHECKING\n", "\n", "import graphviz\n", "import jax.numpy as jnp\n", "import matplotlib.pyplot as plt\n", "import qrules\n", "import sympy as sp\n", - "from ampform.dynamics import EnergyDependentWidth, formulate_form_factor\n", + "from ampform.dynamics import BlattWeisskopfSquared, EnergyDependentWidth\n", "from ampform.kinematics.phasespace import compute_third_mandelstam\n", - "from ampform.sympy import perform_cached_doit, unevaluated\n", + "from ampform.sympy import perform_cached_doit\n", "from IPython.display import Latex, Markdown\n", "from tensorwaves.data.transform import SympyDataTransformer\n", "from tqdm.auto import tqdm\n", @@ -48,6 +48,7 @@ "from ampform_dpd import DalitzPlotDecompositionBuilder, get_particle\n", "from ampform_dpd.adapter.qrules import normalize_state_ids, to_three_body_decay\n", "from ampform_dpd.decay import IsobarNode, Particle, ThreeBodyDecayChain\n", + "from ampform_dpd.dynamics import FormFactor, RelativisticBreitWigner\n", "from ampform_dpd.io import (\n", " as_markdown_table,\n", " aslatex,\n", @@ -68,7 +69,9 @@ }, { "cell_type": "markdown", - "metadata": {}, + "metadata": { + "jp-MarkdownHeadingCollapsed": true + }, "source": [ "## Decay definition" ] @@ -158,7 +161,9 @@ }, { "cell_type": "markdown", - "metadata": {}, + "metadata": { + "jp-MarkdownHeadingCollapsed": true + }, "source": [ "## Lineshapes for dynamics" ] @@ -172,13 +177,6 @@ ":::" ] }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "In the following, we define the **relativistic Breit-Wigner function** as:" - ] - }, { "cell_type": "code", "execution_count": null, @@ -187,48 +185,20 @@ "source_hidden": true }, "tags": [ - "hide-input" + "hide-input", + "full-width" ] }, "outputs": [], "source": [ - "@unevaluated\n", - "class RelativisticBreitWigner(sp.Expr):\n", - " s: Any\n", - " mass0: Any\n", - " gamma0: Any\n", - " m1: Any\n", - " m2: Any\n", - " angular_momentum: Any\n", - " meson_radius: Any\n", - " _latex_repr_ = (\n", - " R\"\\mathcal{{R}}_{{{angular_momentum}}}\\left({s}, {mass0}, {gamma0}\\right)\"\n", - " )\n", - "\n", - " def evaluate(self):\n", - " s, m0, w0, m1, m2, angular_momentum, meson_radius = self.args\n", - " width = EnergyDependentWidth(\n", - " s=s,\n", - " mass0=m0,\n", - " gamma0=w0,\n", - " m_a=m1,\n", - " m_b=m2,\n", - " angular_momentum=angular_momentum,\n", - " meson_radius=meson_radius,\n", - " name=Rf\"\\Gamma_{{{sp.latex(angular_momentum)}}}\",\n", - " )\n", - " return (m0 * w0) / (m0**2 - s - width * m0 * sp.I)\n", - "\n", - "\n", - "bw = RelativisticBreitWigner(*sp.symbols(\"s m0 Gamma0 m1 m2 L R\"))\n", - "Latex(aslatex({bw: bw.doit(deep=False)}))" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "with $\\Gamma_0(s)$ a {class}`~ampform.dynamics.EnergyDependentWidth`, and we define the **form factor** as:" + "s, m0, w0, m1, m2, L, R, z = sp.symbols(\"s m0 Gamma0 m1 m2 L R z\")\n", + "exprs = [\n", + " RelativisticBreitWigner(s, m0, w0, m1, m2, L, R),\n", + " EnergyDependentWidth(s, m0, w0, m1, m2, L, R),\n", + " FormFactor(s, m1, m2, L, R),\n", + " BlattWeisskopfSquared(z, L),\n", + "]\n", + "Latex(aslatex({e: e.doit(deep=False) for e in exprs}))" ] }, { @@ -238,50 +208,8 @@ "jupyter": { "source_hidden": true }, - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "@unevaluated\n", - "class FormFactor(sp.Expr):\n", - " s: Any\n", - " m1: Any\n", - " m2: Any\n", - " angular_momentum: Any\n", - " meson_radius: Any\n", - "\n", - " _latex_repr_ = R\"\\mathcal{{F}}_{{{angular_momentum}}}\\left({s}, {m1}, {m2}\\right)\"\n", - "\n", - " def evaluate(self):\n", - " s, m1, m2, angular_momentum, meson_radius = self.args\n", - " return formulate_form_factor(\n", - " s=s,\n", - " m_a=m1,\n", - " m_b=m2,\n", - " angular_momentum=angular_momentum,\n", - " meson_radius=meson_radius,\n", - " )\n", - "\n", - "\n", - "ff = FormFactor(*sp.symbols(\"s m1 m2 L R\"))\n", - "Latex(aslatex({ff: ff.doit(deep=False)}))" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Here, $B_L^2$ is a {class}`~ampform.dynamics.BlattWeisskopfSquared`." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true + "mystnb": { + "code_prompt_show": "Define dynamics builder functions" }, "tags": [ "hide-input", @@ -400,7 +328,9 @@ }, { "cell_type": "markdown", - "metadata": {}, + "metadata": { + "jp-MarkdownHeadingCollapsed": true + }, "source": [ "## Model formulation" ] @@ -473,12 +403,14 @@ }, "outputs": [], "source": [ - "Latex(aslatex(model.amplitudes))" + "Latex(aslatex({k: v for k, v in model.amplitudes.items() if v}))" ] }, { "cell_type": "markdown", - "metadata": {}, + "metadata": { + "jp-MarkdownHeadingCollapsed": true + }, "source": [ "## Preparing for input data" ] diff --git a/src/ampform_dpd/dynamics/__init__.py b/src/ampform_dpd/dynamics/__init__.py index 13a5a118..18ec31db 100644 --- a/src/ampform_dpd/dynamics/__init__.py +++ b/src/ampform_dpd/dynamics/__init__.py @@ -5,10 +5,41 @@ from typing import Any import sympy as sp +from ampform.dynamics import formulate_form_factor from ampform.kinematics.phasespace import Kallen from ampform.sympy import unevaluated +@unevaluated +class RelativisticBreitWigner(sp.Expr): + s: Any + mass0: Any + gamma0: Any + m1: Any + m2: Any + angular_momentum: Any + meson_radius: Any + _latex_repr_ = ( + R"\mathcal{{R}}_{{{angular_momentum}}}\left({s}, {mass0}, {gamma0}\right)" + ) + + def evaluate(self): + from ampform.dynamics import EnergyDependentWidth # noqa: PLC0415 + + s, m0, w0, m1, m2, angular_momentum, meson_radius = self.args + width = EnergyDependentWidth( + s=s, + mass0=m0, + gamma0=w0, + m_a=m1, + m_b=m2, + angular_momentum=angular_momentum, + meson_radius=meson_radius, + name=Rf"\Gamma_{{{sp.latex(angular_momentum)}}}", + ) + return (m0 * w0) / (m0**2 - s - width * m0 * sp.I) + + @unevaluated class P(sp.Expr): s: Any @@ -150,3 +181,24 @@ def evaluate(self) -> sp.Piecewise: return sp.Piecewise(*[ (sp.sqrt(expr), sp.Eq(L, l_val)) for l_val, expr in cases.items() ]) + + +@unevaluated +class FormFactor(sp.Expr): + s: Any + m1: Any + m2: Any + angular_momentum: Any + meson_radius: Any + + _latex_repr_ = R"\mathcal{{F}}_{{{angular_momentum}}}\left({s}, {m1}, {m2}\right)" + + def evaluate(self): + s, m1, m2, angular_momentum, meson_radius = self.args + return formulate_form_factor( + s=s, + m_a=m1, + m_b=m2, + angular_momentum=angular_momentum, + meson_radius=meson_radius, + )